From 15bcc0ba9f4007c496273ab7472d66a1caab35b8 Mon Sep 17 00:00:00 2001 From: Martin Winks <50446230+uwinx@users.noreply.github.com> Date: Sat, 2 May 2020 17:12:53 +0400 Subject: [PATCH] feat(proxy): proxy for aiohttp,base sessions (#284) * feat(proxy): proxy for aiohttp,base sessions Add support for proxies in aiohttp session with aiohttp_socks library, edit BaseSession class to support proxies for other sessions in future. * fix(annotation): missing underscore before "private" typevar * chore: remove redundant of proxy_url schema for socks version * test: add missing test Add missing test, remove BaseSession.cfg and switch to implementing class' "private" traits, add aiohttp_socks in dependency list as optional and extra. * feat(session): Implement asyncio session for requests [wip] * feat(proxy chain): Chained proxy support in aiohttp session Add ChainProxyConnector support, !pin pydantic to "1.4", add documentation on aiohttp connector. * style(mypy): apply linter changes * tests(mock): remove await for magic mock * fix dangling dependency * refactor(generic): get rid of generic behaviour for base session --- aiogram/api/client/bot.py | 2 +- aiogram/api/client/session/aiohttp.py | 86 +++++++++++++++++-- aiogram/api/client/session/base.py | 5 +- docs/api/client/session/aiohttp.md | 76 ++++++++++++++++ mkdocs.yml | 3 + poetry.lock | 24 +++++- pyproject.toml | 3 + .../test_session/test_aiohttp_session.py | 54 ++++++++++++ 8 files changed, 245 insertions(+), 8 deletions(-) create mode 100644 docs/api/client/session/aiohttp.md diff --git a/aiogram/api/client/bot.py b/aiogram/api/client/bot.py index 27293d83..c33105e4 100644 --- a/aiogram/api/client/bot.py +++ b/aiogram/api/client/bot.py @@ -121,7 +121,7 @@ class Bot(ContextInstanceMixin["Bot"]): """ def __init__( - self, token: str, session: Optional[BaseSession] = None, parse_mode: Optional[str] = None + self, token: str, session: Optional[BaseSession] = None, parse_mode: Optional[str] = None, ) -> None: validate_token(token) diff --git a/aiogram/api/client/session/aiohttp.py b/aiogram/api/client/session/aiohttp.py index b123bcfc..c3c1bba5 100644 --- a/aiogram/api/client/session/aiohttp.py +++ b/aiogram/api/client/session/aiohttp.py @@ -1,29 +1,105 @@ from __future__ import annotations -from typing import AsyncGenerator, Callable, Optional, TypeVar, cast +from typing import ( + Any, + AsyncGenerator, + Callable, + Dict, + Iterable, + List, + Optional, + Tuple, + Type, + TypeVar, + Union, + cast, +) -from aiohttp import ClientSession, ClientTimeout, FormData +from aiohttp import BasicAuth, ClientSession, ClientTimeout, FormData, TCPConnector from aiogram.api.methods import Request, TelegramMethod from .base import PRODUCTION, BaseSession, TelegramAPIServer T = TypeVar("T") +_ProxyBasic = Union[str, Tuple[str, BasicAuth]] +_ProxyChain = Iterable[_ProxyBasic] +_ProxyType = Union[_ProxyChain, _ProxyBasic] + + +def _retrieve_basic(basic: _ProxyBasic) -> Dict[str, Any]: + from aiohttp_socks.utils import parse_proxy_url # type: ignore + + proxy_auth: Optional[BasicAuth] = None + + if isinstance(basic, str): + proxy_url = basic + else: + proxy_url, proxy_auth = basic + + proxy_type, host, port, username, password = parse_proxy_url(proxy_url) + if isinstance(proxy_auth, BasicAuth): + username = proxy_auth.login + password = proxy_auth.password + + return dict( + proxy_type=proxy_type, + host=host, + port=port, + username=username, + password=password, + rdns=True, + ) + + +def _prepare_connector(chain_or_plain: _ProxyType) -> Tuple[Type["TCPConnector"], Dict[str, Any]]: + from aiohttp_socks import ProxyInfo, ProxyConnector, ChainProxyConnector # type: ignore + + # since tuple is Iterable(compatible with _ProxyChain) object, we assume that + # user wants chained proxies if tuple is a pair of string(url) and BasicAuth + if isinstance(chain_or_plain, str) or ( + isinstance(chain_or_plain, tuple) and len(chain_or_plain) == 2 + ): + chain_or_plain = cast(_ProxyBasic, chain_or_plain) + return ProxyConnector, _retrieve_basic(chain_or_plain) + + chain_or_plain = cast(_ProxyChain, chain_or_plain) + infos: List[ProxyInfo] = [] + for basic in chain_or_plain: + infos.append(ProxyInfo(**_retrieve_basic(basic))) + + return ChainProxyConnector, dict(proxy_infos=infos) class AiohttpSession(BaseSession): def __init__( self, api: TelegramAPIServer = PRODUCTION, - json_loads: Optional[Callable[..., str]] = None, + json_loads: Optional[Callable[..., Any]] = None, json_dumps: Optional[Callable[..., str]] = None, + proxy: Optional[_ProxyType] = None, ): - super(AiohttpSession, self).__init__(api=api, json_loads=json_loads, json_dumps=json_dumps) + super(AiohttpSession, self).__init__( + api=api, json_loads=json_loads, json_dumps=json_dumps, proxy=proxy + ) self._session: Optional[ClientSession] = None + self._connector_type: Type[TCPConnector] = TCPConnector + self._connector_init: Dict[str, Any] = {} + + if self.proxy: + try: + self._connector_type, self._connector_init = _prepare_connector( + cast(_ProxyType, self.proxy) + ) + except ImportError as exc: # pragma: no cover + raise UserWarning( + "In order to use aiohttp client for proxy requests, install " + "https://pypi.org/project/aiohttp-socks/" + ) from exc async def create_session(self) -> ClientSession: if self._session is None or self._session.closed: - self._session = ClientSession() + self._session = ClientSession(connector=self._connector_type(**self._connector_init)) return self._session diff --git a/aiogram/api/client/session/base.py b/aiogram/api/client/session/base.py index 700810f1..83e7f3ff 100644 --- a/aiogram/api/client/session/base.py +++ b/aiogram/api/client/session/base.py @@ -12,14 +12,16 @@ from ...methods import Response, TelegramMethod from ..telegram import PRODUCTION, TelegramAPIServer T = TypeVar("T") +PT = TypeVar("PT") class BaseSession(abc.ABC): def __init__( self, api: Optional[TelegramAPIServer] = None, - json_loads: Optional[Callable[..., str]] = None, + json_loads: Optional[Callable[..., Any]] = None, json_dumps: Optional[Callable[..., str]] = None, + proxy: Optional[PT] = None, ) -> None: if api is None: api = PRODUCTION @@ -31,6 +33,7 @@ class BaseSession(abc.ABC): self.api = api self.json_loads = json_loads self.json_dumps = json_dumps + self.proxy = proxy def raise_for_status(self, response: Response[T]) -> None: if response.ok: diff --git a/docs/api/client/session/aiohttp.md b/docs/api/client/session/aiohttp.md new file mode 100644 index 00000000..9eab5ede --- /dev/null +++ b/docs/api/client/session/aiohttp.md @@ -0,0 +1,76 @@ +# Aiohttp session + +AiohttpSession represents a wrapper-class around `ClientSession` from [aiohttp]('https://pypi.org/project/aiohttp/') + +Currently `AiohttpSession` is a default session used in `aiogram.Bot` + +## Usage example + +```python +from aiogram import Bot +from aiogram.api.client.session.aiohttp import AiohttpSession + +session = AiohttpSession() +Bot('token', session=session) +``` + + +## Proxy requests in AiohttpSession + +In order to use AiohttpSession with proxy connector you have to install [aiohttp-socks]('https://pypi.org/project/aiohttp-socks/') + +Binding session to bot: +```python +from aiogram import Bot +from aiogram.api.client.session.aiohttp import AiohttpSession + +session = AiohttpSession(proxy="protocol://host:port/") +Bot(token="bot token", session=session) +``` + +!!! note "Protocols" + Only following protocols are supported: http(tunneling), socks4(a), socks5 as aiohttp_socks documentation claims. + + +### Authorization + +Proxy authorization credentials can be specified in proxy URL or come as an instance of `aiohttp.BasicAuth` containing +login and password. + +Consider examples: +```python +from aiohttp import BasicAuth +from aiogram.api.client.session.aiohttp import AiohttpSession + +auth = BasicAuth(login="user", password="password") +session = AiohttpSession(proxy=("protocol://host:port", auth)) +# or simply include your basic auth credential in URL +session = AiohttpSession(proxy="protocol://user:password@host:port") +``` + +!!! note "Credential priorities" + Aiogram prefers `BasicAuth` over username and password in URL, so + if proxy URL contains login and password and `BasicAuth` object is passed at the same time + aiogram will use login and password from `BasicAuth` instance. + + +### Proxy chains + +Since [aiohttp-socks]('https://pypi.org/project/aiohttp-socks/') supports proxy chains, you're able to use them in aiogram + +Example of chain proxies: +```python +from aiohttp import BasicAuth +from aiogram.api.client.session.aiohttp import AiohttpSession + +auth = BasicAuth(login="user", password="password") +session = AiohttpSession( + proxy={"protocol0://host0:port0", + "protocol1://user:password@host1:port1", + ("protocol2://host2:port2", auth),} # can be any iterable if not set +) +``` + +## Location + +- `from aiogram.api.client.session.aiohttp import AiohttpSession` diff --git a/mkdocs.yml b/mkdocs.yml index 747b5f4e..0d77e640 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -41,6 +41,9 @@ nav: - install.md - Bot API: - api/index.md + - Client: + - Session: + - aiohttp: api/client/session/aiohttp.md - Methods: - Available methods: api/methods/index.md - Getting updates: diff --git a/poetry.lock b/poetry.lock index ef800fbf..0a5da5a0 100644 --- a/poetry.lock +++ b/poetry.lock @@ -24,6 +24,18 @@ yarl = ">=1.0,<2.0" [package.extras] speedups = ["aiodns", "brotlipy", "cchardet"] +[[package]] +category = "main" +description = "Proxy connector for aiohttp" +name = "aiohttp-socks" +optional = false +python-versions = "*" +version = "0.3.9" + +[package.dependencies] +aiohttp = ">=2.3.2" +attrs = ">=19.2.0" + [[package]] category = "dev" description = "A small Python module for determining appropriate platform-specific dirs, e.g. a \"user data dir\"." @@ -931,9 +943,10 @@ testing = ["jaraco.itertools", "func-timeout"] [extras] fast = ["uvloop"] +proxy = ["aiohttp-socks"] [metadata] -content-hash = "6e9fa892cd316d3e39bebc5ad146175716a8da246e7def800d66cd60f63cd630" +content-hash = "8c2dc4c18c8de6ffe48c634e96e9a057d4a2ef211d21459fa4c8e23b87dd8456" python-versions = "^3.7" [metadata.files] @@ -955,6 +968,10 @@ aiohttp = [ {file = "aiohttp-3.6.2-py3-none-any.whl", hash = "sha256:460bd4237d2dbecc3b5ed57e122992f60188afe46e7319116da5eb8a9dfedba4"}, {file = "aiohttp-3.6.2.tar.gz", hash = "sha256:259ab809ff0727d0e834ac5e8a283dc5e3e0ecc30c4d80b3cd17a4139ce1f326"}, ] +aiohttp-socks = [ + {file = "aiohttp_socks-0.3.9-py3-none-any.whl", hash = "sha256:ccd483d7677d7ba80b7ccb738a9be27a3ad6dce4b2756509bc71c9d679d96105"}, + {file = "aiohttp_socks-0.3.9.tar.gz", hash = "sha256:5e5638d0e472baa441eab7990cf19e034960cc803f259748cc359464ccb3c2d6"}, +] appdirs = [ {file = "appdirs-1.4.3-py2.py3-none-any.whl", hash = "sha256:d8b24664561d0d34ddfaec54636d502d7cea6e29c3eaf68f3df6180863e2166e"}, {file = "appdirs-1.4.3.tar.gz", hash = "sha256:9e5896d1372858f8dd3344faf4e5014d21849c756c8d5701f78f8a103b372d92"}, @@ -1166,6 +1183,11 @@ markupsafe = [ {file = "MarkupSafe-1.1.1-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:ba59edeaa2fc6114428f1637ffff42da1e311e29382d81b339c1817d37ec93c6"}, {file = "MarkupSafe-1.1.1-cp37-cp37m-win32.whl", hash = "sha256:b00c1de48212e4cc9603895652c5c410df699856a2853135b3967591e4beebc2"}, {file = "MarkupSafe-1.1.1-cp37-cp37m-win_amd64.whl", hash = "sha256:9bf40443012702a1d2070043cb6291650a0841ece432556f784f004937f0f32c"}, + {file = "MarkupSafe-1.1.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:6788b695d50a51edb699cb55e35487e430fa21f1ed838122d722e0ff0ac5ba15"}, + {file = "MarkupSafe-1.1.1-cp38-cp38-manylinux1_i686.whl", hash = "sha256:cdb132fc825c38e1aeec2c8aa9338310d29d337bebbd7baa06889d09a60a1fa2"}, + {file = "MarkupSafe-1.1.1-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:13d3144e1e340870b25e7b10b98d779608c02016d5184cfb9927a9f10c689f42"}, + {file = "MarkupSafe-1.1.1-cp38-cp38-win32.whl", hash = "sha256:596510de112c685489095da617b5bcbbac7dd6384aeebeda4df6025d0256a81b"}, + {file = "MarkupSafe-1.1.1-cp38-cp38-win_amd64.whl", hash = "sha256:e8313f01ba26fbbe36c7be1966a7b7424942f670f38e666995b88d012765b9be"}, {file = "MarkupSafe-1.1.1.tar.gz", hash = "sha256:29872e92839765e546828bb7754a68c418d927cd064fd4708fab9fe9c8bb116b"}, ] mccabe = [ diff --git a/pyproject.toml b/pyproject.toml index cca84cd1..7cecf764 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ Babel = "^2.7" aiofiles = "^0.4.0" uvloop = {version = "^0.14.0", markers = "sys_platform == 'darwin' or sys_platform == 'linux'", optional = true} async_lru = "^1.0" +aiohttp-socks = {version = "^0.3.8", optional = true} [tool.poetry.dev-dependencies] uvloop = {version = "^0.14.0", markers = "sys_platform == 'darwin' or sys_platform == 'linux'"} @@ -63,9 +64,11 @@ pymdown-extensions = "^6.1" lxml = "^4.4" ipython = "^7.10" markdown-include = "^0.5.1" +aiohttp-socks = "^0.3.4" [tool.poetry.extras] fast = ["uvloop"] +proxy = ["aiohttp-socks"] [tool.black] line-length = 99 diff --git a/tests/test_api/test_client/test_session/test_aiohttp_session.py b/tests/test_api/test_client/test_session/test_aiohttp_session.py index ec1e16a5..2587a686 100644 --- a/tests/test_api/test_client/test_session/test_aiohttp_session.py +++ b/tests/test_api/test_client/test_session/test_aiohttp_session.py @@ -1,6 +1,7 @@ from typing import AsyncContextManager, AsyncGenerator import aiohttp +import aiohttp_socks import pytest from aresponses import ResponsesMockServer @@ -29,6 +30,59 @@ class TestAiohttpSession: assert session._session is not None assert isinstance(aiohttp_session, aiohttp.ClientSession) + @pytest.mark.asyncio + async def test_create_proxy_session(self): + session = AiohttpSession( + proxy=("socks5://proxy.url/", aiohttp.BasicAuth("login", "password", "encoding")) + ) + + assert session._connector_type == aiohttp_socks.ProxyConnector + + assert isinstance(session._connector_init, dict) + assert session._connector_init["proxy_type"] is aiohttp_socks.ProxyType.SOCKS5 + + aiohttp_session = await session.create_session() + assert isinstance(aiohttp_session.connector, aiohttp_socks.ProxyConnector) + + @pytest.mark.asyncio + async def test_create_proxy_session_proxy_url(self): + session = AiohttpSession(proxy="socks4://proxy.url/") + + assert isinstance(session.proxy, str) + + assert isinstance(session._connector_init, dict) + assert session._connector_init["proxy_type"] is aiohttp_socks.ProxyType.SOCKS4 + + aiohttp_session = await session.create_session() + assert isinstance(aiohttp_session.connector, aiohttp_socks.ProxyConnector) + + @pytest.mark.asyncio + async def test_create_proxy_session_chained_proxies(self): + session = AiohttpSession( + proxy=[ + "socks4://proxy.url/", + "socks5://proxy.url/", + "http://user:password@127.0.0.1:3128", + ] + ) + + assert isinstance(session.proxy, list) + + assert isinstance(session._connector_init, dict) + assert isinstance(session._connector_init["proxy_infos"], list) + assert isinstance(session._connector_init["proxy_infos"][0], aiohttp_socks.ProxyInfo) + + assert ( + session._connector_init["proxy_infos"][0].proxy_type is aiohttp_socks.ProxyType.SOCKS4 + ) + assert ( + session._connector_init["proxy_infos"][1].proxy_type is aiohttp_socks.ProxyType.SOCKS5 + ) + assert session._connector_init["proxy_infos"][2].proxy_type is aiohttp_socks.ProxyType.HTTP + + aiohttp_session = await session.create_session() + assert isinstance(aiohttp_session.connector, aiohttp_socks.ChainProxyConnector) + @pytest.mark.asyncio async def test_close_session(self): session = AiohttpSession()