Preserve custom Host headers for WebSocket requests (#6172)

Co-authored-by: cuteNBD <206630735+bingdongni@users.noreply.github.com>

Author: bingdongni

lib/request/connect.py

@@ -521,6 +521,8 @@ def getPage(**kwargs):
 
             # Prepare HTTP headers
             headers = forgeHeaders({HTTP_HEADER.COOKIE: cookie, HTTP_HEADER.USER_AGENT: ua, HTTP_HEADER.REFERER: referer, HTTP_HEADER.HOST: host or getHeader(dict(conf.httpHeaders), HTTP_HEADER.HOST) or getHostHeader(url)}, base=None if target else {})
+            # WebSocket builds its own default Host, but must preserve explicit overrides.
+            customWsHost = webSocket and bool(host or getHeader(dict(conf.httpHeaders), HTTP_HEADER.HOST) or getHeader(auxHeaders, HTTP_HEADER.HOST) is not None)
 
             if HTTP_HEADER.COOKIE in headers:
                 cookie = headers[HTTP_HEADER.COOKIE]
@@ -573,9 +575,13 @@ def getPage(**kwargs):
                     if ':' in line:
                         header, value = line.split(':', 1)
                         headers[header] = value
+                        if webSocket and header.upper() == HTTP_HEADER.HOST.upper():
+                            headers[HTTP_HEADER.HOST] = value
+                            customWsHost = True
 
             if conf.localhost:
                 headers[HTTP_HEADER.HOST] = "localhost"
+                customWsHost = True
 
             for key, value in list(headers.items()):
                 if key.upper() == HTTP_HEADER.ACCEPT_ENCODING.upper():
@@ -600,7 +606,7 @@ def getPage(**kwargs):
                 try:
                     ws.settimeout(WEBSOCKET_INITIAL_TIMEOUT if kb.webSocketRecvCount is None else timeout)
                     wsHeaders = tuple("%s: %s" % (getUnicode(key), getUnicode(value)) for key, value in headers.items() if getUnicode(key).upper() != HTTP_HEADER.HOST.upper())
-                    ws.connect(url, header=wsHeaders, cookie=cookie)  # WebSocket will add Host field of headers automatically
+                    ws.connect(url, header=wsHeaders, cookie=cookie, host_header=getHeader(headers, HTTP_HEADER.HOST) if customWsHost else None)
                     ws.send(urldecode(post or ""))
 
                     _page = []

lib/request/websocket.py

@@ -58,7 +58,7 @@ def settimeout(self, timeout):
         if self.sock is not None:
             self.sock.settimeout(timeout)
 
-    def connect(self, url, header=None, cookie=None):
+    def connect(self, url, header=None, cookie=None, host_header=None):
         parts = _urllib.parse.urlsplit(url)
         secure = parts.scheme == "wss"
         host = parts.hostname
@@ -75,6 +75,8 @@ def connect(self, url, header=None, cookie=None):
         hostport = "[%s]" % host if ":" in host else host       # bracket IPv6 literals
         if port != (443 if secure else 80):    # Note: only the scheme's own default port gets omitted
             hostport = "%s:%d" % (hostport, port)
+        if host_header is not None:
+            hostport = getText(host_header)
 
         key = getText(base64.b64encode(os.urandom(16)))
         lines = ["GET %s HTTP/1.1" % resource,

tests/test_websocket.py

@@ -33,6 +33,9 @@
     OPCODE_PING,
     OPCODE_CLOSE,
 )
+from lib.request.connect import Connect
+
+_GET_PAGE = Connect.__dict__["getPage"]
 
 
 class _FakeSock(object):
@@ -162,6 +165,88 @@ def sendall(self, data):
         finally:
             W.socket.create_connection = original
 
+    def test_getpage_preserves_host_header(self):
+        import socket
+        import tempfile
+        import lib.request.websocket as W
+        from lib.core.data import conf, kb
+        from lib.core.option import _setKnowledgeBaseAttributes
+        from lib.core.threads import getCurrentThreadData
+
+        class _HandshakeSock(_FakeSock):
+            def __init__(self):
+                _FakeSock.__init__(self)
+                self.handshake = None
+                self.response = b""
+
+            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 + _GUID.encode("ascii")).digest())
+                    self.response = b"HTTP/1.1 101 Switching Protocols\r\nSec-WebSocket-Accept: " + accept + b"\r\n\r\n" + _serverFrame("ok")
+
+            def recv(self, size):
+                result, self.response = self.response[:size], self.response[size:]
+                return result
+
+        sockets = []
+        def create_socket(*args, **kwargs):
+            sock = _HandshakeSock()
+            sockets.append(sock)
+            return sock
+
+        oldConf, oldKb = dict(conf), dict(kb)
+        threadData = getCurrentThreadData()
+        oldThreadData = dict(threadData.__dict__)
+        oldTimeout = socket.getdefaulttimeout()
+        original = W.socket.create_connection
+        oldGetPage = Connect.__dict__["getPage"]
+        fd, headersFile = tempfile.mkstemp(prefix="sqlmap-ws-headers-")
+        try:
+            with os.fdopen(fd, "w") as f:
+                f.write("hOsT: file.example\n")
+            threadData.__dict__.clear()
+            threadData.reset()
+            _setKnowledgeBaseAttributes(False)
+            kb.wizardMode, kb.webSocketRecvCount = True, 1
+            conf.url = "ws://example.test:8080/test"
+            conf.hostname, conf.scheme, conf.port, conf.path = "example.test", "ws", 8080, "/test"
+            conf.parameters, conf.data = {}, None
+            conf.abortCode, conf.ignoreCode = [], []
+            conf.timeout, conf.batch, conf.keepAlive = 3, True, False
+            conf.retryOn, conf.ignoreRedirects, conf.scope = None, False, None
+            W.socket.create_connection = create_socket
+            Connect.getPage = _GET_PAGE
+            cases = (
+                ("ws://example.test:8080/test", [], None, None, None, False, "example.test:8080"),
+                ("ws://user:pass@example.test:8080/test", [], None, None, None, False, "example.test:8080"),
+                ("ws://example.test:8080/test", [('hOsT', 'configured.example')], None, None, None, False, "configured.example"),
+                ("ws://example.test:8080/test", [('Host', 'configured.example')], "override.example:8443", None, None, False, "override.example:8443"),
+                ("ws://example.test:8080/test", [], None, {'Host': 'aux.example'}, None, False, "aux.example"),
+                ("ws://example.test:8080/test", [], None, None, headersFile, False, "file.example"),
+                ("ws://example.test:8080/test", [], None, None, None, True, "localhost"),
+            )
+            for url, headers, host, aux, path, localhost, expected in cases:
+                conf.url = url
+                conf.httpHeaders = headers
+                conf.localhost, kb.headersFile = localhost, path
+                self.assertEqual(Connect.getPage(url=url, host=host, auxHeaders=aux)[0], "ok")
+                hostLines = [line for line in sockets[-1].handshake.split(b"\r\n") if line.lower().startswith(b"host:")]
+                self.assertEqual(len(hostLines), 1)
+                self.assertEqual(hostLines[0].split(b":", 1)[1].strip(), expected.encode("ascii"))
+        finally:
+            Connect.getPage = oldGetPage
+            W.socket.create_connection = original
+            socket.setdefaulttimeout(oldTimeout)
+            conf.clear(); conf.update(oldConf)
+            kb.clear(); kb.update(oldKb)
+            threadData.__dict__.clear(); threadData.__dict__.update(oldThreadData)
+            os.unlink(headersFile)
+
 
 if __name__ == "__main__":
     unittest.main()