diff --git a/README.rst b/README.rst index bbb1f4e..b284639 100644 --- a/README.rst +++ b/README.rst @@ -114,7 +114,7 @@ aiohttp usage import asyncio import aiohttp import aiosocks - from aiosocks.connector import SocksConnector + from aiosocks.connector import SocksConnector, proxy_connector async def load_github_main(): @@ -122,11 +122,21 @@ aiohttp usage auth = aiosocks.Socks5Auth('proxyuser1', password='pwd') # remote resolve - # conn = SocksConnector(proxy=addr, proxy_auth=auth, remote_resolve=True) + conn = SocksConnector(proxy=addr, proxy_auth=auth, remote_resolve=True) # or locale resolve conn = SocksConnector(proxy=addr, proxy_auth=auth, remote_resolve=False) + # or use shortcut function for automatically create + # SocksConnector/aiohttp.ProxyConnector (socks or http proxy) + conn = proxy_connector(aiosocks.SocksAddr(...), + remote_resolve=True, verify_ssl=False) + # return SocksConnector + + conn = proxy_connector(aiosocks.HttpProxyAddr('http://proxy'), + aiosocks.HttpProxyAuth('login', 'pwd')) + # return aiohttp.ProxyConnector (http proxy connector) + try: with aiohttp.ClientSession(connector=conn) as ses: async with session.get('http://github.com/') as resp: diff --git a/aiosocks/__init__.py b/aiosocks/__init__.py index 9f20c38..92f263a 100644 --- a/aiosocks/__init__.py +++ b/aiosocks/__init__.py @@ -4,17 +4,19 @@ from .errors import ( SocksConnectionError, InvalidServerReply, InvalidServerVersion ) from .helpers import ( - SocksAddr, Socks4Addr, Socks5Addr, Socks4Auth, Socks5Auth + SocksAddr, Socks4Addr, Socks5Addr, Socks4Auth, + Socks5Auth, HttpProxyAddr, HttpProxyAuth ) from .protocols import Socks4Protocol, Socks5Protocol, DEFAULT_LIMIT __version__ = '0.1.4' __all__ = ('Socks4Protocol', 'Socks5Protocol', 'Socks4Auth', - 'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'SocksError', - 'NoAcceptableAuthMethods', 'LoginAuthenticationFailed', - 'SocksConnectionError', 'InvalidServerVersion', - 'InvalidServerReply', 'create_connection', 'open_connection') + 'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'HttpProxyAddr', + 'HttpProxyAuth', 'SocksError', 'NoAcceptableAuthMethods', + 'LoginAuthenticationFailed', 'SocksConnectionError', + 'InvalidServerVersion', 'InvalidServerReply', + 'create_connection', 'open_connection') @asyncio.coroutine diff --git a/aiosocks/connector.py b/aiosocks/connector.py index 93c376e..948e007 100644 --- a/aiosocks/connector.py +++ b/aiosocks/connector.py @@ -4,6 +4,7 @@ import aiohttp import ipaddress from aiohttp.errors import ProxyConnectionError from .errors import SocksError, SocksConnectionError +from .helpers import HttpProxyAddr, SocksAddr from . import create_connection __all__ = ('SocksConnector',) @@ -108,3 +109,13 @@ class SocksConnector(aiohttp.TCPConnector): raise aiohttp.ClientOSError( exc.errno, 'Can not connect to %s:%s [%s]' % (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') diff --git a/aiosocks/helpers.py b/aiosocks/helpers.py index 730a865..1b97d59 100644 --- a/aiosocks/helpers.py +++ b/aiosocks/helpers.py @@ -1,6 +1,8 @@ from collections import namedtuple +from aiohttp.helpers import BasicAuth as HttpProxyAuth -__all__ = ('Socks4Auth', 'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'SocksAddr') +__all__ = ('Socks4Auth', 'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'SocksAddr', + 'HttpProxyAddr', 'HttpProxyAuth') class Socks4Auth(namedtuple('Socks4Auth', ['login', 'encoding'])): @@ -41,3 +43,10 @@ class Socks4Addr(SocksAddr): class Socks5Addr(SocksAddr): pass + + +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) diff --git a/tests/test_connector.py b/tests/test_connector.py index b0afd33..2024c28 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -2,9 +2,10 @@ import unittest import asyncio import aiosocks import aiohttp +import pytest from unittest import mock from aiohttp.client_reqrep import ClientRequest -from aiosocks.connector import SocksConnector +from aiosocks.connector import SocksConnector, proxy_connector from .helpers import fake_coroutine @@ -131,3 +132,19 @@ class TestSocksConnector(unittest.TestCase): with self.assertRaises(aiosocks.SocksError): self.loop.run_until_complete(connector.connect(req)) + + +def test_proxy_connector(): + socks4_addr = aiosocks.Socks4Addr('h') + socks5_addr = aiosocks.Socks5Addr('h') + http_addr = aiosocks.HttpProxyAddr('http://proxy') + + loop = asyncio.new_event_loop() + + assert isinstance(proxy_connector(socks4_addr, loop=loop), SocksConnector) + assert isinstance(proxy_connector(socks5_addr, loop=loop), SocksConnector) + assert isinstance(proxy_connector(http_addr, loop=loop), + aiohttp.ProxyConnector) + + with pytest.raises(ValueError): + proxy_connector(None) diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 8962060..bd9977c 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -83,3 +83,11 @@ def test_socks5_addr4(): addr = aiosocks.Socks5Addr('localhost', None) assert addr.host == 'localhost' assert addr.port == 1080 + + +def test_http_proxy_addr(): + addr = aiosocks.HttpProxyAddr('http://proxy') + assert addr.url == 'http://proxy' + + with pytest.raises(ValueError): + aiosocks.HttpProxyAddr(None)