jmg
/
aiosocks
1
0
Fork 0
Code Issues 0 Pull Requests 0 Releases 0 Wiki Activity
Browse Source

flake8

main
nibrag 8 years ago
parent
0064efbab2
commit
3c2af6206f
6 changed files with 135 additions and 72 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +28
    -18
      aiosocks/__init__.py
  2. +15
    -11
      aiosocks/connector.py
  3. +4
    -2
      aiosocks/constants.py
  4. +34
    -15
      aiosocks/protocols.py
  5. +9
    -3
      tests/test_connector.py
  6. +45
    -23
      tests/test_protocol.py

+ 28
- 18
aiosocks/__init__.py View File

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


+ 15
- 11
aiosocks/connector.py View File

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

+ 4
- 2
aiosocks/constants.py View File

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


+ 34
- 15
aiosocks/protocols.py View File

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


+ 9
- 3
tests/test_connector.py View File

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



+ 45
- 23
tests/test_protocol.py View File

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


Write Preview
Loading…
Cancel
Save