Merge pull request #14 from nibrag/proxy_from_env

Proxy from env

README.rst

@@ -15,7 +15,7 @@ SOCKS proxy client for asyncio and aiohttp
 Dependencies
 ------------
 python 3.5+
-aiohttp 2.0+
+aiohttp 2.1+
 
 Features
 --------
@@ -185,3 +185,22 @@ aiohttp usage
     loop = asyncio.get_event_loop()
     loop.run_until_complete(load_github_main())
     loop.close()
+
+Proxy from environment
+^^^^^^^^^^^^^^^^^^^^^^
+
+.. code-block:: python
+
+  import os
+  from aiosocks.connector import ProxyConnector, ProxyClientRequest
+
+  os.environ['socks4_proxy'] = 'socks4://127.0.0.1:333'
+  # or
+  os.environ['socks5_proxy'] = 'socks5://127.0.0.1:444'
+
+  conn = ProxyConnector()
+
+  with aiohttp.ClientSession(connector=conn, request_class=ProxyClientRequest) as session:
+        async with session.get('http://github.com/', proxy_from_env=True) as resp:
+          if resp.status == 200:
+            print(await resp.text())

aiosocks/__init__.py

@@ -8,7 +8,7 @@
 )
 from .protocols import Socks4Protocol, Socks5Protocol, DEFAULT_LIMIT
 
-__version__ = '0.2.2'
+__version__ = '0.2.3'
 
 __all__ = ('Socks4Protocol', 'Socks5Protocol', 'Socks4Auth',
            'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'SocksError',

aiosocks/connector.py

@@ -4,6 +4,9 @@
 except ImportError:
     raise ImportError('aiosocks.SocksConnector require aiohttp library')
 
+from yarl import URL
+from urllib.request import getproxies
+
 from .errors import SocksError, SocksConnectionError
 from .helpers import Socks4Auth, Socks5Auth, Socks4Addr, Socks5Addr
 from . import create_connection
@@ -12,7 +15,16 @@
 
 
 class ProxyClientRequest(aiohttp.ClientRequest):
-    def update_proxy(self, proxy, proxy_auth):
+    def update_proxy(self, proxy, proxy_auth, proxy_from_env):
+        if proxy_from_env and not proxy:
+            proxies = getproxies()
+
+            proxy_url = proxies.get(self.original_url.scheme)
+            if not proxy_url:
+                proxy_url = proxies.get('socks4') or proxies.get('socks5')
+
+            proxy = URL(proxy_url) if proxy_url else None
+
         if proxy and proxy.scheme not in ['http', 'socks4', 'socks5']:
             raise ValueError(
                 "Only http, socks4 and socks5 proxies are supported")

tests/test_connector.py

@@ -177,3 +177,38 @@ def test_proxy_client_request_invalid(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)
+
+
+def test_proxy_from_env_http(loop):
+    proxies = {'http': 'http://proxy.org'}
+
+    with mock.patch('aiosocks.connector.getproxies',  return_value=proxies):
+        req = ProxyClientRequest('GET', URL('http://python.org'), loop=loop)
+        req.update_proxy(None, None, True)
+        assert req.proxy == URL('http://proxy.org')
+
+        req.original_url = URL('https://python.org')
+        req.update_proxy(None, None, True)
+        assert req.proxy is None
+
+        proxies.update({'https': 'http://proxy.org',
+                        'socks4': 'socks4://127.0.0.1:33',
+                        'socks5': 'socks5://localhost:44'})
+        req.update_proxy(None, None, True)
+        assert req.proxy == URL('http://proxy.org')
+
+
+def test_proxy_from_env_socks(loop):
+    proxies = {'socks4': 'socks4://127.0.0.1:33',
+               'socks5': 'socks5://localhost:44'}
+
+    with mock.patch('aiosocks.connector.getproxies', return_value=proxies):
+        req = ProxyClientRequest('GET', URL('http://python.org'), loop=loop)
+
+        req.update_proxy(None, None, True)
+        assert req.proxy == URL('socks4://127.0.0.1:33')
+
+        del proxies['socks4']
+
+        req.update_proxy(None, None, True)
+        assert req.proxy == URL('socks5://localhost:44')