flake8

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))

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

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 = {

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)

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()
 

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)))