feat: support socks proxy

Author: joshuagruenstein

examples/example.py

@@ -1,5 +1,7 @@
 import asyncio
 
+from aiosocks import Socks5Addr
+
 from tcp_modbus_aio.client import TCPModbusClient
 from tcp_modbus_aio.exceptions import (
     ModbusCommunicationFailureError,
@@ -13,7 +15,11 @@
 
 async def example() -> None:
 
-    async with TCPModbusClient("192.168.250.207", enforce_pingable=False) as conn:
+    async with TCPModbusClient(
+        "192.168.250.207",
+        enforce_pingable=False,
+        socks_proxy_addr=Socks5Addr("localhost", 1080),
+    ) as conn:
         for _ in range(1000):
             for digital_in_coil in DIGITAL_IN_COILS:
                 example_message = ReadCoils()

poetry.lock

@@ -12,6 +12,7 @@ umodbus = "^1.0.4"
 cachetools = "^5.3.3"
 types-cachetools = "^5.3.0.7"
 typing-extensions = "^4.11.0"
+aiosocks = "^0.2.6"
 
 [tool.poetry.dev-dependencies]
 flake8 = "^6.0.0"

pyproject.toml

@@ -22,3 +22,6 @@ ignore_missing_imports = True
 
 [mypy-umodbus.*]
 ignore_missing_imports = True
+
+[mypy-aiosocks.*]
+ignore_missing_imports = True

setup.cfg

@@ -8,6 +8,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar
 
+import aiosocks
 from cachetools import TTLCache
 
 from tcp_modbus_aio.exceptions import (
@@ -50,7 +51,6 @@ class CoilWatchStatus:
 MBAP_HEADER_STRUCT_FORMAT = ">HHHB"
 
 
-@dataclass
 class TCPModbusClient:
     KEEPALIVE_AFTER_IDLE_SEC: ClassVar = 10
     KEEPALIVE_INTERVAL_SEC: ClassVar = 10
@@ -68,12 +68,16 @@ def __init__(
         logger: logging.Logger | None = None,
         enforce_pingable: bool = True,
         ping_timeout: float = 0.5,
+        socks_proxy_addr: aiosocks.SocksAddr | None = None,
+        socks_proxy_auth: aiosocks.Socks4Auth | aiosocks.Socks5Auth | None = None,
     ) -> None:
         self.host = host
         self.port = port
         self.slave_id = slave_id
         self.logger = logger
         self.ping_timeout = ping_timeout
+        self.socks_proxy_addr = socks_proxy_addr
+        self.socks_proxy_auth = socks_proxy_auth
 
         # If True, will throw an exception if attempting to send a request and the device is not pingable
         self.enforce_pingable = enforce_pingable
@@ -160,9 +164,19 @@ async def _get_tcp_connection(
             )
 
         try:
-            reader, writer = await asyncio.wait_for(
-                asyncio.open_connection(host=self.host, port=self.port), timeout
-            )
+            if self.socks_proxy_addr is None:
+                open_connection_coroutine = asyncio.open_connection(
+                    host=self.host, port=self.port
+                )
+            else:
+                open_connection_coroutine = aiosocks.open_connection(
+                    proxy=self.socks_proxy_addr,
+                    proxy_auth=self.socks_proxy_auth,
+                    dst=(self.host, self.port),
+                    remote_resolve=True,
+                )
+
+            reader, writer = await asyncio.wait_for(open_connection_coroutine, timeout)
         except asyncio.TimeoutError:
             msg = (
                 f"Timed out connecting to TCP modbus device at {self.host}:{self.port}"