SOCKS "CONNECT" command realization

aiosocks/__init__.py

@@ -0,0 +1,268 @@
+import asyncio
+import socket
+import struct
+from . import constants as c
+from .errors import *
+from .helpers import *
+
+__version__ = '0.1a'
+
+__all__ = ('SocksProtocol', 'Socks4Protocol', 'Socks5Protocol', 'Socks4Auth',
+           'Socks5Auth', 'Socks4Server', 'Socks5Server', 'SocksError',
+           'NoAcceptableAuthMethods', 'LoginAuthenticationFailed',
+           'InvalidServerVersion', 'InvalidServerReply', 'create_connection')
+
+
+class SocksProtocol(asyncio.StreamReaderProtocol):
+    def __init__(self, proxy, proxy_auth, dst, remote_resolve=True, loop=None):
+        if not isinstance(dst, (tuple, list)) or len(dst) != 2:
+            raise ValueError('Invalid dst format, tuple("dst_host", dst_port))')
+
+        self._proxy = proxy
+        self._auth = proxy_auth
+        self._dst_host, self._dst_port = dst
+        self._remote_resolve = remote_resolve
+
+        self._loop = loop or asyncio.get_event_loop()
+        self._transport = None
+
+        self._negotiate_done = None
+
+        reader = asyncio.StreamReader(loop=self._loop)
+
+        super().__init__(stream_reader=reader, loop=self._loop)
+
+    def connection_made(self, transport):
+        super().connection_made(transport)
+        self._transport = transport
+
+        req_coro = self.socks_request(c.SOCKS_CMD_CONNECT)
+        self._negotiate_done = asyncio.ensure_future(req_coro, loop=self._loop)
+
+    async def socks_request(self, cmd):
+        raise NotImplementedError
+
+    def write_request(self, request):
+        bdata = bytearray()
+
+        for item in request:
+            if isinstance(item, int):
+                bdata.append(item)
+            elif isinstance(item, (bytearray, bytes)):
+                bdata += item
+            else:
+                raise ValueError('Unsupported item')
+
+        self._transport.write(bdata)
+
+    async def read_response(self, n):
+        return await self._stream_reader.read(n)
+
+    async def _get_dst_addr(self):
+        infos = await self._loop.getaddrinfo(self._dst_host, self._dst_port,
+                                             family=socket.AF_UNSPEC, type=socket.SOCK_STREAM,
+                                             proto=socket.IPPROTO_TCP, flags=socket.AI_ADDRCONFIG)
+        if not infos:
+            raise OSError('getaddrinfo() returned empty list')
+        return infos[0][0], infos[0][4][0]
+
+    async def negotiate_done(self):
+        return await self._negotiate_done
+
+
+class Socks4Protocol(SocksProtocol):
+    def __init__(self, proxy, proxy_auth, dst, remote_resolve=True, loop=None):
+        if not isinstance(proxy, Socks4Server):
+            raise ValueError('Invalid proxy format')
+
+        if proxy_auth is not None and not isinstance(proxy_auth, Socks4Auth):
+            raise ValueError('Invalid proxy_auth format')
+
+        super().__init__(proxy, proxy_auth, dst, remote_resolve, loop)
+
+    async def socks_request(self, cmd):
+        # prepare destination addr/port
+        host, port = self._dst_host, self._dst_port
+        port_bytes = struct.pack(b'>H', port)
+        try:
+            host_bytes = socket.inet_aton(host)
+        except socket.error:
+            if self._remote_resolve:
+                host_bytes = bytes([c.NULL, c.NULL, c.NULL, 0x01])
+            else:
+                # it's not an IP number, so it's probably a DNS name.
+                family, host = await self._get_dst_addr()
+                host_bytes = socket.inet_aton(host)
+
+        # build and send connect command
+        req = [c.SOCKS_VER4, cmd, port_bytes, host_bytes, self._auth.login, c.NULL]
+        if self._remote_resolve:
+            req += [self._dst_host.encode('idna'), c.NULL]
+
+        self.write_request(req)
+
+        # read/process result
+        resp = await self.read_response(8)
+
+        if resp[0] != c.NULL:
+            raise InvalidServerReply('SOCKS4 proxy server sent invalid data')
+        if resp[1] != c.SOCKS4_GRANTED:
+            error = c.SOCKS4_ERRORS.get(resp[1], 'Unknown error')
+            raise SocksError('{0:#04x}: {1}'.format(resp[1], error))
+
+        binded = socket.inet_ntoa(resp[4:]), struct.unpack('>H', resp[2:4])[0]
+        return (host, port), binded
+
+
+class Socks5Protocol(SocksProtocol):
+    def __init__(self, proxy, proxy_auth, dst, remote_resolve=True, loop=None):
+        if not isinstance(proxy, Socks5Server):
+            raise ValueError('Invalid proxy format')
+
+        if proxy_auth is not None and not isinstance(proxy_auth, Socks5Auth):
+            raise ValueError('Invalid proxy_auth format')
+
+        super().__init__(proxy, proxy_auth, dst, remote_resolve, loop)
+
+    async def socks_request(self, cmd):
+        # send available auth methods
+        if self._auth.login and self._auth.password:
+            req = [c.SOCKS_VER5, 0x02, c.SOCKS5_AUTH_ANONYMOUS, c.SOCKS5_AUTH_UNAME_PWD]
+        else:
+            req = [c.SOCKS_VER5, 0x01, c.SOCKS5_AUTH_ANONYMOUS]
+
+        self.write_request(req)
+
+        # read/process response and send auth data if necessary
+        chosen_auth = await self.read_response(2)
+
+        if chosen_auth[0] != c.SOCKS_VER5:
+            raise InvalidServerVersion
+
+        if chosen_auth[1] == c.SOCKS5_AUTH_UNAME_PWD:
+            req = [0x01, chr(len(self._auth.login)).encode(), self._auth.login,
+                   chr(len(self._auth.password)).encode(), self._auth.password]
+            self.write_request(req)
+
+            auth_status = await self.read_response(2)
+            if auth_status[0] != 0x01:
+                raise InvalidServerReply('SOCKS5 proxy server sent invalid data')
+            if auth_status[1] != c.SOCKS5_GRANTED:
+                raise LoginAuthenticationFailed
+        # offered auth methods rejected
+        elif chosen_auth[1] != c.SOCKS5_AUTH_ANONYMOUS:
+            if chosen_auth[1] == c.SOCKS5_AUTH_NO_ACCEPTABLE_METHODS:
+                raise NoAcceptableAuthMethods
+            else:
+                raise InvalidServerReply('SOCKS5 proxy server sent invalid data')
+
+        # build and send command
+        self.write_request([c.SOCKS_VER5, cmd, c.RSV])
+        resolved = await self.write_address(self._dst_host, self._dst_port)
+
+        # read/process command response
+        resp = await self.read_response(3)
+
+        if resp[0] != c.SOCKS_VER5:
+            raise InvalidServerVersion
+        if resp[1] != c.SOCKS5_GRANTED:
+            error = c.SOCKS5_ERRORS.get(resp[1], 'Unknown error')
+            raise SocksError('{0:#04x}: {1}'.format(resp[1], error))
+
+        binded = await self.read_address()
+
+        return resolved, binded
+
+    async def write_address(self, host, port):
+        family_to_byte = {socket.AF_INET: c.SOCKS5_ATYP_IPv4, socket.AF_INET6: c.SOCKS5_ATYP_IPv6}
+        port_bytes = struct.pack('>H', port)
+
+        # if the given destination address is an IP address, we will
+        # use the IP address request even if remote resolving was specified.
+        for family in (socket.AF_INET, socket.AF_INET6):
+            try:
+                host_bytes = socket.inet_pton(family, host)
+                req = [family_to_byte[family], host_bytes, port_bytes]
+                self.write_request(req)
+                return host, port
+            except socket.error:
+                pass
+
+        # it's not an IP number, so it's probably a DNS name.
+        if self._remote_resolve:
+            host_bytes = host.encode('idna')
+            req = [c.SOCKS5_ATYP_DOMAIN, chr(len(host_bytes)).encode(), host_bytes, port_bytes]
+        else:
+            family, host_bytes = await self._get_dst_addr()
+            host_bytes = socket.inet_pton(family, host_bytes)
+            req = [family_to_byte[family], host_bytes, port_bytes]
+            host = socket.inet_ntop(family, host_bytes)
+
+        self.write_request(req)
+        return host, port
+
+    async def read_address(self):
+        atype = await self.read_response(1)
+
+        if atype[0] == c.SOCKS5_ATYP_IPv4:
+            addr = socket.inet_ntoa(await self.read_response(4))
+        elif atype[0] == c.SOCKS5_ATYP_DOMAIN:
+            length = await self.read_response(1)
+            addr = await self.read_response(ord(length))
+        elif atype[0] == c.SOCKS5_ATYP_IPv6:
+            addr = await self.read_response(16)
+            addr = socket.inet_ntop(socket.AF_INET6, addr)
+        else:
+            raise InvalidServerReply('SOCKS5 proxy server sent invalid data')
+
+        port = await self.read_response(2)
+        port = struct.unpack('>H', port)[0]
+
+        return addr, port
+
+
+async def create_connection(protocol_factory, proxy, proxy_auth, dst, *, remote_resolve=True,
+                            loop=None, ssl=None, family=0, proto=0, flags=0, sock=None,
+                            local_addr=None, server_hostname=None):
+
+    assert isinstance(proxy, SocksServer), (
+        'proxy must be Socks4Server() or Socks5Server() tuple'
+    )
+
+    assert proxy_auth is None or isinstance(proxy_auth, (Socks4Auth, Socks5Auth)), (
+        'proxy_auth must be None or Socks4Auth() or Socks5Auth() tuple', proxy_auth
+    )
+    assert isinstance(dst, (tuple, list)) and len(dst) == 2, (
+        'invalid dst format, tuple("dst_host", dst_port))'
+    )
+
+    if (isinstance(proxy, Socks4Server) and not
+            (proxy_auth is None or isinstance(proxy_auth, Socks4Auth))):
+        raise ValueError("proxy is Socks4Server but proxy_auth is not Socks4Auth")
+
+    if (isinstance(proxy, Socks5Server) and not
+            (proxy_auth is None or isinstance(proxy_auth, Socks5Auth))):
+        raise ValueError("proxy is Socks5Server but proxy_auth is not Socks5Auth")
+
+    loop = loop or asyncio.get_event_loop()
+
+    def socks_factory():
+        if isinstance(proxy, Socks4Server):
+            socks_proto = Socks4Protocol
+        else:
+            socks_proto = Socks5Protocol
+
+        return socks_proto(
+            proxy=proxy, proxy_auth=proxy_auth, dst=dst,
+            remote_resolve=remote_resolve, loop=loop)
+
+    transport, protocol = await loop.create_connection(
+        socks_factory, proxy.host, proxy.port, ssl=ssl, family=family, proto=proto,
+        flags=flags, sock=sock, local_addr=local_addr, server_hostname=server_hostname)
+
+    await protocol.negotiate_done()
+
+    sock = transport.get_extra_info('socket')
+
+    return await loop._create_connection_transport(
+        sock, protocol_factory, ssl, server_hostname)

