From 4b8d1ed165c26c65129096febd7f1e167bc062a7 Mon Sep 17 00:00:00 2001
From: animic <448395347@qq.com>
Date: Fri, 22 Nov 2024 23:48:23 +0800
Subject: [PATCH] feat: add proxy support to WS (#1467)

* feat: add https_proxy support for ThreadedWebsocketManager and  BinanceSocketManager

* handle refactor

* remove requirement

* add readme

---------

Co-authored-by: carlosmiei <43336371+carlosmiei@users.noreply.github.com>
---
 README.rst                           |  2 +-
 binance/ws/reconnecting_websocket.py | 24 +++++++++++++++++++++++-
 binance/ws/streams.py                |  5 ++++-
 binance/ws/threaded_stream.py        |  2 ++
 setup.py                             |  1 -
 5 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/README.rst b/README.rst
index 334d9d0ff..7e154db85 100755
--- a/README.rst
+++ b/README.rst
@@ -77,7 +77,7 @@ Features
 - Futures Trading
 - Porfolio Margin Trading
 - Vanilla Options
-- Proxy support
+- Proxy support (REST and WS)
 - Orjson support for faster JSON parsing
 - Support other domains (.us, .jp, etc)
 
diff --git a/binance/ws/reconnecting_websocket.py b/binance/ws/reconnecting_websocket.py
index 03c392e53..424a23835 100644
--- a/binance/ws/reconnecting_websocket.py
+++ b/binance/ws/reconnecting_websocket.py
@@ -19,6 +19,16 @@
 except ImportError:
     from websockets import ConnectionClosedError  # type: ignore
 
+
+Proxy = None
+proxy_connect = None
+try:
+    from websockets_proxy import Proxy as w_Proxy, proxy_connect as w_proxy_connect
+    Proxy = w_Proxy
+    proxy_connect = w_proxy_connect
+except ImportError:
+    pass
+
 import websockets as ws
 
 from binance.exceptions import BinanceWebsocketUnableToConnect
@@ -41,6 +51,7 @@ def __init__(
         prefix: str = "ws/",
         is_binary: bool = False,
         exit_coro=None,
+        https_proxy: Optional[str] = None,
         **kwargs,
     ):
         self._loop = get_loop()
@@ -57,6 +68,7 @@ def __init__(
         self.ws_state = WSListenerState.INITIALISING
         self._queue = asyncio.Queue()
         self._handle_read_loop = None
+        self._https_proxy = https_proxy
         self._ws_kwargs = kwargs
 
     def json_dumps(self, msg):
@@ -93,10 +105,20 @@ async def connect(self):
         self._log.debug("Establishing new WebSocket connection")
         self.ws_state = WSListenerState.RECONNECTING
         await self._before_connect()
+
         ws_url = (
             f"{self._url}{getattr(self, '_prefix', '')}{getattr(self, '_path', '')}"
         )
-        self._conn = ws.connect(ws_url, close_timeout=0.1, **self._ws_kwargs)  # type: ignore
+
+        # handle https_proxy
+        if self._https_proxy:
+            if not Proxy or not proxy_connect:
+                raise ImportError("websockets_proxy is not installed, please install it to use a websockets proxy (pip install websockets_proxy)")
+            proxy = Proxy.from_url(self._https_proxy) # type: ignore
+            self._conn = proxy_connect(ws_url, close_timeout=0.1, proxy=proxy, **self._ws_kwargs) # type: ignore
+        else:
+            self._conn = ws.connect(ws_url, close_timeout=0.1, **self._ws_kwargs)  # type: ignore
+
         try:
             self.ws = await self._conn.__aenter__()
         except Exception as e:  # noqa
diff --git a/binance/ws/streams.py b/binance/ws/streams.py
index 4c9e4a26e..24f597804 100755
--- a/binance/ws/streams.py
+++ b/binance/ws/streams.py
@@ -82,6 +82,7 @@ def _get_socket(
                 prefix=prefix,
                 exit_coro=lambda p: self._exit_socket(f"{socket_type}_{p}"),
                 is_binary=is_binary,
+                https_proxy=self._client.https_proxy,
                 **self.ws_kwargs,
             )
 
@@ -104,6 +105,7 @@ def _get_account_socket(
                 exit_coro=self._exit_socket,
                 is_binary=is_binary,
                 user_timeout=self._user_timeout,
+                https_proxy=self._client.https_proxy,
                 **self.ws_kwargs,
             )
 
@@ -1094,10 +1096,11 @@ def __init__(
         tld: str = "com",
         testnet: bool = False,
         session_params: Optional[Dict[str, Any]] = None,
+        https_proxy: Optional[str] = None,
         loop: Optional[asyncio.AbstractEventLoop] = None,
     ):
         super().__init__(
-            api_key, api_secret, requests_params, tld, testnet, session_params, loop
+            api_key, api_secret, requests_params, tld, testnet, session_params, https_proxy, loop
         )
         self._bsm: Optional[BinanceSocketManager] = None
 
diff --git a/binance/ws/threaded_stream.py b/binance/ws/threaded_stream.py
index e2fc575f9..4ee3254a4 100755
--- a/binance/ws/threaded_stream.py
+++ b/binance/ws/threaded_stream.py
@@ -15,6 +15,7 @@ def __init__(
         tld: str = "com",
         testnet: bool = False,
         session_params: Optional[Dict[str, Any]] = None,
+        https_proxy: Optional[str] = None,
         _loop: Optional[asyncio.AbstractEventLoop] = None,
     ):
         """Initialise the BinanceSocketManager"""
@@ -30,6 +31,7 @@ def __init__(
             "tld": tld,
             "testnet": testnet,
             "session_params": session_params,
+            "https_proxy": https_proxy,
         }
 
     async def _before_socket_listener_start(self): ...
diff --git a/setup.py b/setup.py
index fd1ded1f2..bc659471d 100644
--- a/setup.py
+++ b/setup.py
@@ -34,7 +34,6 @@
         "six",
         "dateparser",
         "aiohttp",
-        "ujson",
         "websockets",
         "pycryptodome",
     ],