From 9b5dcbbbb2f0bc39338fe37bff81510c1a3988e1 Mon Sep 17 00:00:00 2001
From: ricterz <ricterzheng@gmail.com>
Date: Tue, 24 Mar 2015 18:21:50 +0800
Subject: [PATCH] modified error handle #1198

---
 lib/request/connect.py | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/lib/request/connect.py b/lib/request/connect.py
index f556ab674..9b3ae3da9 100644
--- a/lib/request/connect.py
+++ b/lib/request/connect.py
@@ -367,21 +367,12 @@ class Connect(object):
             post = unicodeencode(post, kb.pageEncoding)
 
             if is_websocket:
-                try:
-                    ws = websocket.WebSocket()
-                    ws.connect(url)
-                    ws.send(urldecode(post) if post else '')
-                    response = ws.recv()
-                    ws.close()
-                    return response, {}, 101
-
-                except websocket.WebSocketConnectionClosedException:
-                    # TODO: more exception to handle
-                    warnMsg = "connection was forcibly closed by the target URL"
-                    logger.critical(warnMsg)
-                    return Connect._retryProxy(**kwargs)
-                except Exception:
-                    return None, None, None
+                ws = websocket.WebSocket()
+                ws.connect(url)
+                ws.send(urldecode(post) if post else '')
+                response = ws.recv()
+                ws.close()
+                return response, {}, 101
 
             elif method and method not in (HTTPMETHOD.GET, HTTPMETHOD.POST):
                 method = unicodeencode(method)
@@ -557,7 +548,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), e:
+        except (urllib2.URLError, socket.error, socket.timeout, httplib.BadStatusLine, httplib.IncompleteRead, struct.error, ProxyError, SqlmapCompressionException, websocket.WebSocketException), e:
             tbMsg = traceback.format_exc()
 
             if "no host given" in tbMsg:
@@ -582,6 +573,10 @@ class Connect(object):
             elif "IncompleteRead" in tbMsg:
                 warnMsg = "there was an incomplete read error while retrieving data "
                 warnMsg += "from the target URL"
+            elif "Handshake status" in tbMsg:
+                status = re.search("Handshake status ([\d]{3})", tbMsg)
+                errMsg = "websocket handshake status %s" % status.group(1) if status else 'unknown'
+                raise SqlmapConnectionException(errMsg)
             else:
                 warnMsg = "unable to connect to the target URL"