From bbfdb02a0ef0577ab2312861448a68a373c58921 Mon Sep 17 00:00:00 2001
From: ricterz <ricterzheng@gmail.com>
Date: Tue, 24 Mar 2015 22:25:16 +0800
Subject: [PATCH] fix mandatorily depend of websocket #1198

---
 lib/core/option.py     | 11 +++++++++++
 lib/request/connect.py | 10 ++++++++--
 lib/utils/deps.py      | 11 +++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/lib/core/option.py b/lib/core/option.py
index 0eafd1b20..6bfc91760 100644
--- a/lib/core/option.py
+++ b/lib/core/option.py
@@ -2137,6 +2137,16 @@ def _setTorSocksProxySettings():
     socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS5 if conf.torType == PROXY_TYPE.SOCKS5 else socks.PROXY_TYPE_SOCKS4, LOCALHOST, conf.torPort or DEFAULT_TOR_SOCKS_PORT)
     socks.wrapmodule(urllib2)
 
+def _checkWebSocket():
+    infoMsg = "checking URL is WebSocket or not"
+    logger.debug(infoMsg)
+    if conf.url and (conf.url.startswith("ws:/") or conf.url.startswith("wss:/")):
+        try:
+            from websocket import ABNF
+        except ImportError:
+            errMsg = "it seems that python 'websocket-client' third-party library not be installed. "
+            raise SqlmapMissingDependence(errMsg)
+
 def _checkTor():
     if not conf.checkTor:
         return
@@ -2383,6 +2393,7 @@ def init():
     _setWafFunctions()
     _setTrafficOutputFP()
     _resolveCrossReferences()
+    _checkWebSocket()
 
     parseTargetUrl()
     parseTargetDirect()
diff --git a/lib/request/connect.py b/lib/request/connect.py
index 3adb88bd5..ba2e7909b 100644
--- a/lib/request/connect.py
+++ b/lib/request/connect.py
@@ -18,7 +18,13 @@ import time
 import traceback
 import urllib2
 import urlparse
-import websocket
+
+try:
+    import websocket
+    from websocket import WebSocketException
+except ImportError:
+    class WebSocketException(Exception):
+        pass
 
 from extra.safe2bin.safe2bin import safecharencode
 from lib.core.agent import agent
@@ -551,7 +557,7 @@ class Connect(object):
                 debugMsg = "got HTTP error code: %d (%s)" % (code, status)
                 logger.debug(debugMsg)
 
-        except (urllib2.URLError, socket.error, socket.timeout, httplib.BadStatusLine, httplib.IncompleteRead, struct.error, ProxyError, SqlmapCompressionException, websocket.WebSocketException), e:
+        except (urllib2.URLError, socket.error, socket.timeout, httplib.BadStatusLine, httplib.IncompleteRead, struct.error, ProxyError, SqlmapCompressionException, WebSocketException), e:
             tbMsg = traceback.format_exc()
 
             if "no host given" in tbMsg:
diff --git a/lib/utils/deps.py b/lib/utils/deps.py
index ea2f661e9..d25c8ea94 100644
--- a/lib/utils/deps.py
+++ b/lib/utils/deps.py
@@ -78,6 +78,17 @@ def checkDependencies():
         logger.warn(warnMsg)
         missing_libraries.add('python-ntlm')
 
+    try:
+        from websocket import ABNF
+        debugMsg = "'python websocket-client' library is found"
+        logger.debug(debugMsg)
+    except ImportError:
+        warnMsg = "sqlmap requires 'python websocket-client' third-party library for "
+        warnMsg += "if you plan to attack a web application behind websocket. "
+        warnMsg += "Download from https://pypi.python.org/pypi/websocket-client/"
+        logger.warn(warnMsg)
+        missing_libraries.add('websocket-client')
+
     if IS_WIN:
         try:
             import pyreadline