aiosocks/constants.py

@@ -0,0 +1,34 @@
+RSV = NULL = 0x00
+SOCKS_VER4 = 0x04
+SOCKS_VER5 = 0x05
+
+SOCKS_CMD_CONNECT = 0x01
+SOCKS_CMD_BIND = 0x02
+SOCKS_CMD_UDP_ASSOCIATE = 0x03
+SOCKS4_GRANTED = 0x5A
+SOCKS5_GRANTED = 0x00
+
+SOCKS5_AUTH_ANONYMOUS = 0x00
+SOCKS5_AUTH_UNAME_PWD = 0x02
+SOCKS5_AUTH_NO_ACCEPTABLE_METHODS = 0xFF
+
+SOCKS5_ATYP_IPv4 = 0x01
+SOCKS5_ATYP_DOMAIN = 0x03
+SOCKS5_ATYP_IPv6 = 0x04
+
+SOCKS4_ERRORS = {
+    0x5B: 'Request rejected or failed',
+    0x5C: 'Request rejected because SOCKS server cannot connect to identd on the client',
+    0x5D: 'Request rejected because the client program and identd report different user-ids'
+}
+
+SOCKS5_ERRORS = {
+    0x01: 'General SOCKS server failure',
+    0x02: 'Connection not allowed by ruleset',
+    0x03: 'Network unreachable',
+    0x04: 'Host unreachable',
+    0x05: 'Connection refused',
+    0x06: 'TTL expired',
+    0x07: 'Command not supported, or protocol error',
+    0x08: 'Address type not supported'
+}

