diff --git a/aiosocks/__init__.py b/aiosocks/__init__.py index 3d9a365..9085401 100644 --- a/aiosocks/__init__.py +++ b/aiosocks/__init__.py @@ -1,226 +1,16 @@ import asyncio -import socket -import struct -from . import constants as c from .errors import * from .helpers import * +from .protocols import Socks4Protocol, Socks5Protocol __version__ = '0.1a' -__all__ = ('SocksProtocol', 'Socks4Protocol', 'Socks5Protocol', 'Socks4Auth', +__all__ = ('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): diff --git a/aiosocks/protocols.py b/aiosocks/protocols.py new file mode 100644 index 0000000..c7036a8 --- /dev/null +++ b/aiosocks/protocols.py @@ -0,0 +1,216 @@ +import asyncio +import socket +import struct +from . import constants as c +from .helpers import ( + Socks4Server, Socks5Server, Socks5Auth, Socks4Auth +) +from .errors import * + + +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