From aa289cdd935bc1ce0884f52e6fbc08fe56ddf95a Mon Sep 17 00:00:00 2001 From: mpa Date: Thu, 19 Mar 2020 05:35:22 +0400 Subject: [PATCH] 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. --- aiogram/api/client/session/aiohttp.py | 51 +++++++++++++++++-- aiogram/api/client/session/base.py | 10 +++- poetry.lock | 23 ++++++++- pyproject.toml | 1 + .../test_session/test_aiohttp_session.py | 15 ++++++ .../test_session/test_base_session.py | 9 +++- 6 files changed, 99 insertions(+), 10 deletions(-) diff --git a/aiogram/api/client/session/aiohttp.py b/aiogram/api/client/session/aiohttp.py index 06a1c99d..3a05ca9e 100644 --- a/aiogram/api/client/session/aiohttp.py +++ b/aiogram/api/client/session/aiohttp.py @@ -1,29 +1,70 @@ from __future__ import annotations -from typing import AsyncGenerator, Callable, Optional, TypeVar, cast +from typing import AsyncGenerator, Callable, Optional, TypeVar, Tuple, Dict, Any, cast -from aiohttp import ClientSession, ClientTimeout, FormData +from aiohttp import ClientSession, ClientTimeout, FormData, BasicAuth, TCPConnector from aiogram.api.methods import Request, TelegramMethod from .base import PRODUCTION, BaseSession, TelegramAPIServer T = TypeVar("T") +_ProxyType = Tuple[str, BasicAuth] -class AiohttpSession(BaseSession): +class AiohttpSession(BaseSession[_ProxyType]): def __init__( self, api: TelegramAPIServer = PRODUCTION, + proxy: Optional[_ProxyType] = None, json_loads: Optional[Callable] = None, json_dumps: Optional[Callable] = 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.cfg.connector_type = TCPConnector + self.cfg.connector_init = cast(Dict[str, Any], {}) + + if self.proxy: + proxy_url, proxy_auth = self.proxy + + try: + from aiohttp_socks import ProxyConnector + from aiohttp_socks.utils import parse_proxy_url + except ImportError as exc: + raise UserWarning( + "In order to use aiohttp client for proxy requests, install " + "https://pypi.org/project/aiohttp-socks/0.3.4" + ) from exc + + if proxy_url.startswith('socks5://') or proxy_url.startswith('socks4://'): + self.cfg.connector_type = ProxyConnector + + proxy_type, host, port, username, password = parse_proxy_url(proxy_url) + if proxy_auth: + if not username: + username = proxy_auth.login + if not password: + password = proxy_auth.password + + self.cfg.connector_init.update( + dict( + proxy_type=proxy_type, host=host, port=port, + username=username, password=password, + rdns=True, + ) + ) async def create_session(self) -> ClientSession: if self._session is None or self._session.closed: - self._session = ClientSession() + self._session = ClientSession( + connector=self.cfg.connector_type(**self.cfg.connector_init) + ) return self._session diff --git a/aiogram/api/client/session/base.py b/aiogram/api/client/session/base.py index f31dd451..f14d07bd 100644 --- a/aiogram/api/client/session/base.py +++ b/aiogram/api/client/session/base.py @@ -3,7 +3,8 @@ from __future__ import annotations import abc import datetime import json -from typing import Any, AsyncGenerator, Callable, Optional, TypeVar, Union +import types +from typing import Any, AsyncGenerator, Callable, Optional, TypeVar, Union, Generic from aiogram.utils.exceptions import TelegramAPIError @@ -11,12 +12,14 @@ from ...methods import Response, TelegramMethod from ..telegram import PRODUCTION, TelegramAPIServer T = TypeVar("T") +_ProxyType = TypeVar("_ProxyType") -class BaseSession(abc.ABC): +class BaseSession(abc.ABC, Generic[_ProxyType]): def __init__( self, api: Optional[TelegramAPIServer] = None, + proxy: Optional[ProxyType] = None, json_loads: Optional[Callable[[Any], Any]] = None, json_dumps: Optional[Callable[[Any], Any]] = None, ) -> None: @@ -30,6 +33,9 @@ class BaseSession(abc.ABC): self.api = api self.json_loads = json_loads self.json_dumps = json_dumps + self.proxy = proxy + + self.cfg: types.SimpleNamespace = types.SimpleNamespace() def raise_for_status(self, response: Response[T]) -> None: if response.ok: diff --git a/poetry.lock b/poetry.lock index f1088435..b6b1dbba 100644 --- a/poetry.lock +++ b/poetry.lock @@ -24,6 +24,18 @@ yarl = ">=1.0,<2.0" [package.extras] speedups = ["aiodns", "brotlipy", "cchardet"] +[[package]] +category = "dev" +description = "Proxy connector for aiohttp" +name = "aiohttp-socks" +optional = false +python-versions = "*" +version = "0.3.4" + +[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\"." @@ -861,7 +873,7 @@ testing = ["jaraco.itertools", "func-timeout"] fast = ["uvloop"] [metadata] -content-hash = "2eb50b5b57d0fac4780f1eb3f92ff129d891fd346e0c00856c1a56c58feffb03" +content-hash = "2a95c67ca1bea20c29bf63542bf5ace592b9ee5dcc70e5bf8ed304bd9b37ded2" python-versions = "^3.7" [metadata.files] @@ -883,6 +895,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.4-py3-none-any.whl", hash = "sha256:654203308e24aa4f012f96d4af7675404565f4cc39a113d110e3d345ede0838f"}, + {file = "aiohttp_socks-0.3.4.tar.gz", hash = "sha256:cdea1d99c14fd3884968a34a6ecb2784f4c48e657aec99ec1b986a10812287bd"}, +] appdirs = [ {file = "appdirs-1.4.3-py2.py3-none-any.whl", hash = "sha256:d8b24664561d0d34ddfaec54636d502d7cea6e29c3eaf68f3df6180863e2166e"}, {file = "appdirs-1.4.3.tar.gz", hash = "sha256:9e5896d1372858f8dd3344faf4e5014d21849c756c8d5701f78f8a103b372d92"}, @@ -1083,6 +1099,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 119fe577..51d7a605 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ 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"] 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..9d61f92f 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,20 @@ 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.cfg.connector_type == aiohttp_socks.ProxyConnector + + assert isinstance(session.cfg.connector_init, dict) + assert session.cfg.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_close_session(self): session = AiohttpSession() diff --git a/tests/test_api/test_client/test_session/test_base_session.py b/tests/test_api/test_client/test_session/test_base_session.py index cc565a42..1af89ae3 100644 --- a/tests/test_api/test_client/test_session/test_base_session.py +++ b/tests/test_api/test_client/test_session/test_base_session.py @@ -1,5 +1,6 @@ +import types import datetime -from typing import AsyncContextManager, AsyncGenerator +from typing import AsyncContextManager, AsyncGenerator, Any import pytest @@ -14,7 +15,7 @@ except ImportError: from unittest.mock import AsyncMock as CoroutineMock, patch # type: ignore -class CustomSession(BaseSession): +class CustomSession(BaseSession[Any]): async def close(self): pass @@ -44,6 +45,10 @@ class TestBaseSession(DataMixin): session = CustomSession(api=api) assert session.api == api + def test_init_cfg_namespace(self): + session = CustomSession() + assert isinstance(session.cfg, types.SimpleNamespace) + def test_prepare_value(self): session = CustomSession()