From 3c2af6206f571aa3d302148c1221a0f46c7aa824 Mon Sep 17 00:00:00 2001 From: nibrag Date: Tue, 10 May 2016 03:10:56 +0300 Subject: [PATCH] flake8 --- aiosocks/__init__.py | 46 +++++++++++++++++----------- aiosocks/connector.py | 26 +++++++++------- aiosocks/constants.py | 6 ++-- aiosocks/protocols.py | 49 ++++++++++++++++++++--------- tests/test_connector.py | 12 ++++++-- tests/test_protocol.py | 68 +++++++++++++++++++++++++++-------------- 6 files changed, 135 insertions(+), 72 deletions(-) diff --git a/aiosocks/__init__.py b/aiosocks/__init__.py index 7cc8434..da207e9 100644 --- a/aiosocks/__init__.py +++ b/aiosocks/__init__.py @@ -1,26 +1,30 @@ import asyncio -from .errors import * -from .helpers import * +from .errors import * # noqa +from .helpers import * # noqa from .protocols import Socks4Protocol, Socks5Protocol __version__ = '0.1.2' __all__ = ('Socks4Protocol', 'Socks5Protocol', 'Socks4Auth', 'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'SocksError', - 'NoAcceptableAuthMethods', 'LoginAuthenticationFailed', 'SocksConnectionError', - 'InvalidServerVersion', 'InvalidServerReply', 'create_connection') + 'NoAcceptableAuthMethods', 'LoginAuthenticationFailed', + 'SocksConnectionError', 'InvalidServerVersion', + 'InvalidServerReply', 'create_connection') -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): - +@asyncio.coroutine +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, SocksAddr), ( 'proxy must be Socks4Addr() or Socks5Addr() 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 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))' @@ -28,11 +32,15 @@ async def create_connection(protocol_factory, proxy, proxy_auth, dst, *, remote_ if (isinstance(proxy, Socks4Addr) and not (proxy_auth is None or isinstance(proxy_auth, Socks4Auth))): - raise ValueError("proxy is Socks4Addr but proxy_auth is not Socks4Auth") + raise ValueError( + "proxy is Socks4Addr but proxy_auth is not Socks4Auth" + ) if (isinstance(proxy, Socks5Addr) and not (proxy_auth is None or isinstance(proxy_auth, Socks5Auth))): - raise ValueError("proxy is Socks5Addr but proxy_auth is not Socks5Auth") + raise ValueError( + "proxy is Socks5Addr but proxy_auth is not Socks5Auth" + ) loop = loop or asyncio.get_event_loop() @@ -47,16 +55,18 @@ async def create_connection(protocol_factory, proxy, proxy_auth, dst, *, remote_ remote_resolve=remote_resolve, loop=loop) try: - 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) + transport, protocol = yield from 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) except OSError as exc: - raise SocksConnectionError('[Errno %s] Can not connect to proxy %s:%d [%s]' % - (exc.errno, proxy.host, proxy.port, exc.strerror)) from exc + raise SocksConnectionError( + '[Errno %s] Can not connect to proxy %s:%d [%s]' % + (exc.errno, proxy.host, proxy.port, exc.strerror)) from exc # Wait until communication with proxy server is finished try: - await protocol.negotiate_done() + yield from protocol.negotiate_done() except SocksError as exc: raise SocksError('Can not connect to %s:%s [%s]' % (dst[0], dst[1], exc)) diff --git a/aiosocks/connector.py b/aiosocks/connector.py index 9d37942..29cb01a 100644 --- a/aiosocks/connector.py +++ b/aiosocks/connector.py @@ -10,8 +10,8 @@ __all__ = ('SocksConnector',) class SocksConnector(aiohttp.TCPConnector): - def __init__(self, proxy, proxy_auth=None, *, remote_resolve=True, **kwargs): - super().__init__(**kwargs) + def __init__(self, proxy, proxy_auth=None, *, remote_resolve=True, **kwgs): + super().__init__(**kwgs) self._proxy = proxy self._proxy_auth = proxy_auth @@ -44,7 +44,8 @@ class SocksConnector(aiohttp.TCPConnector): # It's aiohttp bug? Hot fix: try: ipaddress.ip_address(self._proxy.host) - proxy_hosts = await self._loop.getaddrinfo(self._proxy.host, self._proxy.port) + proxy_hosts = await self._loop.getaddrinfo(self._proxy.host, + self._proxy.port) family, _, proto, _, address = proxy_hosts[0] proxy_hosts = ({'hostname': self._proxy.host, @@ -52,16 +53,19 @@ class SocksConnector(aiohttp.TCPConnector): 'family': family, 'proto': proto, 'flags': socket.AI_NUMERICHOST},) except ValueError: - proxy_hosts = await self._resolve_host(self._proxy.host, self._proxy.port) + proxy_hosts = await self._resolve_host(self._proxy.host, + self._proxy.port) for hinfo in proxy_hosts: try: - proxy = self._proxy.__class__(host=hinfo['host'], port=hinfo['port']) + proxy = self._proxy.__class__(host=hinfo['host'], + port=hinfo['port']) transp, proto = await create_connection( - self._factory, proxy, self._proxy_auth, dst, loop=self._loop, - remote_resolve=self._remote_resolve, ssl=None, family=hinfo['family'], - proto=hinfo['proto'], flags=hinfo['flags'], local_addr=self._local_addr) + self._factory, proxy, self._proxy_auth, dst, + loop=self._loop, remote_resolve=self._remote_resolve, + ssl=None, family=hinfo['family'], proto=hinfo['proto'], + flags=hinfo['flags'], local_addr=self._local_addr) return transp, proto except (OSError, SocksError, SocksConnectionError) as e: @@ -72,6 +76,6 @@ class SocksConnector(aiohttp.TCPConnector): if isinstance(exc, SocksError): raise exc else: - raise aiohttp.ClientOSError(exc.errno, - 'Can not connect to %s:%s [%s]' % - (req.host, req.port, exc.strerror)) from exc + raise aiohttp.ClientOSError( + exc.errno, 'Can not connect to %s:%s [%s]' % + (req.host, req.port, exc.strerror)) from exc diff --git a/aiosocks/constants.py b/aiosocks/constants.py index e092290..19437d7 100644 --- a/aiosocks/constants.py +++ b/aiosocks/constants.py @@ -18,8 +18,10 @@ 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' + 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 = { diff --git a/aiosocks/protocols.py b/aiosocks/protocols.py index fa500a3..044cf92 100644 --- a/aiosocks/protocols.py +++ b/aiosocks/protocols.py @@ -5,13 +5,15 @@ from . import constants as c from .helpers import ( Socks4Addr, Socks5Addr, Socks5Auth, Socks4Auth ) -from .errors import * +from .errors import * # noqa class BaseSocksProtocol(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))') + raise ValueError( + 'Invalid dst format, tuple("dst_host", dst_port))' + ) self._proxy = proxy self._auth = proxy_auth @@ -53,9 +55,10 @@ class BaseSocksProtocol(asyncio.StreamReaderProtocol): 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) + 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] @@ -94,7 +97,8 @@ class Socks4Protocol(BaseSocksProtocol): 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] + req = [c.SOCKS_VER4, cmd, port_bytes, + host_bytes, self._auth.login, c.NULL] if include_hostname: req += [self._dst_host.encode('idna'), c.NULL] @@ -136,7 +140,9 @@ class Socks5Protocol(BaseSocksProtocol): resp = await self.read_response(3) if resp[0] != c.SOCKS_VER5: - raise InvalidServerVersion('SOCKS5 proxy server sent invalid version') + raise InvalidServerVersion( + 'SOCKS5 proxy server sent invalid version' + ) if resp[1] != c.SOCKS5_GRANTED: error = c.SOCKS5_ERRORS.get(resp[1], 'Unknown error') raise SocksError('[Errno {0:#04x}]: {1}'.format(resp[1], error)) @@ -148,7 +154,8 @@ class Socks5Protocol(BaseSocksProtocol): async def authenticate(self): # 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] + req = [c.SOCKS_VER5, 0x02, + c.SOCKS5_AUTH_ANONYMOUS, c.SOCKS5_AUTH_UNAME_PWD] else: req = [c.SOCKS_VER5, 0x01, c.SOCKS5_AUTH_ANONYMOUS] @@ -158,7 +165,9 @@ class Socks5Protocol(BaseSocksProtocol): chosen_auth = await self.read_response(2) if chosen_auth[0] != c.SOCKS_VER5: - raise InvalidServerVersion('SOCKS5 proxy server sent invalid version') + raise InvalidServerVersion( + 'SOCKS5 proxy server sent invalid version' + ) if chosen_auth[1] == c.SOCKS5_AUTH_UNAME_PWD: req = [0x01, chr(len(self._auth.login)).encode(), self._auth.login, @@ -167,18 +176,27 @@ class Socks5Protocol(BaseSocksProtocol): auth_status = await self.read_response(2) if auth_status[0] != 0x01: - raise InvalidServerReply('SOCKS5 proxy server sent invalid data') + raise InvalidServerReply( + 'SOCKS5 proxy server sent invalid data' + ) if auth_status[1] != c.SOCKS5_GRANTED: - raise LoginAuthenticationFailed('SOCKS5 authentication failed') + raise LoginAuthenticationFailed( + "SOCKS5 authentication failed" + ) # offered auth methods rejected elif chosen_auth[1] != c.SOCKS5_AUTH_ANONYMOUS: if chosen_auth[1] == c.SOCKS5_AUTH_NO_ACCEPTABLE_METHODS: - raise NoAcceptableAuthMethods('All offered SOCKS5 authentication methods were rejected') + raise NoAcceptableAuthMethods( + 'All offered SOCKS5 authentication methods were rejected' + ) else: - raise InvalidServerReply('SOCKS5 proxy server sent invalid data') + raise InvalidServerReply( + 'SOCKS5 proxy server sent invalid data' + ) async def write_address(self, host, port): - family_to_byte = {socket.AF_INET: c.SOCKS5_ATYP_IPv4, socket.AF_INET6: c.SOCKS5_ATYP_IPv6} + 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 @@ -195,7 +213,8 @@ class Socks5Protocol(BaseSocksProtocol): # 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] + 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) diff --git a/tests/test_connector.py b/tests/test_connector.py index 81fe6ad..c16bec1 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -41,7 +41,9 @@ class TestSocksConnector(unittest.TestCase): self.assertTrue(loop_mock.getaddrinfo.is_called) self.assertIs(conn._transport, tr) - self.assertTrue(isinstance(conn._protocol, aiohttp.parsers.StreamProtocol)) + self.assertTrue( + isinstance(conn._protocol, aiohttp.parsers.StreamProtocol) + ) conn.close() @@ -63,7 +65,9 @@ class TestSocksConnector(unittest.TestCase): self.assertTrue(connector._resolve_host.is_called) self.assertEqual(connector._resolve_host.call_count, 1) self.assertIs(conn._transport, tr) - self.assertTrue(isinstance(conn._protocol, aiohttp.parsers.StreamProtocol)) + self.assertTrue( + isinstance(conn._protocol, aiohttp.parsers.StreamProtocol) + ) conn.close() @@ -85,7 +89,9 @@ class TestSocksConnector(unittest.TestCase): self.assertTrue(connector._resolve_host.is_called) self.assertEqual(connector._resolve_host.call_count, 2) self.assertIs(conn._transport, tr) - self.assertTrue(isinstance(conn._protocol, aiohttp.parsers.StreamProtocol)) + self.assertTrue( + isinstance(conn._protocol, aiohttp.parsers.StreamProtocol) + ) conn.close() diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 58c3196..3e12656 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -67,7 +67,8 @@ class TestBaseSocksProtocol(unittest.TestCase): BaseSocksProtocol(None, None, ('python.org',), loop=self.loop) def test_write_request(self): - proto = BaseSocksProtocol(None, None, ('python.org', 80), loop=self.loop) + proto = BaseSocksProtocol(None, None, ('python.org', 80), + loop=self.loop) proto._transport = mock.Mock() proto.write_request([b'\x00', b'\x01\x02', 0x03]) @@ -97,10 +98,12 @@ class TestSocks4Protocol(unittest.TestCase): aiosocks.Socks4Protocol(None, auth, dst, loop=self.loop) with self.assertRaises(ValueError): - aiosocks.Socks4Protocol(aiosocks.Socks5Addr('host'), auth, dst, loop=self.loop) + aiosocks.Socks4Protocol(aiosocks.Socks5Addr('host'), auth, dst, + loop=self.loop) with self.assertRaises(ValueError): - aiosocks.Socks4Protocol(addr, aiosocks.Socks5Auth('l', 'p'), dst, loop=self.loop) + aiosocks.Socks4Protocol(addr, aiosocks.Socks5Auth('l', 'p'), dst, + loop=self.loop) aiosocks.Socks4Protocol(addr, None, dst, loop=self.loop) aiosocks.Socks4Protocol(addr, auth, dst, loop=self.loop) @@ -119,7 +122,8 @@ class TestSocks4Protocol(unittest.TestCase): ) # dst = domain, remote resolve = false - proto = make_socks4(self.loop, dst=('python.org', 80), rr=False, r=resp) + proto = make_socks4(self.loop, dst=('python.org', 80), + rr=False, r=resp) req = proto.socks_request(c.SOCKS_CMD_CONNECT) self.loop.run_until_complete(req) @@ -138,7 +142,8 @@ class TestSocks4Protocol(unittest.TestCase): ) # dst = ip, remote resolve = false - proto = make_socks4(self.loop, dst=('127.0.0.1', 8800), rr=False, r=resp) + proto = make_socks4(self.loop, dst=('127.0.0.1', 8800), + rr=False, r=resp) req = proto.socks_request(c.SOCKS_CMD_CONNECT) self.loop.run_until_complete(req) @@ -147,8 +152,8 @@ class TestSocks4Protocol(unittest.TestCase): ) # dst = domain, without user - proto = make_socks4( - self.loop, auth=aiosocks.Socks4Auth(''), dst=('python.org', 80), r=resp) + proto = make_socks4(self.loop, auth=aiosocks.Socks4Auth(''), + dst=('python.org', 80), r=resp) req = proto.socks_request(c.SOCKS_CMD_CONNECT) self.loop.run_until_complete(req) @@ -157,8 +162,8 @@ class TestSocks4Protocol(unittest.TestCase): ) # dst = ip, without user - proto = make_socks4( - self.loop, auth=aiosocks.Socks4Auth(''), dst=('127.0.0.1', 8800), r=resp) + proto = make_socks4(self.loop, auth=aiosocks.Socks4Auth(''), + dst=('127.0.0.1', 8800), r=resp) req = proto.socks_request(c.SOCKS_CMD_CONNECT) self.loop.run_until_complete(req) @@ -226,10 +231,12 @@ class TestSocks5Protocol(unittest.TestCase): aiosocks.Socks5Protocol(None, auth, dst, loop=self.loop) with self.assertRaises(ValueError): - aiosocks.Socks5Protocol(aiosocks.Socks4Addr('host'), auth, dst, loop=self.loop) + aiosocks.Socks5Protocol(aiosocks.Socks4Addr('host'), + auth, dst, loop=self.loop) with self.assertRaises(ValueError): - aiosocks.Socks5Protocol(addr, aiosocks.Socks4Auth('l'), dst, loop=self.loop) + aiosocks.Socks5Protocol(addr, aiosocks.Socks4Auth('l'), + dst, loop=self.loop) aiosocks.Socks5Protocol(addr, None, dst, loop=self.loop) aiosocks.Socks5Protocol(addr, auth, dst, loop=self.loop) @@ -264,8 +271,10 @@ class TestSocks5Protocol(unittest.TestCase): proto = make_socks5(self.loop, r=(b'\x05\x02', b'\x01\x00',)) req = proto.authenticate() self.loop.run_until_complete(req) - proto._transport.write.assert_has_calls( - [mock.call(b'\x05\x02\x00\x02'), mock.call(b'\x01\x04user\x03pwd')]) + proto._transport.write.assert_has_calls([ + mock.call(b'\x05\x02\x00\x02'), + mock.call(b'\x01\x04user\x03pwd') + ]) # invalid reply proto = make_socks5(self.loop, r=(b'\x05\x02', b'\x00\x00',)) @@ -289,7 +298,8 @@ class TestSocks5Protocol(unittest.TestCase): # ipv6 proto = make_socks5(self.loop) - req = proto.write_address('2001:0db8:11a3:09d7:1f34:8a2e:07a0:765d', 80) + req = proto.write_address( + '2001:0db8:11a3:09d7:1f34:8a2e:07a0:765d', 80) self.loop.run_until_complete(req) proto._transport.write.assert_called_with( @@ -311,22 +321,29 @@ class TestSocks5Protocol(unittest.TestCase): def test_read_address(self): # ipv4 - proto = make_socks5(self.loop, r=[b'\x01', b'\x7f\x00\x00\x01', b'\x00P']) + proto = make_socks5( + self.loop, r=[b'\x01', b'\x7f\x00\x00\x01', b'\x00P']) req = asyncio.ensure_future(proto.read_address(), loop=self.loop) self.loop.run_until_complete(req) self.assertEqual(req.result(), ('127.0.0.1', 80)) # ipv6 - proto = make_socks5( - self.loop, r=[b'\x04', b' \x01\r\xb8\x11\xa3\t\xd7\x1f4\x8a.\x07\xa0v]', b'\x00P']) + resp = [ + b'\x04', + b' \x01\r\xb8\x11\xa3\t\xd7\x1f4\x8a.\x07\xa0v]', + b'\x00P' + ] + proto = make_socks5(self.loop, r=resp) req = asyncio.ensure_future(proto.read_address(), loop=self.loop) self.loop.run_until_complete(req) - self.assertEqual(req.result(), ('2001:db8:11a3:9d7:1f34:8a2e:7a0:765d', 80)) + self.assertEqual( + req.result(), ('2001:db8:11a3:9d7:1f34:8a2e:7a0:765d', 80)) # domain - proto = make_socks5(self.loop, r=[b'\x03', b'\n', b'python.org', b'\x00P']) + proto = make_socks5( + self.loop, r=[b'\x03', b'\n', b'python.org', b'\x00P']) req = asyncio.ensure_future(proto.read_address(), loop=self.loop) self.loop.run_until_complete(req) @@ -345,7 +362,8 @@ class TestSocks5Protocol(unittest.TestCase): with self.assertRaises(aiosocks.SocksError) as ct: self.loop.run_until_complete(req) - self.assertTrue('Connection not allowed by ruleset' in str(ct.exception)) + self.assertTrue( + 'Connection not allowed by ruleset' in str(ct.exception)) # socks unknown error proto = make_socks5(self.loop, r=[b'\x05\x00', b'\x05\xFF\x00']) @@ -356,9 +374,13 @@ class TestSocks5Protocol(unittest.TestCase): self.assertTrue('Unknown error' in str(ct.exception)) # cmd granted - proto = make_socks5( - self.loop, r=[b'\x05\x00', b'\x05\x00\x00', b'\x01', b'\x7f\x00\x00\x01', b'\x00P']) - req = asyncio.ensure_future(proto.socks_request(c.SOCKS_CMD_CONNECT), loop=self.loop) + resp = [b'\x05\x00', + b'\x05\x00\x00', + b'\x01', b'\x7f\x00\x00\x01', + b'\x00P'] + proto = make_socks5(self.loop, r=resp) + req = asyncio.ensure_future(proto.socks_request(c.SOCKS_CMD_CONNECT), + loop=self.loop) self.loop.run_until_complete(req) self.assertEqual(req.result(), (('python.org', 80), ('127.0.0.1', 80)))