diff --git a/lib/request/connect.py b/lib/request/connect.py index 04fc3f3a413..2f3695e45dd 100644 --- a/lib/request/connect.py +++ b/lib/request/connect.py @@ -613,17 +613,17 @@ def getPage(**kwargs): warnMsg = "too large websocket response detected. Automatically trimming it" singleTimeWarnMessage(warnMsg) break - except websocket.WebSocketTimeoutException: + except (websocket.WebSocketTimeoutException, websocket.WebSocketNormalClosureException): kb.webSocketRecvCount = len(_page) break else: for i in xrange(max(1, kb.webSocketRecvCount)): # Note: a later response may carry fewer frames than the calibration one - # (e.g. a FALSE boolean-blind reply); stop on timeout instead of letting it - # bubble up as a failed request, mirroring the initial recv loop above + # (e.g. a FALSE boolean-blind reply); stop on timeout or normal closure, + # mirroring the initial recv loop above try: _page.append(ws.recv()) - except websocket.WebSocketTimeoutException: + except (websocket.WebSocketTimeoutException, websocket.WebSocketNormalClosureException): break page = "\n".join(_page) diff --git a/lib/request/websocket.py b/lib/request/websocket.py index 88d58a790f1..580a862e5d3 100644 --- a/lib/request/websocket.py +++ b/lib/request/websocket.py @@ -39,6 +39,9 @@ class WebSocketTimeoutException(WebSocketException): class WebSocketConnectionClosedException(WebSocketException): pass +class WebSocketNormalClosureException(WebSocketConnectionClosedException): + pass + class WebSocket(object): """ Minimal RFC 6455 client exposing the websocket-client subset used by sqlmap: settimeout(), connect(), @@ -157,10 +160,14 @@ def _sendFrame(self, data, opcode): def recv(self): data = bytearray() + receivedFragment = False while True: fin, opcode, payload = self._recvFrame() if opcode == OPCODE_CLOSE: self._closed = True + normal = not payload or (len(payload) >= 2 and struct.unpack("!H", bytes(payload[:2]))[0] in (1000, 1001)) + if normal and not receivedFragment: + raise WebSocketNormalClosureException("WebSocket connection closed") raise WebSocketConnectionClosedException("WebSocket connection closed") elif opcode == OPCODE_PING: self._sendFrame(bytes(payload), OPCODE_PONG) @@ -169,6 +176,7 @@ def recv(self): continue data += payload + receivedFragment = True if fin: break diff --git a/tests/test_websocket_response.py b/tests/test_websocket_response.py new file mode 100644 index 00000000000..0d56c95ed7a --- /dev/null +++ b/tests/test_websocket_response.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python + +""" +Copyright (c) 2006-2026 sqlmap developers (https://sqlmap.org) +See the file 'LICENSE' for copying permission +""" + +import base64 +import hashlib +import os +import socket +import struct +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from _testutils import bootstrap +bootstrap() + +from lib.core.data import conf, kb +from lib.core.exception import SqlmapConnectionException +from lib.core.option import _setKnowledgeBaseAttributes +from lib.core.threads import getCurrentThreadData +from lib.request.connect import Connect +import lib.request.websocket as websocket + +_GET_PAGE = Connect.__dict__["getPage"] + + +def _frame(payload, opcode=1, fin=1): + return bytes(bytearray([(fin << 7) | opcode, len(payload)])) + payload + + +class _Socket(object): + def __init__(self, frames): + self.frames = frames + self.response = b"" + self.handshake = None + + def settimeout(self, value): + pass + + def sendall(self, data): + if self.handshake is None: + self.handshake = data + key = next(line.split(b":", 1)[1].strip() for line in data.split(b"\r\n") if line.lower().startswith(b"sec-websocket-key:")) + accept = base64.b64encode(hashlib.sha1(key + websocket._GUID.encode("ascii")).digest()) + self.response = b"HTTP/1.1 101 Switching Protocols\r\nSec-WebSocket-Accept: " + accept + b"\r\n\r\n" + self.frames + + def recv(self, size): + result, self.response = self.response[:size], self.response[size:] + return result + + def close(self): + pass + + +class TestWebSocketResponseClose(unittest.TestCase): + def setUp(self): + self.oldConf, self.oldKb = dict(conf), dict(kb) + self.threadData = getCurrentThreadData() + self.oldThreadData = dict(self.threadData.__dict__) + self.oldTimeout = socket.getdefaulttimeout() + self.oldSocket = websocket.socket.create_connection + self.oldGetPage = Connect.__dict__["getPage"] + self.threadData.__dict__.clear() + self.threadData.reset() + _setKnowledgeBaseAttributes(False) + kb.wizardMode = True + conf.url = "ws://example.test:8080/test" + conf.hostname, conf.scheme, conf.port, conf.path = "example.test", "ws", 8080, "/test" + conf.httpHeaders, conf.parameters, conf.data = [], {}, None + conf.abortCode, conf.ignoreCode = [], [] + conf.timeout, conf.batch, conf.keepAlive, conf.retries = 3, True, False, 0 + conf.retryOn, conf.ignoreRedirects, conf.scope = None, False, None + Connect.getPage = _GET_PAGE + websocket.socket.create_connection = lambda *args, **kwargs: _Socket(self.frames) + + def tearDown(self): + Connect.getPage = self.oldGetPage + websocket.socket.create_connection = self.oldSocket + socket.setdefaulttimeout(self.oldTimeout) + conf.clear(); conf.update(self.oldConf) + kb.clear(); kb.update(self.oldKb) + self.threadData.__dict__.clear(); self.threadData.__dict__.update(self.oldThreadData) + + def test_normal_close_after_response(self): + for status in (b"", struct.pack("!H", 1000), struct.pack("!H", 1001)): + for count in (None, 3): + self.frames = _frame(b"one") + _frame(b"two") + _frame(status, opcode=8) + kb.webSocketRecvCount = count + self.assertEqual(Connect.getPage(url=conf.url)[0], "one\ntwo") + if count is None: + self.assertEqual(kb.webSocketRecvCount, 2) + + def test_normal_close_without_response(self): + self.frames = _frame(struct.pack("!H", 1000), opcode=8) + kb.webSocketRecvCount = None + self.assertEqual(Connect.getPage(url=conf.url)[0], "") + self.assertEqual(kb.webSocketRecvCount, 0) + + def test_abnormal_close_remains_an_error(self): + for status in (b"\x03", struct.pack("!H", 1002), struct.pack("!H", 1011)): + self.frames = _frame(b"one") + _frame(status, opcode=8) + for count in (None, 3): + kb.webSocketRecvCount = count + self.assertRaises(SqlmapConnectionException, Connect.getPage, url=conf.url) + + def test_truncated_frame_remains_an_error(self): + for frames in (_frame(b"one"), _frame(b"one") + b"\x81\x05ab"): + self.frames = frames + for count in (None, 3): + kb.webSocketRecvCount = count + self.assertRaises(SqlmapConnectionException, Connect.getPage, url=conf.url) + + def test_close_during_fragmented_message_remains_an_error(self): + for fragment in (b"", b"partial"): + self.frames = _frame(fragment, fin=0) + _frame(struct.pack("!H", 1000), opcode=8) + kb.webSocketRecvCount = None + self.assertRaises(SqlmapConnectionException, Connect.getPage, url=conf.url) + + +if __name__ == "__main__": + unittest.main()