aiosocks/errors.py

@@ -0,0 +1,18 @@
+class SocksError(Exception):
+    pass
+
+
+class NoAcceptableAuthMethods(SocksError):
+    pass
+
+
+class LoginAuthenticationFailed(SocksError):
+    pass
+
+
+class InvalidServerVersion(SocksError):
+    pass
+
+
+class InvalidServerReply(SocksError):
+    pass

aiosocks/helpers.py

@@ -0,0 +1,43 @@
+from collections import namedtuple
+
+__all__ = ('Socks4Auth', 'Socks5Auth', 'Socks4Server', 'Socks5Server', 'SocksServer')
+
+
+class Socks4Auth(namedtuple('Socks4Auth', ['login', 'encoding'])):
+    def __new__(cls, login, encoding='utf-8'):
+        if login is None:
+            raise ValueError('None is not allowed as login value')
+
+        return super().__new__(cls, login.encode(encoding), encoding)
+
+
+class Socks5Auth(namedtuple('Socks5Auth', ['login', 'password', 'encoding'])):
+    def __new__(cls, login, password, encoding='utf-8'):
+        if login is None:
+            raise ValueError('None is not allowed as login value')
+
+        if password is None:
+            raise ValueError('None is not allowed as password value')
+
+        return super().__new__(cls,
+                               login.encode(encoding),
+                               password.encode(encoding), encoding)
+
+
+class SocksServer(namedtuple('SocksServer', ['host', 'port'])):
+    def __new__(cls, host, port):
+        if host is None:
+            raise ValueError('None is not allowed as host value')
+
+        if port is None:
+            port = 1080  # default socks server port
+
+        return super().__new__(cls, host, port)
+
+
+class Socks4Server(SocksServer):
+    pass
+
+
+class Socks5Server(SocksServer):
+    pass