@@ -11,6 +11,12 @@ SOCKS proxy client for asyncio and aiohttp | |||||
.. image:: https://badge.fury.io/py/aiosocks.svg | .. image:: https://badge.fury.io/py/aiosocks.svg | ||||
:target: https://badge.fury.io/py/aiosocks | :target: https://badge.fury.io/py/aiosocks | ||||
Dependencies | |||||
------------ | |||||
python 3.5+ | |||||
aiohttp 2.0+ | |||||
Features | Features | ||||
-------- | -------- | ||||
- SOCKS4, SOCKS4a and SOCKS5 version | - SOCKS4, SOCKS4a and SOCKS5 version | ||||
@@ -136,56 +142,47 @@ aiohttp usage | |||||
import asyncio | import asyncio | ||||
import aiohttp | import aiohttp | ||||
import aiosocks | import aiosocks | ||||
from aiosocks.connector import ( | |||||
SocksConnector, HttpProxyAddr, HttpProxyAuth | |||||
) | |||||
from yarl import URL | |||||
from aiosocks.connector import ProxyConnecotr, ProxyClientRequest | |||||
async def load_github_main(): | async def load_github_main(): | ||||
addr = aiosocks.Socks5Addr('127.0.0.1', 1080) | |||||
auth = aiosocks.Socks5Auth('proxyuser1', password='pwd') | |||||
auth5 = aiosocks.Socks5Auth('proxyuser1', password='pwd') | |||||
auth4 = aiosocks.Socks4Auth('proxyuser1') | |||||
ba = aiohttp.BasicAuth('login') | |||||
# remote resolve | # remote resolve | ||||
conn = SocksConnector(proxy=addr, proxy_auth=auth, remote_resolve=True) | |||||
conn = ProxyConnecotr(remote_resolve=True) | |||||
# or locale resolve | # or locale resolve | ||||
conn = SocksConnector(proxy=addr, proxy_auth=auth, remote_resolve=False) | |||||
conn = SocksConnector(remote_resolve=False) | |||||
try: | try: | ||||
with aiohttp.ClientSession(connector=conn) as session: | |||||
async with session.get('http://github.com/') as resp: | |||||
with aiohttp.ClientSession(connector=conn, request_class=ProxyClientRequest) as session: | |||||
# socks5 proxy | |||||
async with session.get('http://github.com/', proxy=URL('socks5://127.0.0.1:1080'), | |||||
proxy_auth=auth5) as resp: | |||||
if resp.status == 200: | |||||
print(await resp.text()) | |||||
# socks4 proxy | |||||
async with session.get('http://github.com/', proxy=URL('socks4://127.0.0.1:1081'), | |||||
proxy_auth=auth4) as resp: | |||||
if resp.status == 200: | |||||
print(await resp.text()) | |||||
# http proxy | |||||
async with session.get('http://github.com/', proxy=URL('http://127.0.0.1:8080'), | |||||
proxy_auth=ba) as resp: | |||||
if resp.status == 200: | if resp.status == 200: | ||||
print(await resp.text()) | print(await resp.text()) | ||||
except aiohttp.ProxyConnectionError: | except aiohttp.ProxyConnectionError: | ||||
# connection problem | # connection problem | ||||
except aiosocks.SocksError: | except aiosocks.SocksError: | ||||
# communication problem | # communication problem | ||||
if __name__ == '__main__': | if __name__ == '__main__': | ||||
loop = asyncio.get_event_loop() | loop = asyncio.get_event_loop() | ||||
loop.run_until_complete(load_github_main()) | loop.run_until_complete(load_github_main()) | ||||
loop.close() | loop.close() | ||||
proxy_connector | |||||
^^^^^^^^^^^^^^^ | |||||
A unified method to create `connector`. | |||||
.. code-block:: python | |||||
import asyncio | |||||
import aiohttp | |||||
import aiosocks | |||||
from aiosocks.connector import ( | |||||
proxy_connector, HttpProxyAddr, HttpProxyAuth | |||||
) | |||||
# make SocksConnector | |||||
conn = proxy_connector(aiosocks.Socks5Addr(...), | |||||
remote_resolve=True, verify_ssl=False) | |||||
# return SocksConnector instance | |||||
# make aiohttp.ProxyConnector (http proxy) | |||||
conn = proxy_connector(HttpProxyAddr('http://proxy'), | |||||
HttpProxyAuth('login', 'pwd'), verify_ssl=True) | |||||
# return aiohttp.ProxyConnector instance |
@@ -17,11 +17,10 @@ __all__ = ('Socks4Protocol', 'Socks5Protocol', 'Socks4Auth', | |||||
'InvalidServerReply', 'create_connection', 'open_connection') | 'InvalidServerReply', 'create_connection', 'open_connection') | ||||
@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, reader_limit=DEFAULT_LIMIT): | |||||
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, reader_limit=DEFAULT_LIMIT): | |||||
assert isinstance(proxy, SocksAddr), ( | assert isinstance(proxy, SocksAddr), ( | ||||
'proxy must be Socks4Addr() or Socks5Addr() tuple' | 'proxy must be Socks4Addr() or Socks5Addr() tuple' | ||||
) | ) | ||||
@@ -70,7 +69,7 @@ def create_connection(protocol_factory, proxy, proxy_auth, dst, *, | |||||
reader_limit=reader_limit) | reader_limit=reader_limit) | ||||
try: | try: | ||||
transport, protocol = yield from loop.create_connection( | |||||
transport, protocol = await loop.create_connection( | |||||
socks_factory, proxy.host, proxy.port, family=family, | socks_factory, proxy.host, proxy.port, family=family, | ||||
proto=proto, flags=flags, sock=sock, local_addr=local_addr) | proto=proto, flags=flags, sock=sock, local_addr=local_addr) | ||||
except OSError as exc: | except OSError as exc: | ||||
@@ -79,7 +78,7 @@ def create_connection(protocol_factory, proxy, proxy_auth, dst, *, | |||||
(exc.errno, proxy.host, proxy.port, exc.strerror)) from exc | (exc.errno, proxy.host, proxy.port, exc.strerror)) from exc | ||||
try: | try: | ||||
yield from waiter | |||||
await waiter | |||||
except: | except: | ||||
transport.close() | transport.close() | ||||
raise | raise | ||||
@@ -87,10 +86,9 @@ def create_connection(protocol_factory, proxy, proxy_auth, dst, *, | |||||
return protocol.app_transport, protocol.app_protocol | return protocol.app_transport, protocol.app_protocol | ||||
@asyncio.coroutine | |||||
def open_connection(proxy, proxy_auth, dst, *, remote_resolve=True, | |||||
loop=None, limit=DEFAULT_LIMIT, **kwds): | |||||
_, protocol = yield from create_connection( | |||||
async def open_connection(proxy, proxy_auth, dst, *, remote_resolve=True, | |||||
loop=None, limit=DEFAULT_LIMIT, **kwds): | |||||
_, protocol = await create_connection( | |||||
None, proxy, proxy_auth, dst, reader_limit=limit, | None, proxy, proxy_auth, dst, reader_limit=limit, | ||||
remote_resolve=remote_resolve, loop=loop, **kwds) | remote_resolve=remote_resolve, loop=loop, **kwds) | ||||
@@ -1,105 +1,99 @@ | |||||
try: | try: | ||||
import aiohttp | import aiohttp | ||||
from aiohttp.errors import ProxyConnectionError | |||||
from aiohttp.helpers import BasicAuth as HttpProxyAuth | |||||
from aiohttp.connector import sentinel | |||||
except ImportError: | except ImportError: | ||||
raise ImportError('aiosocks.SocksConnector require aiohttp library') | raise ImportError('aiosocks.SocksConnector require aiohttp library') | ||||
import asyncio | |||||
from collections import namedtuple | |||||
from .errors import SocksError, SocksConnectionError | from .errors import SocksError, SocksConnectionError | ||||
from .helpers import SocksAddr | |||||
from .helpers import Socks4Auth, Socks5Auth, Socks4Addr, Socks5Addr | |||||
from . import create_connection | from . import create_connection | ||||
__all__ = ('SocksConnector', 'HttpProxyAddr', 'HttpProxyAuth') | |||||
__all__ = ('ProxyConnector', 'ProxyClientRequest') | |||||
class ProxyClientRequest(aiohttp.ClientRequest): | |||||
def update_proxy(self, proxy, proxy_auth): | |||||
if proxy and proxy.scheme not in ['http', 'socks4', 'socks5']: | |||||
raise ValueError( | |||||
"Only http, socks4 and socks5 proxies are supported") | |||||
if proxy and proxy_auth: | |||||
if proxy.scheme == 'http' and \ | |||||
not isinstance(proxy_auth, aiohttp.BasicAuth): | |||||
raise ValueError("proxy_auth must be None or " | |||||
"BasicAuth() tuple for http proxy") | |||||
if proxy.scheme == 'socks4' and \ | |||||
not isinstance(proxy_auth, Socks4Auth): | |||||
raise ValueError("proxy_auth must be None or Socks4Auth() " | |||||
"tuple for socks4 proxy") | |||||
if proxy.scheme == 'socks5' and \ | |||||
not isinstance(proxy_auth, Socks5Auth): | |||||
raise ValueError("proxy_auth must be None or Socks5Auth() " | |||||
"tuple for socks5 proxy") | |||||
self.proxy = proxy | |||||
self.proxy_auth = proxy_auth | |||||
class ProxyConnector(aiohttp.TCPConnector): | |||||
def __init__(self, *, verify_ssl=True, fingerprint=None, | |||||
resolve=sentinel, use_dns_cache=True, | |||||
family=0, ssl_context=None, local_addr=None, | |||||
resolver=None, keepalive_timeout=sentinel, | |||||
force_close=False, limit=100, limit_per_host=0, | |||||
enable_cleanup_closed=False, loop=None, remote_resolve=True): | |||||
super().__init__( | |||||
verify_ssl=verify_ssl, fingerprint=fingerprint, resolve=resolve, | |||||
family=family, ssl_context=ssl_context, local_addr=local_addr, | |||||
resolver=resolver, keepalive_timeout=keepalive_timeout, | |||||
force_close=force_close, limit=limit, loop=loop, | |||||
limit_per_host=limit_per_host, use_dns_cache=use_dns_cache, | |||||
enable_cleanup_closed=enable_cleanup_closed) | |||||
class HttpProxyAddr(namedtuple('HttpProxyAddr', ['url'])): | |||||
def __new__(cls, url): | |||||
if url is None: | |||||
raise ValueError('None is not allowed as url value') | |||||
return super().__new__(cls, url) | |||||
class SocksConnector(aiohttp.TCPConnector): | |||||
def __init__(self, proxy, proxy_auth=None, *, remote_resolve=True, **kwgs): | |||||
super().__init__(**kwgs) | |||||
self._proxy = proxy | |||||
self._proxy_auth = proxy_auth | |||||
self._remote_resolve = remote_resolve | self._remote_resolve = remote_resolve | ||||
@property | |||||
def proxy(self): | |||||
"""Proxy info. | |||||
Should be Socks4Server/Socks5Server instance. | |||||
""" | |||||
return self._proxy | |||||
@property | |||||
def proxy_auth(self): | |||||
"""Proxy auth info. | |||||
Should be Socks4Auth/Socks5Auth instance. | |||||
""" | |||||
return self._proxy_auth | |||||
def _validate_ssl_fingerprint(self, transport, host): | |||||
has_cert = transport.get_extra_info('sslcontext') | |||||
if has_cert and self._fingerprint: | |||||
sock = transport.get_extra_info('socket') | |||||
if not hasattr(sock, 'getpeercert'): | |||||
# Workaround for asyncio 3.5.0 | |||||
# Starting from 3.5.1 version | |||||
# there is 'ssl_object' extra info in transport | |||||
sock = transport._ssl_protocol._sslpipe.ssl_object | |||||
# gives DER-encoded cert as a sequence of bytes (or None) | |||||
cert = sock.getpeercert(binary_form=True) | |||||
assert cert | |||||
got = self._hashfunc(cert).digest() | |||||
expected = self._fingerprint | |||||
if got != expected: | |||||
transport.close() | |||||
raise aiohttp.FingerprintMismatch( | |||||
expected, got, host, 80 | |||||
) | |||||
async def _create_proxy_connection(self, req): | |||||
if req.proxy.scheme == 'http': | |||||
return await super()._create_proxy_connection(req) | |||||
else: | |||||
return await self._create_socks_connection(req) | |||||
@asyncio.coroutine | |||||
def _create_connection(self, req): | |||||
async def _create_socks_connection(self, req): | |||||
if req.ssl: | if req.ssl: | ||||
sslcontext = self.ssl_context | sslcontext = self.ssl_context | ||||
else: | else: | ||||
sslcontext = None | sslcontext = None | ||||
if not self._remote_resolve: | if not self._remote_resolve: | ||||
dst_hosts = yield from self._resolve_host(req.host, req.port) | |||||
dst_hosts = await self._resolve_host(req.host, req.port) | |||||
dst = dst_hosts[0]['host'], dst_hosts[0]['port'] | dst = dst_hosts[0]['host'], dst_hosts[0]['port'] | ||||
else: | else: | ||||
dst = req.host, req.port | dst = req.host, req.port | ||||
proxy_hosts = yield from self._resolve_host(self._proxy.host, | |||||
self._proxy.port) | |||||
proxy_hosts = await self._resolve_host(req.proxy.host, req.proxy.port) | |||||
exc = None | exc = None | ||||
for hinfo in proxy_hosts: | for hinfo in proxy_hosts: | ||||
try: | |||||
proxy = self._proxy.__class__(host=hinfo['host'], | |||||
port=hinfo['port']) | |||||
if req.proxy.scheme == 'socks4': | |||||
proxy = Socks4Addr(hinfo['host'], hinfo['port']) | |||||
else: | |||||
proxy = Socks5Addr(hinfo['host'], hinfo['port']) | |||||
transp, proto = yield from create_connection( | |||||
self._factory, proxy, self._proxy_auth, dst, | |||||
try: | |||||
transp, proto = await create_connection( | |||||
self._factory, proxy, req.proxy_auth, dst, | |||||
loop=self._loop, remote_resolve=self._remote_resolve, | loop=self._loop, remote_resolve=self._remote_resolve, | ||||
ssl=sslcontext, family=hinfo['family'], | ssl=sslcontext, family=hinfo['family'], | ||||
proto=hinfo['proto'], flags=hinfo['flags'], | proto=hinfo['proto'], flags=hinfo['flags'], | ||||
local_addr=self._local_addr, | local_addr=self._local_addr, | ||||
server_hostname=req.host if sslcontext else None) | server_hostname=req.host if sslcontext else None) | ||||
self._validate_ssl_fingerprint(transp, req.host) | |||||
self._validate_ssl_fingerprint(transp, req.host, req.port) | |||||
return transp, proto | return transp, proto | ||||
except (OSError, SocksError, SocksConnectionError) as e: | except (OSError, SocksError, SocksConnectionError) as e: | ||||
exc = e | exc = e | ||||
else: | else: | ||||
if isinstance(exc, SocksConnectionError): | if isinstance(exc, SocksConnectionError): | ||||
raise ProxyConnectionError(*exc.args) | |||||
raise aiohttp.ClientProxyConnectionError(*exc.args) | |||||
if isinstance(exc, SocksError): | if isinstance(exc, SocksError): | ||||
raise exc | raise exc | ||||
else: | else: | ||||
@@ -107,12 +101,23 @@ class SocksConnector(aiohttp.TCPConnector): | |||||
exc.errno, 'Can not connect to %s:%s [%s]' % | exc.errno, 'Can not connect to %s:%s [%s]' % | ||||
(req.host, req.port, exc.strerror)) from exc | (req.host, req.port, exc.strerror)) from exc | ||||
def proxy_connector(proxy, proxy_auth=None, **kwargs): | |||||
if isinstance(proxy, HttpProxyAddr): | |||||
return aiohttp.ProxyConnector( | |||||
proxy.url, proxy_auth=proxy_auth, **kwargs) | |||||
elif isinstance(proxy, SocksAddr): | |||||
return SocksConnector(proxy, proxy_auth, **kwargs) | |||||
else: | |||||
raise ValueError('Unsupported `proxy` format') | |||||
def _validate_ssl_fingerprint(self, transp, host, port): | |||||
has_cert = transp.get_extra_info('sslcontext') | |||||
if has_cert and self._fingerprint: | |||||
sock = transp.get_extra_info('socket') | |||||
if not hasattr(sock, 'getpeercert'): | |||||
# Workaround for asyncio 3.5.0 | |||||
# Starting from 3.5.1 version | |||||
# there is 'ssl_object' extra info in transport | |||||
sock = transp._ssl_protocol._sslpipe.ssl_object | |||||
# gives DER-encoded cert as a sequence of bytes (or None) | |||||
cert = sock.getpeercert(binary_form=True) | |||||
assert cert | |||||
got = self._hashfunc(cert).digest() | |||||
expected = self._fingerprint | |||||
if got != expected: | |||||
transp.close() | |||||
if not self._cleanup_closed_disabled: | |||||
self._cleanup_closed_transports.append(transp) | |||||
raise aiohttp.ServerFingerprintMismatch( | |||||
expected, got, host, port) |
@@ -10,11 +10,6 @@ from .errors import ( | |||||
InvalidServerReply, InvalidServerVersion | InvalidServerReply, InvalidServerVersion | ||||
) | ) | ||||
try: | |||||
from asyncio import ensure_future | |||||
except ImportError: | |||||
ensure_future = asyncio.async | |||||
DEFAULT_LIMIT = getattr(asyncio.streams, '_DEFAULT_LIMIT', 2**16) | DEFAULT_LIMIT = getattr(asyncio.streams, '_DEFAULT_LIMIT', 2**16) | ||||
@@ -54,11 +49,10 @@ class BaseSocksProtocol(asyncio.StreamReaderProtocol): | |||||
super().__init__(stream_reader=reader, | super().__init__(stream_reader=reader, | ||||
client_connected_cb=self.negotiate, loop=self._loop) | client_connected_cb=self.negotiate, loop=self._loop) | ||||
@asyncio.coroutine | |||||
def negotiate(self, reader, writer): | |||||
async def negotiate(self, reader, writer): | |||||
try: | try: | ||||
req = self.socks_request(c.SOCKS_CMD_CONNECT) | req = self.socks_request(c.SOCKS_CMD_CONNECT) | ||||
self._proxy_peername, self._proxy_sockname = yield from req | |||||
self._proxy_peername, self._proxy_sockname = await req | |||||
except SocksError as exc: | except SocksError as exc: | ||||
exc = SocksError('Can not connect to %s:%s. %s' % | exc = SocksError('Can not connect to %s:%s. %s' % | ||||
(self._dst_host, self._dst_port, exc)) | (self._dst_host, self._dst_port, exc)) | ||||
@@ -134,8 +128,7 @@ class BaseSocksProtocol(asyncio.StreamReaderProtocol): | |||||
self._app_protocol.eof_received() | self._app_protocol.eof_received() | ||||
super().eof_received() | super().eof_received() | ||||
@asyncio.coroutine | |||||
def socks_request(self, cmd): | |||||
async def socks_request(self, cmd): | |||||
raise NotImplementedError | raise NotImplementedError | ||||
def write_request(self, request): | def write_request(self, request): | ||||
@@ -150,17 +143,15 @@ class BaseSocksProtocol(asyncio.StreamReaderProtocol): | |||||
raise ValueError('Unsupported item') | raise ValueError('Unsupported item') | ||||
self._stream_writer.write(bdata) | self._stream_writer.write(bdata) | ||||
@asyncio.coroutine | |||||
def read_response(self, n): | |||||
async def read_response(self, n): | |||||
try: | try: | ||||
return (yield from self._stream_reader.readexactly(n)) | |||||
return (await self._stream_reader.readexactly(n)) | |||||
except asyncio.IncompleteReadError as e: | except asyncio.IncompleteReadError as e: | ||||
raise InvalidServerReply( | raise InvalidServerReply( | ||||
'Server sent fewer bytes than required (%s)' % str(e)) | 'Server sent fewer bytes than required (%s)' % str(e)) | ||||
@asyncio.coroutine | |||||
def _get_dst_addr(self): | |||||
infos = yield from self._loop.getaddrinfo( | |||||
async def _get_dst_addr(self): | |||||
infos = await self._loop.getaddrinfo( | |||||
self._dst_host, self._dst_port, family=socket.AF_UNSPEC, | self._dst_host, self._dst_port, family=socket.AF_UNSPEC, | ||||
type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP, | type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP, | ||||
flags=socket.AI_ADDRCONFIG) | flags=socket.AI_ADDRCONFIG) | ||||
@@ -227,8 +218,7 @@ class Socks4Protocol(BaseSocksProtocol): | |||||
reader_limit=reader_limit, | reader_limit=reader_limit, | ||||
negotiate_done_cb=negotiate_done_cb) | negotiate_done_cb=negotiate_done_cb) | ||||
@asyncio.coroutine | |||||
def socks_request(self, cmd): | |||||
async def socks_request(self, cmd): | |||||
# prepare destination addr/port | # prepare destination addr/port | ||||
host, port = self._dst_host, self._dst_port | host, port = self._dst_host, self._dst_port | ||||
port_bytes = struct.pack(b'>H', port) | port_bytes = struct.pack(b'>H', port) | ||||
@@ -242,7 +232,7 @@ class Socks4Protocol(BaseSocksProtocol): | |||||
include_hostname = True | include_hostname = True | ||||
else: | else: | ||||
# it's not an IP number, so it's probably a DNS name. | # it's not an IP number, so it's probably a DNS name. | ||||
family, host = yield from self._get_dst_addr() | |||||
family, host = await self._get_dst_addr() | |||||
host_bytes = socket.inet_aton(host) | host_bytes = socket.inet_aton(host) | ||||
# build and send connect command | # build and send connect command | ||||
@@ -254,7 +244,7 @@ class Socks4Protocol(BaseSocksProtocol): | |||||
self.write_request(req) | self.write_request(req) | ||||
# read/process result | # read/process result | ||||
resp = yield from self.read_response(8) | |||||
resp = await self.read_response(8) | |||||
if resp[0] != c.NULL: | if resp[0] != c.NULL: | ||||
raise InvalidServerReply('SOCKS4 proxy server sent invalid data') | raise InvalidServerReply('SOCKS4 proxy server sent invalid data') | ||||
@@ -285,17 +275,16 @@ class Socks5Protocol(BaseSocksProtocol): | |||||
reader_limit=reader_limit, | reader_limit=reader_limit, | ||||
negotiate_done_cb=negotiate_done_cb) | negotiate_done_cb=negotiate_done_cb) | ||||
@asyncio.coroutine | |||||
def socks_request(self, cmd): | |||||
yield from self.authenticate() | |||||
async def socks_request(self, cmd): | |||||
await self.authenticate() | |||||
# build and send command | # build and send command | ||||
dst_addr, resolved = yield from self.build_dst_address( | |||||
dst_addr, resolved = await self.build_dst_address( | |||||
self._dst_host, self._dst_port) | self._dst_host, self._dst_port) | ||||
self.write_request([c.SOCKS_VER5, cmd, c.RSV] + dst_addr) | self.write_request([c.SOCKS_VER5, cmd, c.RSV] + dst_addr) | ||||
# read/process command response | # read/process command response | ||||
resp = yield from self.read_response(3) | |||||
resp = await self.read_response(3) | |||||
if resp[0] != c.SOCKS_VER5: | if resp[0] != c.SOCKS_VER5: | ||||
raise InvalidServerVersion( | raise InvalidServerVersion( | ||||
@@ -305,12 +294,11 @@ class Socks5Protocol(BaseSocksProtocol): | |||||
error = c.SOCKS5_ERRORS.get(resp[1], 'Unknown error') | error = c.SOCKS5_ERRORS.get(resp[1], 'Unknown error') | ||||
raise SocksError('[Errno {0:#04x}]: {1}'.format(resp[1], error)) | raise SocksError('[Errno {0:#04x}]: {1}'.format(resp[1], error)) | ||||
binded = yield from self.read_address() | |||||
binded = await self.read_address() | |||||
return resolved, binded | return resolved, binded | ||||
@asyncio.coroutine | |||||
def authenticate(self): | |||||
async def authenticate(self): | |||||
# send available auth methods | # send available auth methods | ||||
if self._auth.login and self._auth.password: | if self._auth.login and self._auth.password: | ||||
req = [c.SOCKS_VER5, 0x02, | req = [c.SOCKS_VER5, 0x02, | ||||
@@ -321,7 +309,7 @@ class Socks5Protocol(BaseSocksProtocol): | |||||
self.write_request(req) | self.write_request(req) | ||||
# read/process response and send auth data if necessary | # read/process response and send auth data if necessary | ||||
chosen_auth = yield from self.read_response(2) | |||||
chosen_auth = await self.read_response(2) | |||||
if chosen_auth[0] != c.SOCKS_VER5: | if chosen_auth[0] != c.SOCKS_VER5: | ||||
raise InvalidServerVersion( | raise InvalidServerVersion( | ||||
@@ -333,7 +321,7 @@ class Socks5Protocol(BaseSocksProtocol): | |||||
chr(len(self._auth.password)).encode(), self._auth.password] | chr(len(self._auth.password)).encode(), self._auth.password] | ||||
self.write_request(req) | self.write_request(req) | ||||
auth_status = yield from self.read_response(2) | |||||
auth_status = await self.read_response(2) | |||||
if auth_status[0] != 0x01: | if auth_status[0] != 0x01: | ||||
raise InvalidServerReply( | raise InvalidServerReply( | ||||
'SOCKS5 proxy server sent invalid data' | 'SOCKS5 proxy server sent invalid data' | ||||
@@ -353,8 +341,7 @@ class Socks5Protocol(BaseSocksProtocol): | |||||
'SOCKS5 proxy server sent invalid data' | 'SOCKS5 proxy server sent invalid data' | ||||
) | ) | ||||
@asyncio.coroutine | |||||
def build_dst_address(self, host, port): | |||||
async def build_dst_address(self, host, port): | |||||
family_to_byte = {socket.AF_INET: c.SOCKS5_ATYP_IPv4, | family_to_byte = {socket.AF_INET: c.SOCKS5_ATYP_IPv4, | ||||
socket.AF_INET6: c.SOCKS5_ATYP_IPv6} | socket.AF_INET6: c.SOCKS5_ATYP_IPv6} | ||||
port_bytes = struct.pack('>H', port) | port_bytes = struct.pack('>H', port) | ||||
@@ -375,29 +362,28 @@ class Socks5Protocol(BaseSocksProtocol): | |||||
req = [c.SOCKS5_ATYP_DOMAIN, chr(len(host_bytes)).encode(), | req = [c.SOCKS5_ATYP_DOMAIN, chr(len(host_bytes)).encode(), | ||||
host_bytes, port_bytes] | host_bytes, port_bytes] | ||||
else: | else: | ||||
family, host_bytes = yield from self._get_dst_addr() | |||||
family, host_bytes = await self._get_dst_addr() | |||||
host_bytes = socket.inet_pton(family, host_bytes) | host_bytes = socket.inet_pton(family, host_bytes) | ||||
req = [family_to_byte[family], host_bytes, port_bytes] | req = [family_to_byte[family], host_bytes, port_bytes] | ||||
host = socket.inet_ntop(family, host_bytes) | host = socket.inet_ntop(family, host_bytes) | ||||
return req, (host, port) | return req, (host, port) | ||||
@asyncio.coroutine | |||||
def read_address(self): | |||||
atype = yield from self.read_response(1) | |||||
async def read_address(self): | |||||
atype = await self.read_response(1) | |||||
if atype[0] == c.SOCKS5_ATYP_IPv4: | if atype[0] == c.SOCKS5_ATYP_IPv4: | ||||
addr = socket.inet_ntoa((yield from self.read_response(4))) | |||||
addr = socket.inet_ntoa((await self.read_response(4))) | |||||
elif atype[0] == c.SOCKS5_ATYP_DOMAIN: | elif atype[0] == c.SOCKS5_ATYP_DOMAIN: | ||||
length = yield from self.read_response(1) | |||||
addr = yield from self.read_response(ord(length)) | |||||
length = await self.read_response(1) | |||||
addr = await self.read_response(ord(length)) | |||||
elif atype[0] == c.SOCKS5_ATYP_IPv6: | elif atype[0] == c.SOCKS5_ATYP_IPv6: | ||||
addr = yield from self.read_response(16) | |||||
addr = await self.read_response(16) | |||||
addr = socket.inet_ntop(socket.AF_INET6, addr) | addr = socket.inet_ntop(socket.AF_INET6, addr) | ||||
else: | else: | ||||
raise InvalidServerReply('SOCKS5 proxy server sent invalid data') | raise InvalidServerReply('SOCKS5 proxy server sent invalid data') | ||||
port = yield from self.read_response(2) | |||||
port = await self.read_response(2) | |||||
port = struct.unpack('>H', port)[0] | port = struct.unpack('>H', port)[0] | ||||
return addr, port | return addr, port |
@@ -0,0 +1,136 @@ | |||||
import asyncio | |||||
import struct | |||||
import socket | |||||
from aiohttp.test_utils import unused_port | |||||
class FakeSocksSrv: | |||||
def __init__(self, loop, write_buff): | |||||
self._loop = loop | |||||
self._write_buff = write_buff | |||||
self._transports = [] | |||||
self._srv = None | |||||
self.port = unused_port() | |||||
async def __aenter__(self): | |||||
transports = self._transports | |||||
write_buff = self._write_buff | |||||
class SocksPrimitiveProtocol(asyncio.Protocol): | |||||
_transport = None | |||||
def connection_made(self, transport): | |||||
self._transport = transport | |||||
transports.append(transport) | |||||
def data_received(self, data): | |||||
self._transport.write(write_buff) | |||||
def factory(): | |||||
return SocksPrimitiveProtocol() | |||||
self._srv = await self._loop.create_server( | |||||
factory, '127.0.0.1', self.port) | |||||
return self | |||||
async def __aexit__(self, exc_type, exc_val, exc_tb): | |||||
for tr in self._transports: | |||||
tr.close() | |||||
self._srv.close() | |||||
await self._srv.wait_closed() | |||||
class FakeSocks4Srv: | |||||
def __init__(self, loop): | |||||
self._loop = loop | |||||
self._transports = [] | |||||
self._futures = [] | |||||
self._srv = None | |||||
self.port = unused_port() | |||||
async def __aenter__(self): | |||||
transports = self._transports | |||||
futures = self._futures | |||||
class Socks4Protocol(asyncio.StreamReaderProtocol): | |||||
def __init__(self, _loop): | |||||
self._loop = _loop | |||||
reader = asyncio.StreamReader(loop=self._loop) | |||||
super().__init__(reader, client_connected_cb=self.negotiate, | |||||
loop=self._loop) | |||||
def connection_made(self, transport): | |||||
transports.append(transport) | |||||
super().connection_made(transport) | |||||
async def negotiate(self, reader, writer): | |||||
writer.write(b'\x00\x5a\x04W\x01\x01\x01\x01') | |||||
data = await reader.read(9) | |||||
dst_port = struct.unpack('>H', data[2:4])[0] | |||||
dst_addr = data[4:8] | |||||
if data[-1] != 0x00: | |||||
while True: | |||||
byte = await reader.read(1) | |||||
if byte == 0x00: | |||||
break | |||||
if dst_addr == b'\x00\x00\x00\x01': | |||||
dst_addr = bytearray() | |||||
while True: | |||||
byte = await reader.read(1) | |||||
if byte == 0x00: | |||||
break | |||||
dst_addr.append(byte) | |||||
else: | |||||
dst_addr = socket.inet_ntoa(dst_addr) | |||||
cl_reader, cl_writer = await asyncio.open_connection( | |||||
host=dst_addr, port=dst_port, loop=self._loop | |||||
) | |||||
transports.append(cl_writer) | |||||
cl_fut = asyncio.ensure_future( | |||||
self.retranslator(reader, cl_writer), loop=self._loop) | |||||
dst_fut = asyncio.ensure_future( | |||||
self.retranslator(cl_reader, writer), loop=self._loop) | |||||
futures.append(cl_fut) | |||||
futures.append(dst_fut) | |||||
async def retranslator(self, reader, writer): | |||||
data = bytearray() | |||||
while True: | |||||
try: | |||||
byte = await reader.read(10) | |||||
if not byte: | |||||
break | |||||
data.append(byte[0]) | |||||
writer.write(byte) | |||||
await writer.drain() | |||||
except: | |||||
break | |||||
def factory(): | |||||
return Socks4Protocol(_loop=self._loop) | |||||
self._srv = await self._loop.create_server( | |||||
factory, '127.0.0.1', self.port) | |||||
return self | |||||
async def __aexit__(self, exc_type, exc_val, exc_tb): | |||||
for tr in self._transports: | |||||
tr.close() | |||||
self._srv.close() | |||||
await self._srv.wait_closed() | |||||
for f in self._futures: | |||||
if not f.cancelled() or not f.done(): | |||||
f.cancel() |
@@ -0,0 +1 @@ | |||||
from aiohttp.pytest_plugin import * # noqa |
@@ -1,242 +0,0 @@ | |||||
import asyncio | |||||
import aiohttp | |||||
import contextlib | |||||
import gc | |||||
import os | |||||
import socket | |||||
import ssl | |||||
import struct | |||||
import threading | |||||
from unittest import mock | |||||
from aiohttp.server import ServerHttpProtocol | |||||
try: | |||||
from asyncio import ensure_future | |||||
except ImportError: | |||||
ensure_future = asyncio.async | |||||
def fake_coroutine(return_value): | |||||
def coro(*args, **kwargs): | |||||
if isinstance(return_value, Exception): | |||||
raise return_value | |||||
return return_value | |||||
return mock.Mock(side_effect=asyncio.coroutine(coro)) | |||||
def find_unused_port(): | |||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | |||||
s.bind(('127.0.0.1', 0)) | |||||
port = s.getsockname()[1] | |||||
s.close() | |||||
return port | |||||
@contextlib.contextmanager | |||||
def fake_socks_srv(loop, write_buff): | |||||
transports = [] | |||||
class SocksPrimitiveProtocol(asyncio.Protocol): | |||||
_transport = None | |||||
def connection_made(self, transport): | |||||
self._transport = transport | |||||
transports.append(transport) | |||||
def data_received(self, data): | |||||
self._transport.write(write_buff) | |||||
port = find_unused_port() | |||||
def factory(): | |||||
return SocksPrimitiveProtocol() | |||||
srv = loop.run_until_complete( | |||||
loop.create_server(factory, '127.0.0.1', port)) | |||||
yield port | |||||
for tr in transports: | |||||
tr.close() | |||||
srv.close() | |||||
loop.run_until_complete(srv.wait_closed()) | |||||
gc.collect() | |||||
@contextlib.contextmanager | |||||
def fake_socks4_srv(loop): | |||||
port = find_unused_port() | |||||
transports = [] | |||||
futures = [] | |||||
class Socks4Protocol(asyncio.StreamReaderProtocol): | |||||
def __init__(self, _loop): | |||||
self._loop = _loop | |||||
reader = asyncio.StreamReader(loop=self._loop) | |||||
super().__init__(reader, client_connected_cb=self.negotiate, | |||||
loop=self._loop) | |||||
def connection_made(self, transport): | |||||
transports.append(transport) | |||||
super().connection_made(transport) | |||||
@asyncio.coroutine | |||||
def negotiate(self, reader, writer): | |||||
writer.write(b'\x00\x5a\x04W\x01\x01\x01\x01') | |||||
data = yield from reader.read(9) | |||||
dst_port = struct.unpack('>H', data[2:4])[0] | |||||
dst_addr = data[4:8] | |||||
if data[-1] != 0x00: | |||||
while True: | |||||
byte = yield from reader.read(1) | |||||
if byte == 0x00: | |||||
break | |||||
if dst_addr == b'\x00\x00\x00\x01': | |||||
dst_addr = bytearray() | |||||
while True: | |||||
byte = yield from reader.read(1) | |||||
if byte == 0x00: | |||||
break | |||||
dst_addr.append(byte) | |||||
else: | |||||
dst_addr = socket.inet_ntoa(dst_addr) | |||||
cl_reader, cl_writer = yield from asyncio.open_connection( | |||||
host=dst_addr, port=dst_port, loop=self._loop | |||||
) | |||||
transports.append(cl_writer) | |||||
cl_fut = ensure_future( | |||||
self.retranslator(reader, cl_writer), loop=self._loop) | |||||
dst_fut = ensure_future( | |||||
self.retranslator(cl_reader, writer), loop=self._loop) | |||||
futures.append(cl_fut) | |||||
futures.append(dst_fut) | |||||
@asyncio.coroutine | |||||
def retranslator(self, reader, writer): | |||||
data = bytearray() | |||||
while True: | |||||
try: | |||||
byte = yield from reader.read(1) | |||||
if not byte: | |||||
break | |||||
data.append(byte[0]) | |||||
writer.write(byte) | |||||
yield from writer.drain() | |||||
except: | |||||
break | |||||
def run(_fut): | |||||
thread_loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(thread_loop) | |||||
srv_coroutine = thread_loop.create_server( | |||||
lambda: Socks4Protocol(thread_loop), '127.0.0.1', port) | |||||
srv = thread_loop.run_until_complete(srv_coroutine) | |||||
waiter = asyncio.Future(loop=thread_loop) | |||||
loop.call_soon_threadsafe( | |||||
_fut.set_result, (thread_loop, waiter)) | |||||
try: | |||||
thread_loop.run_until_complete(waiter) | |||||
finally: | |||||
# close opened transports | |||||
for tr in transports: | |||||
tr.close() | |||||
for ft in futures: | |||||
if not ft.done(): | |||||
ft.set_result(1) | |||||
srv.close() | |||||
thread_loop.stop() | |||||
thread_loop.close() | |||||
gc.collect() | |||||
fut = asyncio.Future(loop=loop) | |||||
srv_thread = threading.Thread(target=run, args=(fut,)) | |||||
srv_thread.start() | |||||
_thread_loop, _waiter = loop.run_until_complete(fut) | |||||
yield port | |||||
_thread_loop.call_soon_threadsafe(_waiter.set_result, None) | |||||
srv_thread.join() | |||||
@contextlib.contextmanager | |||||
def http_srv(loop, *, listen_addr=('127.0.0.1', 0), use_ssl=False): | |||||
transports = [] | |||||
class TestHttpServer(ServerHttpProtocol): | |||||
def connection_made(self, transport): | |||||
transports.append(transport) | |||||
super().connection_made(transport) | |||||
@asyncio.coroutine | |||||
def handle_request(self, message, payload): | |||||
response = aiohttp.Response(self.writer, 200, message.version) | |||||
text = b'Test message' | |||||
response.add_header('Content-type', 'text/plain') | |||||
response.add_header('Content-length', str(len(text))) | |||||
response.send_headers() | |||||
response.write(text) | |||||
response.write_eof() | |||||
if use_ssl: | |||||
here = os.path.join(os.path.dirname(__file__), '..', 'tests') | |||||
keyfile = os.path.join(here, 'sample.key') | |||||
certfile = os.path.join(here, 'sample.crt') | |||||
sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23) | |||||
sslcontext.load_cert_chain(certfile, keyfile) | |||||
else: | |||||
sslcontext = None | |||||
def run(_fut): | |||||
thread_loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(thread_loop) | |||||
host, port = listen_addr | |||||
srv_coroutine = thread_loop.create_server( | |||||
lambda: TestHttpServer(), host, port, ssl=sslcontext) | |||||
srv = thread_loop.run_until_complete(srv_coroutine) | |||||
waiter = asyncio.Future(loop=thread_loop) | |||||
loop.call_soon_threadsafe( | |||||
_fut.set_result, (thread_loop, waiter, | |||||
srv.sockets[0].getsockname())) | |||||
try: | |||||
thread_loop.run_until_complete(waiter) | |||||
finally: | |||||
# close opened transports | |||||
for tr in transports: | |||||
tr.close() | |||||
srv.close() | |||||
thread_loop.stop() | |||||
thread_loop.close() | |||||
gc.collect() | |||||
fut = asyncio.Future(loop=loop) | |||||
srv_thread = threading.Thread(target=run, args=(fut,)) | |||||
srv_thread.start() | |||||
_thread_loop, _waiter, _addr = loop.run_until_complete(fut) | |||||
url = '{}://{}:{}'.format( | |||||
'https' if use_ssl else 'http', *_addr) | |||||
yield url | |||||
_thread_loop.call_soon_threadsafe(_waiter.set_result, None) | |||||
srv_thread.join() |
@@ -1,155 +1,179 @@ | |||||
import unittest | |||||
import asyncio | |||||
import aiosocks | import aiosocks | ||||
import aiohttp | import aiohttp | ||||
import pytest | |||||
from yarl import URL | |||||
from unittest import mock | from unittest import mock | ||||
from aiohttp.client_reqrep import ClientRequest | |||||
from aiosocks.connector import SocksConnector, proxy_connector, HttpProxyAddr | |||||
from .helpers import fake_coroutine | |||||
from aiohttp.test_utils import make_mocked_coro | |||||
from aiohttp import BasicAuth | |||||
from aiosocks.connector import ProxyConnector, ProxyClientRequest | |||||
from aiosocks.helpers import Socks4Auth, Socks5Auth | |||||
class TestSocksConnector(unittest.TestCase): | |||||
def setUp(self): | |||||
self.loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(None) | |||||
def tearDown(self): | |||||
self.loop.close() | |||||
def test_properties(self): | |||||
addr = aiosocks.Socks4Addr('localhost') | |||||
auth = aiosocks.Socks4Auth('login') | |||||
conn = SocksConnector(addr, auth, loop=self.loop) | |||||
self.assertIs(conn.proxy, addr) | |||||
self.assertIs(conn.proxy_auth, auth) | |||||
@mock.patch('aiosocks.connector.create_connection') | |||||
def test_connect_proxy_ip(self, cr_conn_mock): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
cr_conn_mock.side_effect = \ | |||||
fake_coroutine((tr, proto)).side_effect | |||||
async def test_connect_proxy_ip(): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
with mock.patch('aiosocks.connector.create_connection', | |||||
make_mocked_coro((tr, proto))): | |||||
loop_mock = mock.Mock() | loop_mock = mock.Mock() | ||||
loop_mock.getaddrinfo = make_mocked_coro( | |||||
[[0, 0, 0, 0, ['127.0.0.1', 1080]]]) | |||||
req = ClientRequest('GET', 'http://python.org', loop=self.loop) | |||||
connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), | |||||
None, loop=loop_mock) | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop_mock, | |||||
proxy=URL('socks5://proxy.org')) | |||||
connector = ProxyConnector(loop=loop_mock) | |||||
conn = await connector.connect(req) | |||||
loop_mock.getaddrinfo = fake_coroutine([mock.MagicMock()]) | |||||
assert loop_mock.getaddrinfo.called | |||||
assert conn.protocol is proto | |||||
conn = self.loop.run_until_complete(connector.connect(req)) | |||||
conn.close() | |||||
self.assertTrue(loop_mock.getaddrinfo.is_called) | |||||
self.assertIs(conn._transport, tr) | |||||
conn.close() | |||||
async def test_connect_proxy_domain(): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
@mock.patch('aiosocks.connector.create_connection') | |||||
def test_connect_proxy_domain(self, cr_conn_mock): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
cr_conn_mock.side_effect = \ | |||||
fake_coroutine((tr, proto)).side_effect | |||||
with mock.patch('aiosocks.connector.create_connection', | |||||
make_mocked_coro((tr, proto))): | |||||
loop_mock = mock.Mock() | loop_mock = mock.Mock() | ||||
req = ClientRequest('GET', 'http://python.org', loop=self.loop) | |||||
connector = SocksConnector(aiosocks.Socks5Addr('proxy.example'), | |||||
None, loop=loop_mock) | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop_mock, | |||||
proxy=URL('socks5://proxy.example')) | |||||
connector = ProxyConnector(loop=loop_mock) | |||||
connector._resolve_host = fake_coroutine([mock.MagicMock()]) | |||||
connector._resolve_host = make_mocked_coro([mock.MagicMock()]) | |||||
conn = await connector.connect(req) | |||||
conn = self.loop.run_until_complete(connector.connect(req)) | |||||
assert connector._resolve_host.call_count == 1 | |||||
assert conn.protocol is proto | |||||
self.assertTrue(connector._resolve_host.is_called) | |||||
self.assertEqual(connector._resolve_host.call_count, 1) | |||||
self.assertIs(conn._transport, tr) | |||||
conn.close() | |||||
conn.close() | |||||
@mock.patch('aiosocks.connector.create_connection') | |||||
def test_connect_remote_resolve(self, cr_conn_mock): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
cr_conn_mock.side_effect = \ | |||||
fake_coroutine((tr, proto)).side_effect | |||||
async def test_connect_remote_resolve(loop): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
req = ClientRequest('GET', 'http://python.org', loop=self.loop) | |||||
connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), | |||||
None, loop=self.loop, remote_resolve=True) | |||||
with mock.patch('aiosocks.connector.create_connection', | |||||
make_mocked_coro((tr, proto))): | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('socks5://127.0.0.1')) | |||||
connector = ProxyConnector(loop=loop, remote_resolve=True) | |||||
connector._resolve_host = make_mocked_coro([mock.MagicMock()]) | |||||
connector._resolve_host = fake_coroutine([mock.MagicMock()]) | |||||
conn = await connector.connect(req) | |||||
conn = self.loop.run_until_complete(connector.connect(req)) | |||||
assert connector._resolve_host.call_count == 1 | |||||
assert conn.protocol is proto | |||||
self.assertEqual(connector._resolve_host.call_count, 1) | |||||
conn.close() | |||||
conn.close() | |||||
@mock.patch('aiosocks.connector.create_connection') | |||||
def test_connect_locale_resolve(self, cr_conn_mock): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
cr_conn_mock.side_effect = \ | |||||
fake_coroutine((tr, proto)).side_effect | |||||
async def test_connect_locale_resolve(loop): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
req = ClientRequest('GET', 'http://python.org', loop=self.loop) | |||||
connector = SocksConnector(aiosocks.Socks5Addr('proxy.example'), | |||||
None, loop=self.loop, remote_resolve=False) | |||||
with mock.patch('aiosocks.connector.create_connection', | |||||
make_mocked_coro((tr, proto))): | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('socks5://proxy.example')) | |||||
connector = ProxyConnector(loop=loop, remote_resolve=False) | |||||
connector._resolve_host = make_mocked_coro([mock.MagicMock()]) | |||||
connector._resolve_host = fake_coroutine([mock.MagicMock()]) | |||||
conn = await connector.connect(req) | |||||
conn = self.loop.run_until_complete(connector.connect(req)) | |||||
assert connector._resolve_host.call_count == 2 | |||||
assert conn.protocol is proto | |||||
self.assertTrue(connector._resolve_host.is_called) | |||||
self.assertEqual(connector._resolve_host.call_count, 2) | |||||
conn.close() | |||||
conn.close() | |||||
@mock.patch('aiosocks.connector.create_connection') | |||||
def test_proxy_connect_fail(self, cr_conn_mock): | |||||
loop_mock = mock.Mock() | |||||
cr_conn_mock.side_effect = \ | |||||
fake_coroutine(aiosocks.SocksConnectionError()).side_effect | |||||
async def test_proxy_connect_fail(loop): | |||||
loop_mock = mock.Mock() | |||||
loop_mock.getaddrinfo = make_mocked_coro( | |||||
[[0, 0, 0, 0, ['127.0.0.1', 1080]]]) | |||||
cc_coro = make_mocked_coro( | |||||
raise_exception=aiosocks.SocksConnectionError()) | |||||
req = ClientRequest('GET', 'http://python.org', loop=self.loop) | |||||
connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), | |||||
None, loop=loop_mock) | |||||
with mock.patch('aiosocks.connector.create_connection', cc_coro): | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('socks5://127.0.0.1')) | |||||
connector = ProxyConnector(loop=loop_mock) | |||||
loop_mock.getaddrinfo = fake_coroutine([mock.MagicMock()]) | |||||
with pytest.raises(aiohttp.ClientConnectionError): | |||||
await connector.connect(req) | |||||
with self.assertRaises(aiohttp.ProxyConnectionError): | |||||
self.loop.run_until_complete(connector.connect(req)) | |||||
@mock.patch('aiosocks.connector.create_connection') | |||||
def test_proxy_negotiate_fail(self, cr_conn_mock): | |||||
loop_mock = mock.Mock() | |||||
cr_conn_mock.side_effect = \ | |||||
fake_coroutine(aiosocks.SocksError()).side_effect | |||||
async def test_proxy_negotiate_fail(loop): | |||||
loop_mock = mock.Mock() | |||||
loop_mock.getaddrinfo = make_mocked_coro( | |||||
[[0, 0, 0, 0, ['127.0.0.1', 1080]]]) | |||||
req = ClientRequest('GET', 'http://python.org', loop=self.loop) | |||||
connector = SocksConnector(aiosocks.Socks5Addr('127.0.0.1'), | |||||
None, loop=loop_mock) | |||||
with mock.patch('aiosocks.connector.create_connection', | |||||
make_mocked_coro(raise_exception=aiosocks.SocksError())): | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('socks5://127.0.0.1')) | |||||
connector = ProxyConnector(loop=loop_mock) | |||||
loop_mock.getaddrinfo = fake_coroutine([mock.MagicMock()]) | |||||
with pytest.raises(aiosocks.SocksError): | |||||
await connector.connect(req) | |||||
with self.assertRaises(aiosocks.SocksError): | |||||
self.loop.run_until_complete(connector.connect(req)) | |||||
def test_proxy_connector(self): | |||||
socks4_addr = aiosocks.Socks4Addr('h') | |||||
socks5_addr = aiosocks.Socks5Addr('h') | |||||
http_addr = HttpProxyAddr('http://proxy') | |||||
async def test_proxy_connect_http(loop): | |||||
tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') | |||||
loop_mock = mock.Mock() | |||||
loop_mock.getaddrinfo = make_mocked_coro([ | |||||
[0, 0, 0, 0, ['127.0.0.1', 1080]]]) | |||||
loop_mock.create_connection = make_mocked_coro((tr, proto)) | |||||
self.assertIsInstance(proxy_connector(socks4_addr, loop=self.loop), | |||||
SocksConnector) | |||||
self.assertIsInstance(proxy_connector(socks5_addr, loop=self.loop), | |||||
SocksConnector) | |||||
self.assertIsInstance(proxy_connector(http_addr, loop=self.loop), | |||||
aiohttp.ProxyConnector) | |||||
req = ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('http://127.0.0.1')) | |||||
connector = ProxyConnector(loop=loop_mock) | |||||
with self.assertRaises(ValueError): | |||||
proxy_connector(None) | |||||
await connector.connect(req) | |||||
def test_http_proxy_addr(self): | |||||
addr = HttpProxyAddr('http://proxy') | |||||
self.assertEqual(addr.url, 'http://proxy') | |||||
with self.assertRaises(ValueError): | |||||
HttpProxyAddr(None) | |||||
@pytest.mark.parametrize('proxy', [ | |||||
(URL('socks4://proxy.org'), Socks4Auth('login')), | |||||
(URL('socks5://proxy.org'), Socks5Auth('login', 'password')), | |||||
(URL('http://proxy.org'), BasicAuth('login')), (None, BasicAuth('login')), | |||||
(URL('socks4://proxy.org'), None), (None, None)]) | |||||
def test_proxy_client_request_valid(proxy, loop): | |||||
proxy, proxy_auth = proxy | |||||
p = ProxyClientRequest('GET', URL('http://python.org'), | |||||
proxy=proxy, proxy_auth=proxy_auth, loop=loop) | |||||
assert p.proxy is proxy | |||||
assert p.proxy_auth is proxy_auth | |||||
def test_proxy_client_request_invalid(loop): | |||||
with pytest.raises(ValueError) as cm: | |||||
ProxyClientRequest( | |||||
'GET', URL('http://python.org'), | |||||
proxy=URL('socks6://proxy.org'), proxy_auth=None, loop=loop) | |||||
assert 'Only http, socks4 and socks5 proxies are supported' in str(cm) | |||||
with pytest.raises(ValueError) as cm: | |||||
ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('http://proxy.org'), proxy_auth=Socks4Auth('l')) | |||||
assert 'proxy_auth must be None or BasicAuth() ' \ | |||||
'tuple for http proxy' in str(cm) | |||||
with pytest.raises(ValueError) as cm: | |||||
ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('socks4://proxy.org'), proxy_auth=BasicAuth('l')) | |||||
assert 'proxy_auth must be None or Socks4Auth() ' \ | |||||
'tuple for socks4 proxy' in str(cm) | |||||
with pytest.raises(ValueError) as cm: | |||||
ProxyClientRequest( | |||||
'GET', URL('http://python.org'), loop=loop, | |||||
proxy=URL('socks5://proxy.org'), proxy_auth=Socks4Auth('l')) | |||||
assert 'proxy_auth must be None or Socks5Auth() ' \ | |||||
'tuple for socks5 proxy' in str(cm) |
@@ -1,134 +1,97 @@ | |||||
import unittest | |||||
import pytest | |||||
import aiosocks | import aiosocks | ||||
import asyncio | |||||
from aiohttp.test_utils import make_mocked_coro | |||||
from unittest import mock | from unittest import mock | ||||
from .helpers import fake_coroutine | |||||
try: | |||||
from asyncio import ensure_future | |||||
except ImportError: | |||||
ensure_future = asyncio.async | |||||
class TestCreateConnection(unittest.TestCase): | |||||
def setUp(self): | |||||
self.loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(None) | |||||
def tearDown(self): | |||||
self.loop.close() | |||||
def test_init(self): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
# proxy argument | |||||
with self.assertRaises(AssertionError) as ct: | |||||
conn = aiosocks.create_connection(None, None, auth, dst) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertEqual(str(ct.exception), | |||||
'proxy must be Socks4Addr() or Socks5Addr() tuple') | |||||
with self.assertRaises(AssertionError) as ct: | |||||
conn = aiosocks.create_connection(None, auth, auth, dst) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertEqual(str(ct.exception), | |||||
'proxy must be Socks4Addr() or Socks5Addr() tuple') | |||||
# proxy_auth | |||||
with self.assertRaises(AssertionError) as ct: | |||||
conn = aiosocks.create_connection(None, addr, addr, dst) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertIn('proxy_auth must be None or Socks4Auth()', | |||||
str(ct.exception)) | |||||
# dst | |||||
with self.assertRaises(AssertionError) as ct: | |||||
conn = aiosocks.create_connection(None, addr, auth, None) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertIn('invalid dst format, tuple("dst_host", dst_port))', | |||||
str(ct.exception)) | |||||
# addr and auth compatibility | |||||
with self.assertRaises(ValueError) as ct: | |||||
conn = aiosocks.create_connection( | |||||
None, addr, aiosocks.Socks4Auth(''), dst | |||||
) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertIn('proxy is Socks5Addr but proxy_auth is not Socks5Auth', | |||||
str(ct.exception)) | |||||
with self.assertRaises(ValueError) as ct: | |||||
conn = aiosocks.create_connection( | |||||
None, aiosocks.Socks4Addr(''), auth, dst | |||||
) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertIn('proxy is Socks4Addr but proxy_auth is not Socks4Auth', | |||||
str(ct.exception)) | |||||
# test ssl, server_hostname | |||||
with self.assertRaises(ValueError) as ct: | |||||
conn = aiosocks.create_connection( | |||||
None, addr, auth, dst, server_hostname='python.org' | |||||
) | |||||
self.loop.run_until_complete(conn) | |||||
self.assertIn('server_hostname is only meaningful with ssl', | |||||
str(ct.exception)) | |||||
def test_connection_fail(self): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
loop_mock = mock.Mock() | |||||
loop_mock.create_connection = fake_coroutine(OSError()) | |||||
with self.assertRaises(aiosocks.SocksConnectionError): | |||||
conn = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop_mock | |||||
) | |||||
self.loop.run_until_complete(conn) | |||||
@mock.patch('aiosocks.asyncio.Future') | |||||
def test_negotiate_fail(self, future_mock): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
loop_mock = mock.Mock() | |||||
loop_mock.create_connection = fake_coroutine( | |||||
(mock.Mock(), mock.Mock()) | |||||
) | |||||
fut = fake_coroutine(aiosocks.SocksError()) | |||||
future_mock.side_effect = fut.side_effect | |||||
with self.assertRaises(aiosocks.SocksError): | |||||
conn = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop_mock | |||||
) | |||||
self.loop.run_until_complete(conn) | |||||
@mock.patch('aiosocks.asyncio.Future') | |||||
def test_open_connection(self, future_mock): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
transp, proto = mock.Mock(), mock.Mock() | |||||
reader, writer = mock.Mock(), mock.Mock() | |||||
proto.app_protocol.reader, proto.app_protocol.writer = reader, writer | |||||
loop_mock = mock.Mock() | |||||
loop_mock.create_connection = fake_coroutine((transp, proto)) | |||||
fut = fake_coroutine(True) | |||||
future_mock.side_effect = fut.side_effect | |||||
conn = aiosocks.open_connection(addr, auth, dst, loop=loop_mock) | |||||
r, w = self.loop.run_until_complete(conn) | |||||
self.assertIs(reader, r) | |||||
self.assertIs(writer, w) | |||||
async def test_create_connection_init(): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
# proxy argument | |||||
with pytest.raises(AssertionError) as ct: | |||||
await aiosocks.create_connection(None, None, auth, dst) | |||||
assert 'proxy must be Socks4Addr() or Socks5Addr() tuple' in str(ct) | |||||
with pytest.raises(AssertionError) as ct: | |||||
await aiosocks.create_connection(None, auth, auth, dst) | |||||
assert 'proxy must be Socks4Addr() or Socks5Addr() tuple' in str(ct) | |||||
# proxy_auth | |||||
with pytest.raises(AssertionError) as ct: | |||||
await aiosocks.create_connection(None, addr, addr, dst) | |||||
assert 'proxy_auth must be None or Socks4Auth()' in str(ct) | |||||
# dst | |||||
with pytest.raises(AssertionError) as ct: | |||||
await aiosocks.create_connection(None, addr, auth, None) | |||||
assert 'invalid dst format, tuple("dst_host", dst_port))' in str(ct) | |||||
# addr and auth compatibility | |||||
with pytest.raises(ValueError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, aiosocks.Socks4Auth(''), dst) | |||||
assert 'proxy is Socks5Addr but proxy_auth is not Socks5Auth' in str(ct) | |||||
with pytest.raises(ValueError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, aiosocks.Socks4Addr(''), auth, dst) | |||||
assert 'proxy is Socks4Addr but proxy_auth is not Socks4Auth' in str(ct) | |||||
# test ssl, server_hostname | |||||
with pytest.raises(ValueError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, server_hostname='python.org') | |||||
assert 'server_hostname is only meaningful with ssl' in str(ct) | |||||
async def test_connection_fail(): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
loop_mock = mock.Mock() | |||||
loop_mock.create_connection = make_mocked_coro(raise_exception=OSError()) | |||||
with pytest.raises(aiosocks.SocksConnectionError): | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop_mock) | |||||
async def test_negotiate_fail(): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
loop_mock = mock.Mock() | |||||
loop_mock.create_connection = make_mocked_coro((mock.Mock(), mock.Mock())) | |||||
with mock.patch('aiosocks.asyncio.Future') as future_mock: | |||||
future_mock.side_effect = make_mocked_coro( | |||||
raise_exception=aiosocks.SocksError()) | |||||
with pytest.raises(aiosocks.SocksError): | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop_mock) | |||||
async def test_open_connection(): | |||||
addr = aiosocks.Socks5Addr('localhost') | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
transp, proto = mock.Mock(), mock.Mock() | |||||
reader, writer = mock.Mock(), mock.Mock() | |||||
proto.app_protocol.reader, proto.app_protocol.writer = reader, writer | |||||
loop_mock = mock.Mock() | |||||
loop_mock.create_connection = make_mocked_coro((transp, proto)) | |||||
with mock.patch('aiosocks.asyncio.Future') as future_mock: | |||||
future_mock.side_effect = make_mocked_coro(True) | |||||
r, w = await aiosocks.open_connection(addr, auth, dst, loop=loop_mock) | |||||
assert reader is r | |||||
assert writer is w |
@@ -1,368 +1,293 @@ | |||||
import pytest | |||||
import aiosocks | import aiosocks | ||||
import asyncio | |||||
import aiohttp | import aiohttp | ||||
import unittest | |||||
from aiosocks.connector import SocksConnector | |||||
try: | |||||
from asyncio import ensure_future | |||||
except ImportError: | |||||
ensure_future = asyncio.async | |||||
from .helpers import fake_socks_srv, fake_socks4_srv, http_srv | |||||
class TestCreateSocks4Connection(unittest.TestCase): | |||||
def setUp(self): | |||||
self.loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(None) | |||||
def tearDown(self): | |||||
self.loop.close() | |||||
def test_connect_success(self): | |||||
with fake_socks_srv(self.loop, | |||||
b'\x00\x5a\x04W\x01\x01\x01\x01test') as port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks4Auth('usr') | |||||
dst = ('python.org', 80) | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
transport, protocol = self.loop.run_until_complete(coro) | |||||
self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) | |||||
data = self.loop.run_until_complete( | |||||
protocol._stream_reader.read(4)) | |||||
self.assertEqual(data, b'test') | |||||
transport.close() | |||||
def test_invalid_data(self): | |||||
with fake_socks_srv(self.loop, | |||||
b'\x01\x5a\x04W\x01\x01\x01\x01') as port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks4Auth('usr') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('invalid data', str(ct.exception)) | |||||
def test_socks_srv_error(self): | |||||
with fake_socks_srv(self.loop, | |||||
b'\x00\x5b\x04W\x01\x01\x01\x01') as port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks4Auth('usr') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('0x5b', str(ct.exception)) | |||||
class TestCreateSocks5Connect(unittest.TestCase): | |||||
def setUp(self): | |||||
self.loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(None) | |||||
def tearDown(self): | |||||
self.loop.close() | |||||
def test_connect_success_anonymous(self): | |||||
with fake_socks_srv( | |||||
self.loop, | |||||
b'\x05\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
transport, protocol = self.loop.run_until_complete(coro) | |||||
self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) | |||||
data = self.loop.run_until_complete( | |||||
protocol._stream_reader.read(4)) | |||||
self.assertEqual(data, b'test') | |||||
transport.close() | |||||
def test_connect_success_usr_pwd(self): | |||||
with fake_socks_srv( | |||||
self.loop, | |||||
b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest' | |||||
) as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
transport, protocol = self.loop.run_until_complete(coro) | |||||
self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) | |||||
data = self.loop.run_until_complete( | |||||
protocol._stream_reader.read(4)) | |||||
self.assertEqual(data, b'test') | |||||
transport.close() | |||||
def test_auth_ver_err(self): | |||||
with fake_socks_srv(self.loop, b'\x04\x02') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('invalid version', str(ct.exception)) | |||||
def test_auth_method_rejected(self): | |||||
with fake_socks_srv(self.loop, b'\x05\xFF') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('authentication methods were rejected', | |||||
str(ct.exception)) | |||||
def test_auth_status_invalid(self): | |||||
with fake_socks_srv(self.loop, b'\x05\xF0') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('invalid data', str(ct.exception)) | |||||
def test_auth_status_invalid2(self): | |||||
with fake_socks_srv(self.loop, b'\x05\x02\x02\x00') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('invalid data', str(ct.exception)) | |||||
def test_auth_failed(self): | |||||
with fake_socks_srv(self.loop, b'\x05\x02\x01\x01') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('authentication failed', str(ct.exception)) | |||||
def test_cmd_ver_err(self): | |||||
with fake_socks_srv(self.loop, | |||||
b'\x05\x02\x01\x00\x04\x00\x00') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('invalid version', str(ct.exception)) | |||||
def test_cmd_not_granted(self): | |||||
with fake_socks_srv(self.loop, | |||||
b'\x05\x02\x01\x00\x05\x01\x00') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('General SOCKS server failure', str(ct.exception)) | |||||
def test_invalid_address_type(self): | |||||
with fake_socks_srv(self.loop, | |||||
b'\x05\x02\x01\x00\x05\x00\x00\xFF') as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with self.assertRaises(aiosocks.SocksError) as ct: | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
self.loop.run_until_complete(coro) | |||||
self.assertIn('invalid data', str(ct.exception)) | |||||
def test_atype_ipv4(self): | |||||
with fake_socks_srv( | |||||
self.loop, | |||||
b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04W' | |||||
) as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
transport, protocol = self.loop.run_until_complete(coro) | |||||
self.assertEqual(protocol.proxy_sockname, ('1.1.1.1', 1111)) | |||||
transport.close() | |||||
def test_atype_ipv6(self): | |||||
with fake_socks_srv( | |||||
self.loop, | |||||
b'\x05\x02\x01\x00\x05\x00\x00\x04\x00\x00\x00\x00' | |||||
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x11\x04W' | |||||
) as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
transport, protocol = self.loop.run_until_complete(coro) | |||||
self.assertEqual(protocol.proxy_sockname, ('::111', 1111)) | |||||
transport.close() | |||||
def test_atype_domain(self): | |||||
with fake_socks_srv( | |||||
self.loop, | |||||
b'\x05\x02\x01\x00\x05\x00\x00\x03\x0apython.org\x04W' | |||||
) as port: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
coro = aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=self.loop) | |||||
transport, protocol = self.loop.run_until_complete(coro) | |||||
self.assertEqual(protocol.proxy_sockname, (b'python.org', 1111)) | |||||
transport.close() | |||||
class TestSocksConnector(unittest.TestCase): | |||||
def setUp(self): | |||||
self.loop = asyncio.new_event_loop() | |||||
asyncio.set_event_loop(None) | |||||
def tearDown(self): | |||||
self.loop.close() | |||||
def test_http_connect(self): | |||||
with fake_socks4_srv(self.loop) as proxy_port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', proxy_port) | |||||
conn = SocksConnector(proxy=addr, proxy_auth=None, loop=self.loop, | |||||
remote_resolve=False) | |||||
with http_srv(self.loop) as url: | |||||
with aiohttp.ClientSession(connector=conn, | |||||
loop=self.loop) as ses: | |||||
@asyncio.coroutine | |||||
def make_req(): | |||||
return (yield from ses.request('get', url=url)) | |||||
resp = self.loop.run_until_complete(make_req()) | |||||
self.assertEqual(resp.status, 200) | |||||
import os | |||||
import ssl | |||||
from aiohttp import web | |||||
from aiohttp.test_utils import RawTestServer | |||||
from aiosocks.test_utils import FakeSocksSrv, FakeSocks4Srv | |||||
from aiosocks.connector import ProxyConnector, ProxyClientRequest | |||||
content = self.loop.run_until_complete(resp.text()) | |||||
self.assertEqual(content, 'Test message') | |||||
resp.close() | |||||
async def test_socks4_connect_success(loop): | |||||
pld = b'\x00\x5a\x04W\x01\x01\x01\x01test' | |||||
def test_https_connect(self): | |||||
with fake_socks4_srv(self.loop) as proxy_port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', proxy_port) | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks4Auth('usr') | |||||
dst = ('python.org', 80) | |||||
conn = SocksConnector(proxy=addr, proxy_auth=None, loop=self.loop, | |||||
remote_resolve=False, verify_ssl=False) | |||||
transport, protocol = await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
with http_srv(self.loop, use_ssl=True) as url: | |||||
with aiohttp.ClientSession(connector=conn, | |||||
loop=self.loop) as ses: | |||||
@asyncio.coroutine | |||||
def make_req(): | |||||
return (yield from ses.request('get', url=url)) | |||||
assert protocol.proxy_sockname == ('1.1.1.1', 1111) | |||||
resp = self.loop.run_until_complete(make_req()) | |||||
data = await protocol._stream_reader.read(4) | |||||
assert data == b'test' | |||||
self.assertEqual(resp.status, 200) | |||||
transport.close() | |||||
content = self.loop.run_until_complete(resp.text()) | |||||
self.assertEqual(content, 'Test message') | |||||
resp.close() | |||||
async def test_socks4_invalid_data(loop): | |||||
pld = b'\x01\x5a\x04W\x01\x01\x01\x01' | |||||
def test_fingerprint_success(self): | |||||
with fake_socks4_srv(self.loop) as proxy_port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', proxy_port) | |||||
fp = (b's\x93\xfd:\xed\x08\x1do\xa9\xaeq9' | |||||
b'\x1a\xe3\xc5\x7f\x89\xe7l\xf9') | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks4Auth('usr') | |||||
dst = ('python.org', 80) | |||||
conn = SocksConnector(proxy=addr, proxy_auth=None, loop=self.loop, | |||||
remote_resolve=False, verify_ssl=False, | |||||
fingerprint=fp) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'invalid data' in str(ct) | |||||
with http_srv(self.loop, use_ssl=True) as url: | |||||
with aiohttp.ClientSession(connector=conn, | |||||
loop=self.loop) as ses: | |||||
@asyncio.coroutine | |||||
def make_req(): | |||||
return (yield from ses.request('get', url=url)) | |||||
resp = self.loop.run_until_complete(make_req()) | |||||
async def test_socks4_srv_error(loop): | |||||
pld = b'\x00\x5b\x04W\x01\x01\x01\x01' | |||||
self.assertEqual(resp.status, 200) | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks4Auth('usr') | |||||
dst = ('python.org', 80) | |||||
content = self.loop.run_until_complete(resp.text()) | |||||
self.assertEqual(content, 'Test message') | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert '0x5b' in str(ct) | |||||
resp.close() | |||||
def test_fingerprint_fail(self): | |||||
with fake_socks4_srv(self.loop) as proxy_port: | |||||
addr = aiosocks.Socks4Addr('127.0.0.1', proxy_port) | |||||
fp = (b's\x93\xfd:\xed\x08\x1do\xa9\xaeq9' | |||||
b'\x1a\xe3\xc5\x7f\x89\xe7l\x10') | |||||
async def test_socks5_connect_success_anonymous(loop): | |||||
pld = b'\x05\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest' | |||||
conn = SocksConnector(proxy=addr, proxy_auth=None, loop=self.loop, | |||||
remote_resolve=False, verify_ssl=False, | |||||
fingerprint=fp) | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with http_srv(self.loop, use_ssl=True) as url: | |||||
with aiohttp.ClientSession(connector=conn, | |||||
loop=self.loop) as ses: | |||||
@asyncio.coroutine | |||||
def make_req(): | |||||
return (yield from ses.request('get', url=url)) | |||||
transport, protocol = await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
with self.assertRaises(aiohttp.FingerprintMismatch): | |||||
self.loop.run_until_complete(make_req()) | |||||
assert protocol.proxy_sockname == ('1.1.1.1', 1111) | |||||
data = await protocol._stream_reader.read(4) | |||||
assert data == b'test' | |||||
transport.close() | |||||
async def test_socks5_connect_success_usr_pwd(loop): | |||||
pld = b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04Wtest' | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
transport, protocol = await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert protocol.proxy_sockname == ('1.1.1.1', 1111) | |||||
data = await protocol._stream_reader.read(4) | |||||
assert data == b'test' | |||||
transport.close() | |||||
async def test_socks5_auth_ver_err(loop): | |||||
async with FakeSocksSrv(loop, b'\x04\x02') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'invalid version' in str(ct) | |||||
async def test_socks5_auth_method_rejected(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\xFF') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'authentication methods were rejected' in str(ct) | |||||
async def test_socks5_auth_status_invalid(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\xF0') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'invalid data' in str(ct) | |||||
async def test_socks5_auth_status_invalid2(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\x02\x02\x00') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'invalid data' in str(ct) | |||||
async def test_socks5_auth_failed(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\x02\x01\x01') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'authentication failed' in str(ct) | |||||
async def test_socks5_cmd_ver_err(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\x02\x01\x00\x04\x00\x00') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'invalid version' in str(ct) | |||||
async def test_socks5_cmd_not_granted(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\x02\x01\x00\x05\x01\x00') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'General SOCKS server failure' in str(ct) | |||||
async def test_socks5_invalid_address_type(loop): | |||||
async with FakeSocksSrv(loop, b'\x05\x02\x01\x00\x05\x00\x00\xFF') as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
with pytest.raises(aiosocks.SocksError) as ct: | |||||
await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert 'invalid data' in str(ct) | |||||
async def test_socks5_atype_ipv4(loop): | |||||
pld = b'\x05\x02\x01\x00\x05\x00\x00\x01\x01\x01\x01\x01\x04W' | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
transport, protocol = await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert protocol.proxy_sockname == ('1.1.1.1', 1111) | |||||
transport.close() | |||||
async def test_socks5_atype_ipv6(loop): | |||||
pld = b'\x05\x02\x01\x00\x05\x00\x00\x04\x00\x00\x00\x00' \ | |||||
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x11\x04W' | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
transport, protocol = await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert protocol.proxy_sockname == ('::111', 1111) | |||||
transport.close() | |||||
async def test_socks5_atype_domain(loop): | |||||
pld = b'\x05\x02\x01\x00\x05\x00\x00\x03\x0apython.org\x04W' | |||||
async with FakeSocksSrv(loop, pld) as srv: | |||||
addr = aiosocks.Socks5Addr('127.0.0.1', srv.port) | |||||
auth = aiosocks.Socks5Auth('usr', 'pwd') | |||||
dst = ('python.org', 80) | |||||
transport, protocol = await aiosocks.create_connection( | |||||
None, addr, auth, dst, loop=loop) | |||||
assert protocol.proxy_sockname == (b'python.org', 1111) | |||||
transport.close() | |||||
async def test_http_connect(loop): | |||||
async def handler(request): | |||||
return web.Response(text='Test message') | |||||
async with RawTestServer(handler, host='127.0.0.1', loop=loop) as ws: | |||||
async with FakeSocks4Srv(loop) as srv: | |||||
conn = ProxyConnector(loop=loop, remote_resolve=False) | |||||
async with aiohttp.ClientSession( | |||||
connector=conn, loop=loop, | |||||
request_class=ProxyClientRequest) as ses: | |||||
proxy = 'socks4://127.0.0.1:{}'.format(srv.port) | |||||
async with ses.get(ws.make_url('/'), proxy=proxy) as resp: | |||||
assert resp.status == 200 | |||||
assert (await resp.text()) == 'Test message' | |||||
async def test_https_connect(loop): | |||||
async def handler(request): | |||||
return web.Response(text='Test message') | |||||
here = os.path.join(os.path.dirname(__file__), '..', 'tests') | |||||
keyfile = os.path.join(here, 'sample.key') | |||||
certfile = os.path.join(here, 'sample.crt') | |||||
sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23) | |||||
sslcontext.load_cert_chain(certfile, keyfile) | |||||
ws = RawTestServer(handler, scheme='https', host='127.0.0.1', loop=loop) | |||||
await ws.start_server(loop=loop, ssl=sslcontext) | |||||
v_fp = b's\x93\xfd:\xed\x08\x1do\xa9\xaeq9\x1a\xe3\xc5\x7f\x89\xe7l\xf9' | |||||
inv_fp = b's\x93\xfd:\xed\x08\x1do\xa9\xaeq9\x1a\xe3\xc5\x7f\x89\xe7l\x10' | |||||
async with FakeSocks4Srv(loop) as srv: | |||||
v_conn = ProxyConnector(loop=loop, remote_resolve=False, | |||||
verify_ssl=False, fingerprint=v_fp) | |||||
inv_conn = ProxyConnector(loop=loop, remote_resolve=False, | |||||
verify_ssl=False, fingerprint=inv_fp) | |||||
async with aiohttp.ClientSession( | |||||
connector=v_conn, loop=loop, | |||||
request_class=ProxyClientRequest) as ses: | |||||
proxy = 'socks4://127.0.0.1:{}'.format(srv.port) | |||||
async with ses.get(ws.make_url('/'), proxy=proxy) as resp: | |||||
assert resp.status == 200 | |||||
assert (await resp.text()) == 'Test message' | |||||
async with aiohttp.ClientSession( | |||||
connector=inv_conn, loop=loop, | |||||
request_class=ProxyClientRequest) as ses: | |||||
proxy = 'socks4://127.0.0.1:{}'.format(srv.port) | |||||
with pytest.raises(aiohttp.ServerFingerprintMismatch): | |||||
async with ses.get(ws.make_url('/'), proxy=proxy) as resp: | |||||
assert resp.status == 200 |