connect upstream socks5

proxy/core/connection/__init__.py

@@ -16,14 +16,15 @@
 from .pool import UpstreamConnectionPool
 from .types import tcpConnectionTypes
 from .client import TcpClientConnection
-from .server import TcpServerConnection
+from .server import TcpServerConnection, Socks5ServerConnection
 from .connection import TcpConnection, TcpConnectionUninitializedException
 
 
 __all__ = [
     'TcpConnection',
     'TcpConnectionUninitializedException',
     'TcpServerConnection',
+    'Socks5ServerConnection',
     'TcpClientConnection',
     'tcpConnectionTypes',
     'UpstreamConnectionPool',

proxy/core/connection/server.py

@@ -9,14 +9,22 @@
     :license: BSD, see LICENSE for more details.
 """
 import ssl
+import socket
+import logging
 from typing import Optional
 
+import socks
+
 from .types import tcpConnectionTypes
 from .connection import TcpConnection, TcpConnectionUninitializedException
 from ...common.types import HostPort, TcpOrTlsSocket
 from ...common.utils import new_socket_connection
-from ...common.constants import DEFAULT_SSL_CONTEXT_OPTIONS
+from ...common.constants import (
+    DEFAULT_BUFFER_SIZE, DEFAULT_MAX_SEND_SIZE, DEFAULT_SSL_CONTEXT_OPTIONS,
+)
+
 
+logger = logging.getLogger(__name__)
 
 class TcpServerConnection(TcpConnection):
     """A buffered server connection object."""
@@ -66,3 +74,58 @@ def wrap(
         )
         if as_non_blocking:
             self.connection.setblocking(False)
+
+class Socks5ServerConnection(TcpConnection):
+    """A buffered SOCKS5 upstream connection object."""
+
+    def __init__(
+            self,
+            host: str,
+            port: int,
+            username: Optional[str] = None,
+            password: Optional[str] = None,
+    ) -> None:
+        super().__init__(tcpConnectionTypes.SERVER)
+        self._conn: Optional[TcpOrTlsSocket] = None
+        self.addr: HostPort = (host, port, username, password)
+        self.closed = True
+
+    @property
+    def connection(self) -> TcpOrTlsSocket:
+        if self._conn is None:
+            raise TcpConnectionUninitializedException()
+        return self._conn
+
+    def connect(
+            self,
+            addr: Optional[HostPort] = None,
+            source_address: Optional[HostPort] = None,
+    ) -> None:
+        assert self._conn is None
+
+        sock = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.set_proxy(
+            proxy_type=socks.SOCKS5,
+            addr=self.addr[0],
+            port=self.addr[1],
+            username=self.addr[2] if len(self.addr) > 2 else None,
+            password=self.addr[3] if len(self.addr) > 3 else None,
+        )
+        sock.connect(addr)
+        sock.setblocking(False)
+
+        self._conn = sock
+        self.closed = False
+
+    def recv(
+            self,
+            buffer_size: int = DEFAULT_BUFFER_SIZE,
+    ) -> Optional[memoryview]:
+        try:
+            data: bytes = self._conn.recv(buffer_size)
+        except BlockingIOError:
+            return memoryview(b'')
+        if len(data) == 0:
+            return None
+        logger.debug('received %d bytes from server', len(data))
+        return memoryview(data)

setup.cfg

@@ -107,6 +107,7 @@ include_package_data = True
 zip_safe = False
 
 install_requires =
+  PySocks>=1.7.1
 
 [options.extras_require]
 metrics =

tests/core/test_connection.py

@@ -16,7 +16,8 @@
 from proxy.common.types import TcpOrTlsSocket
 from proxy.core.connection import (
     TcpConnection, TcpClientConnection, TcpServerConnection,
-    TcpConnectionUninitializedException, tcpConnectionTypes,
+    Socks5ServerConnection, TcpConnectionUninitializedException,
+    tcpConnectionTypes,
 )
 from proxy.common.constants import (
     DEFAULT_PORT, DEFAULT_IPV4_HOSTNAME, DEFAULT_IPV6_HOSTNAME,
@@ -136,3 +137,207 @@ def testTcpClientRaisesTcpConnectionUninitializedException(self) -> None:
         conn._conn = None
         with self.assertRaises(TcpConnectionUninitializedException):
             _ = conn.connection
+
+
+class TestSocks5ServerConnection(unittest.TestCase):
+
+    def testInitWithBasicParams(self) -> None:
+        conn = Socks5ServerConnection('127.0.0.1', 1080)
+        self.assertEqual(conn.addr, ('127.0.0.1', 1080, None, None))
+        self.assertTrue(conn.closed)
+        self.assertEqual(conn.tag, 'server')
+
+    def testInitWithAuthParams(self) -> None:
+        conn = Socks5ServerConnection(
+            '127.0.0.1', 1080,
+            username='user', password='pass',
+        )
+        self.assertEqual(conn.addr, ('127.0.0.1', 1080, 'user', 'pass'))
+        self.assertTrue(conn.closed)
+
+    def testRaisesTcpConnectionUninitializedException(self) -> None:
+        conn = Socks5ServerConnection('127.0.0.1', 1080)
+        with self.assertRaises(TcpConnectionUninitializedException):
+            _ = conn.connection
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testConnectWithoutAuth(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        target_addr = ('target.example.com', 80)
+        conn.connect(addr=target_addr)
+
+        mock_socksocket.assert_called_once()
+        mock_sock.set_proxy.assert_called_once_with(
+            proxy_type=2,
+            addr='proxy.example.com',
+            port=1080,
+            username=None,
+            password=None,
+        )
+        mock_sock.connect.assert_called_once_with(target_addr)
+        mock_sock.setblocking.assert_called_once_with(False)
+        self.assertFalse(conn.closed)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testConnectWithAuth(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection(
+            'proxy.example.com', 1080,
+            username='testuser', password='testpass',
+        )
+        target_addr = ('target.example.com', 443)
+        conn.connect(addr=target_addr)
+
+        mock_sock.set_proxy.assert_called_once_with(
+            proxy_type=2,
+            addr='proxy.example.com',
+            port=1080,
+            username='testuser',
+            password='testpass',
+        )
+        self.assertFalse(conn.closed)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testConnectWillNotIgnoreDoubleConnectAttempts(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        with self.assertRaises(AssertionError):
+            conn.connect(addr=('target.example.com', 80))
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testConnectionProperty(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        self.assertEqual(conn.connection, mock_sock)
+
+    def testClose(self) -> None:
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn._conn = mock.MagicMock()
+        conn.closed = False
+
+        conn.close()
+        conn._conn.close.assert_called_once()
+        self.assertTrue(conn.closed)
+
+    def testCloseNoOpIfAlreadyClosed(self) -> None:
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn._conn = mock.MagicMock()
+        conn.closed = True
+
+        conn.close()
+        conn._conn.close.assert_not_called()
+        self.assertTrue(conn.closed)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testRecvReturnsData(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_sock.recv.return_value = b'test data'
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        result = conn.recv()
+        self.assertIsNotNone(result)
+        self.assertEqual(result.tobytes(), b'test data')
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testRecvReturnsEmptyOnBlockingIOError(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_sock.recv.side_effect = BlockingIOError()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        result = conn.recv()
+        self.assertIsNotNone(result)
+        self.assertEqual(result.tobytes(), b'')
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testRecvReturnsNoneOnEmptyData(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_sock.recv.return_value = b''
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        result = conn.recv()
+        self.assertIsNone(result)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testSend(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_sock.send.return_value = 10
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        result = conn.send(b'test data')
+        self.assertEqual(result, 10)
+        mock_sock.send.assert_called_once_with(b'test data')
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testSendWithMemoryView(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_sock.send.return_value = 9
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        data = memoryview(b'test data')
+        result = conn.send(data)
+        self.assertEqual(result, 9)
+        mock_sock.send.assert_called_once_with(data)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testRecvWithCustomBufferSize(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_sock.recv.return_value = b'x' * 4096
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect(addr=('target.example.com', 80))
+
+        result = conn.recv(buffer_size=4096)
+        self.assertIsNotNone(result)
+        self.assertEqual(len(result), 4096)
+        mock_sock.recv.assert_called_once_with(4096)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testConnectWithSourceAddress(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        source_addr = ('192.168.1.100', 0)
+        target_addr = ('target.example.com', 80)
+        conn.connect(addr=target_addr, source_address=source_addr)
+
+        mock_sock.connect.assert_called_once_with(target_addr)
+
+    @mock.patch('proxy.core.connection.server.socks.socksocket')
+    def testConnectWithoutAddrParameter(self, mock_socksocket: mock.Mock) -> None:
+        mock_sock = mock.MagicMock()
+        mock_socksocket.return_value = mock_sock
+
+        conn = Socks5ServerConnection('proxy.example.com', 1080)
+        conn.connect()
+
+        mock_sock.connect.assert_called_once()