codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2#3import os4import os.path5import socket6import ssl7import unittest8 9import websocket10from websocket._exceptions import WebSocketProxyException, WebSocketException11from websocket._http import (12    _get_addrinfo_list,13    _start_proxied_socket,14    _tunnel,15    connect,16    proxy_info,17    read_headers,18    HAVE_PYTHON_SOCKS,19)20 21"""22test_http.py23websocket - WebSocket client library for Python24 25Copyright 2025 engn33r26 27Licensed under the Apache License, Version 2.0 (the "License");28you may not use this file except in compliance with the License.29You may obtain a copy of the License at30 31    http://www.apache.org/licenses/LICENSE-2.032 33Unless required by applicable law or agreed to in writing, software34distributed under the License is distributed on an "AS IS" BASIS,35WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.36See the License for the specific language governing permissions and37limitations under the License.38"""39 40try:41    from python_socks._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError42except:43    from websocket._http import ProxyConnectionError, ProxyError, ProxyTimeoutError44 45# Skip test to access the internet unless TEST_WITH_INTERNET == 146TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"47TEST_WITH_PROXY = os.environ.get("TEST_WITH_PROXY", "0") == "1"48# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -149LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")50TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"51 52 53class SockMock:54    def __init__(self):55        self.data = []56        self.sent = []57 58    def add_packet(self, data):59        self.data.append(data)60 61    def gettimeout(self):62        return None63 64    def recv(self, bufsize):65        if self.data:66            e = self.data.pop(0)67            if isinstance(e, Exception):68                raise e69            if len(e) > bufsize:70                self.data.insert(0, e[bufsize:])71            return e[:bufsize]72 73    def send(self, data):74        self.sent.append(data)75        return len(data)76 77    def close(self):78        pass79 80 81class HeaderSockMock(SockMock):82    def __init__(self, fname):83        SockMock.__init__(self)84        path = os.path.join(os.path.dirname(__file__), fname)85        with open(path, "rb") as f:86            self.add_packet(f.read())87 88 89class OptsList:90    def __init__(self):91        self.timeout = 192        self.sockopt = []93        self.sslopt = {"cert_reqs": ssl.CERT_NONE}94 95 96class HttpTest(unittest.TestCase):97    def test_read_header(self):98        status, header, _ = read_headers(HeaderSockMock("data/header01.txt"))99        self.assertEqual(status, 101)100        self.assertEqual(header["connection"], "Upgrade")101        # header02.txt is intentionally malformed102        self.assertRaises(103            WebSocketException, read_headers, HeaderSockMock("data/header02.txt")104        )105 106    def test_tunnel(self):107        self.assertRaises(108            WebSocketProxyException,109            _tunnel,110            HeaderSockMock("data/header01.txt"),111            "example.com",112            80,113            ("username", "password"),114        )115        self.assertRaises(116            WebSocketProxyException,117            _tunnel,118            HeaderSockMock("data/header02.txt"),119            "example.com",120            80,121            ("username", "password"),122        )123 124    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")125    def test_connect(self):126        # Not currently testing an actual proxy connection, so just check whether proxy errors are raised. This requires internet for a DNS lookup127        if HAVE_PYTHON_SOCKS:128            # Need this check, otherwise case where python_socks is not installed triggers129            # websocket._exceptions.WebSocketException: Python Socks is needed for SOCKS proxying but is not available130            self.assertRaises(131                (ProxyTimeoutError, OSError),132                _start_proxied_socket,133                "wss://example.com",134                OptsList(),135                proxy_info(136                    http_proxy_host="example.com",137                    http_proxy_port="8080",138                    proxy_type="socks4",139                    http_proxy_timeout=1,140                ),141            )142            self.assertRaises(143                (ProxyTimeoutError, OSError),144                _start_proxied_socket,145                "wss://example.com",146                OptsList(),147                proxy_info(148                    http_proxy_host="example.com",149                    http_proxy_port="8080",150                    proxy_type="socks4a",151                    http_proxy_timeout=1,152                ),153            )154            self.assertRaises(155                (ProxyTimeoutError, OSError),156                _start_proxied_socket,157                "wss://example.com",158                OptsList(),159                proxy_info(160                    http_proxy_host="example.com",161                    http_proxy_port="8080",162                    proxy_type="socks5",163                    http_proxy_timeout=1,164                ),165            )166            self.assertRaises(167                (ProxyTimeoutError, OSError),168                _start_proxied_socket,169                "wss://example.com",170                OptsList(),171                proxy_info(172                    http_proxy_host="example.com",173                    http_proxy_port="8080",174                    proxy_type="socks5h",175                    http_proxy_timeout=1,176                ),177            )178            self.assertRaises(179                ProxyConnectionError,180                connect,181                "wss://example.com",182                OptsList(),183                proxy_info(184                    http_proxy_host="127.0.0.1",185                    http_proxy_port=9999,186                    proxy_type="socks4",187                    http_proxy_timeout=1,188                ),189                None,190            )191 192        self.assertRaises(193            TypeError,194            _get_addrinfo_list,195            None,196            80,197            True,198            proxy_info(199                http_proxy_host="127.0.0.1", http_proxy_port="9999", proxy_type="http"200            ),201        )202        self.assertRaises(203            TypeError,204            _get_addrinfo_list,205            None,206            80,207            True,208            proxy_info(209                http_proxy_host="127.0.0.1", http_proxy_port="9999", proxy_type="http"210            ),211        )212        self.assertRaises(213            socket.timeout,214            connect,215            "wss://google.com",216            OptsList(),217            proxy_info(218                http_proxy_host="8.8.8.8",219                http_proxy_port=9999,220                proxy_type="http",221                http_proxy_timeout=1,222            ),223            None,224        )225        self.assertEqual(226            connect(227                "wss://google.com",228                OptsList(),229                proxy_info(230                    http_proxy_host="8.8.8.8", http_proxy_port=8080, proxy_type="http"231                ),232                True,233            ),234            (True, ("google.com", 443, "/")),235        )236        # The following test fails on Mac OS with a gaierror, not an OverflowError237        # self.assertRaises(OverflowError, connect, "wss://example.com", OptsList(), proxy_info(http_proxy_host="127.0.0.1", http_proxy_port=99999, proxy_type="socks4", timeout=2), False)238 239    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")240    @unittest.skipUnless(241        TEST_WITH_PROXY, "This test requires a HTTP proxy to be running on port 8899"242    )243    @unittest.skipUnless(244        TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"245    )246    def test_proxy_connect(self):247        ws = websocket.WebSocket()248        ws.connect(249            f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",250            http_proxy_host="127.0.0.1",251            http_proxy_port="8899",252            proxy_type="http",253        )254        ws.send("Hello, Server")255        server_response = ws.recv()256        self.assertEqual(server_response, "Hello, Server")257        # self.assertEqual(_start_proxied_socket("wss://api.bitfinex.com/ws/2", OptsList(), proxy_info(http_proxy_host="127.0.0.1", http_proxy_port="8899", proxy_type="http"))[1], ("api.bitfinex.com", 443, '/ws/2'))258        self.assertEqual(259            _get_addrinfo_list(260                "api.bitfinex.com",261                443,262                True,263                proxy_info(264                    http_proxy_host="127.0.0.1",265                    http_proxy_port="8899",266                    proxy_type="http",267                ),268            ),269            (270                socket.getaddrinfo(271                    "127.0.0.1", 8899, 0, socket.SOCK_STREAM, socket.SOL_TCP272                ),273                True,274                None,275            ),276        )277        self.assertEqual(278            connect(279                "wss://api.bitfinex.com/ws/2",280                OptsList(),281                proxy_info(282                    http_proxy_host="127.0.0.1", http_proxy_port=8899, proxy_type="http"283                ),284                None,285            )[1],286            ("api.bitfinex.com", 443, "/ws/2"),287        )288        # TODO: Test SOCKS4 and SOCK5 proxies with unit tests289 290    @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")291    def test_sslopt(self):292        ssloptions = {293            "check_hostname": False,294            "server_hostname": "ServerName",295            "ssl_version": ssl.PROTOCOL_TLS_CLIENT,296            "ciphers": "TLS_AES_256_GCM_SHA384:TLS_CHACHA20_POLY1305_SHA256:\297                        TLS_AES_128_GCM_SHA256:ECDHE-ECDSA-AES256-GCM-SHA384:\298                        ECDHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384:\299                        ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:\300                        DHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-GCM-SHA256:\301                        ECDHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES128-GCM-SHA256:\302                        ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:\303                        DHE-RSA-AES256-SHA256:ECDHE-ECDSA-AES128-SHA256:\304                        ECDHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA256:\305                        ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA",306            "ecdh_curve": "prime256v1",307        }308        ws_ssl1 = websocket.WebSocket(sslopt=ssloptions)309        ws_ssl1.connect("wss://api.bitfinex.com/ws/2")310        ws_ssl1.send("Hello")311        ws_ssl1.close()312 313        ws_ssl2 = websocket.WebSocket(sslopt={"check_hostname": True})314        ws_ssl2.connect("wss://api.bitfinex.com/ws/2")315        ws_ssl2.close316 317    def test_proxy_info(self):318        self.assertEqual(319            proxy_info(320                http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"321            ).proxy_protocol,322            "http",323        )324        self.assertRaises(325            ProxyError,326            proxy_info,327            http_proxy_host="127.0.0.1",328            http_proxy_port="8080",329            proxy_type="badval",330        )331        self.assertEqual(332            proxy_info(333                http_proxy_host="example.com", http_proxy_port="8080", proxy_type="http"334            ).proxy_host,335            "example.com",336        )337        self.assertEqual(338            proxy_info(339                http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"340            ).proxy_port,341            "8080",342        )343        self.assertEqual(344            proxy_info(345                http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"346            ).auth,347            None,348        )349        self.assertEqual(350            proxy_info(351                http_proxy_host="127.0.0.1",352                http_proxy_port="8080",353                proxy_type="http",354                http_proxy_auth=("my_username123", "my_pass321"),355            ).auth[0],356            "my_username123",357        )358        self.assertEqual(359            proxy_info(360                http_proxy_host="127.0.0.1",361                http_proxy_port="8080",362                proxy_type="http",363                http_proxy_auth=("my_username123", "my_pass321"),364            ).auth[1],365            "my_pass321",366        )367 368 369if __name__ == "__main__":370    unittest.main()371