Adding native websocket support

tests/test_websocket.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python
+
+"""
+Copyright (c) 2006-2026 sqlmap developers (https://sqlmap.org)
+See the file 'LICENSE' for copying permission
+
+Unit coverage for the PURE (network-free) parts of the native WebSocket client in
+lib/request/websocket.py: the RFC 6455 accept-key computation, client frame masking,
+the length-encoding boundaries (7/16/64-bit), fragment reassembly and control-frame
+handling. No socket is opened - frames are fed through a primed buffer and a fake sink.
+
+stdlib unittest only (no pytest / no pip); works on Python 2.7 and 3.x.
+"""
+
+import base64
+import hashlib
+import os
+import struct
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _testutils import bootstrap
+bootstrap()
+
+from lib.request.websocket import (
+    WebSocket,
+    WebSocketConnectionClosedException,
+    WebSocketTimeoutException,
+    _GUID,
+    OPCODE_TEXT,
+    OPCODE_CONTINUATION,
+    OPCODE_PING,
+    OPCODE_CLOSE,
+)
+
+
+class _FakeSock(object):
+    """Captures everything the client sends, so masked client frames / PONGs can be inspected."""
+    def __init__(self):
+        self.sent = b""
+
+    def sendall(self, data):
+        self.sent += data
+
+    def close(self):
+        pass
+
+
+def _serverFrame(data, opcode=OPCODE_TEXT, fin=1):
+    """Build an (unmasked, server->client) frame carrying data."""
+    if not isinstance(data, bytes):
+        data = data.encode("utf-8")
+    frame = bytearray([(fin << 7) | opcode])
+    length = len(data)
+    if length < 126:
+        frame.append(length)
+    elif length < 65536:
+        frame.append(126); frame += struct.pack("!H", length)
+    else:
+        frame.append(127); frame += struct.pack("!Q", length)
+    frame += data
+    return bytes(frame)
+
+
+def _client(buffer=b""):
+    ws = WebSocket.__new__(WebSocket)   # bypass connect(): no socket
+    ws.sock = _FakeSock()
+    ws.status = 101
+    ws._headers = {}
+    ws._timeout = None
+    ws._buffer = buffer
+    ws._closed = False
+    return ws
+
+
+class TestWebSocket(unittest.TestCase):
+    def test_accept_key_rfc6455_vector(self):
+        # RFC 6455 section 1.3 canonical example
+        key = "dGhlIHNhbXBsZSBub25jZQ=="
+        accept = base64.b64encode(hashlib.sha1((key + _GUID).encode("ascii")).digest()).decode("ascii")
+        self.assertEqual(accept, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
+
+    def test_client_frame_is_masked_and_roundtrips(self):
+        ws = _client()
+        ws._sendFrame(b"hello", OPCODE_TEXT)
+        raw = ws.sock.sent
+        self.assertEqual(bytearray(raw)[0], 0x80 | OPCODE_TEXT)   # FIN + text
+        self.assertTrue(bytearray(raw)[1] & 0x80, "client frame must set the mask bit")
+
+        # feeding the client's own (masked) frame back through the parser must recover the payload
+        fin, opcode, payload = _client(raw)._recvFrame()
+        self.assertEqual((fin, opcode, bytes(payload)), (1, OPCODE_TEXT, b"hello"))
+
+    def test_length_encoding_boundaries(self):
+        for size in (125, 126, 65535, 65536):
+            ws = _client()
+            ws._sendFrame(b"A" * size, OPCODE_TEXT)
+            fin, opcode, payload = _client(ws.sock.sent)._recvFrame()
+            self.assertEqual(len(payload), size, msg="round-trip failed at length %d" % size)
+
+    def test_recv_reassembles_fragments(self):
+        buf = _serverFrame("ab", OPCODE_TEXT, fin=0) + _serverFrame("cd", OPCODE_CONTINUATION, fin=1)
+        self.assertEqual(_client(buf).recv(), "abcd")
+
+    def test_recv_answers_ping_then_returns_data(self):
+        ws = _client(_serverFrame("hi", OPCODE_PING) + _serverFrame("data", OPCODE_TEXT))
+        self.assertEqual(ws.recv(), "data")
+        # a PONG (opcode 0xA) carrying the ping payload must have been sent back
+        pong = bytearray(ws.sock.sent)
+        self.assertEqual(pong[0], 0x80 | 0xA)
+
+    def test_recv_close_raises(self):
+        ws = _client(_serverFrame(struct.pack("!H", 1000), OPCODE_CLOSE))
+        self.assertRaises(WebSocketConnectionClosedException, ws.recv)
+
+    def test_read_timeout_maps_to_ws_timeout(self):
+        import socket as _socket
+        import ssl as _ssl
+
+        class _RaisingSock(object):
+            def __init__(self, exc):
+                self.exc = exc
+            def recv(self, n):
+                raise self.exc
+
+        # both a plain socket timeout and Python 2's TLS 'read operation timed out' must surface as
+        # WebSocketTimeoutException (sqlmap's frame loop relies on it), while other SSL errors propagate
+        for exc in (_socket.timeout("timed out"), _ssl.SSLError("The read operation timed out")):
+            ws = _client(); ws.sock = _RaisingSock(exc)
+            self.assertRaises(WebSocketTimeoutException, ws.recv)
+
+        ws = _client(); ws.sock = _RaisingSock(_ssl.SSLError("decryption failed"))
+        self.assertRaises(_ssl.SSLError, ws.recv)
+
+
+if __name__ == "__main__":
+    unittest.main()