From d5290647c5be324587bcba14a2902500c160881b Mon Sep 17 00:00:00 2001 From: Alex Root Junior Date: Mon, 13 Aug 2018 23:25:17 +0300 Subject: [PATCH] Remove Connectors mechanism and use aiohttp_socks instead of aiosocksy for Socks4/Socks5 proxies --- aiogram/bot/api.py | 91 ++++------------------------------- aiogram/bot/base.py | 44 +++++++++++------ examples/proxy_and_emojize.py | 2 +- 3 files changed, 41 insertions(+), 96 deletions(-) diff --git a/aiogram/bot/api.py b/aiogram/bot/api.py index 7d9ecf51..4acd4e46 100644 --- a/aiogram/bot/api.py +++ b/aiogram/bot/api.py @@ -40,29 +40,7 @@ def check_token(token: str) -> bool: return True -class AbstractConnector(abc.ABC): - """ - Abstract connector class - """ - - def __init__(self, loop: Optional[AbstractEventLoop] = None, *args, **kwargs): - if loop is None: - loop = asyncio.get_event_loop() - self.loop = loop - self._args = args - self._kwargs = kwargs - - async def make_request(self, token, method, data=None, files=None, **kwargs): - log.debug(f"Make request: '{method}' with data: {data} and files {files}") - url = Methods.api_url(token=token, method=method) - content_type, status, data = await self.request(url, data, files, **kwargs) - return await self.check_result(method, content_type, status, data) - - @abc.abstractmethod - async def request(self, url, data=None, files=None, **kwargs) -> Tuple[str, int, str]: - pass - - async def check_result(self, method_name: str, content_type: str, status_code: int, body: str): +async def check_result(method_name: str, content_type: str, status_code: int, body: str): """ Checks whether `result` is a valid API response. A result is considered invalid if: @@ -113,66 +91,17 @@ class AbstractConnector(abc.ABC): raise exceptions.TelegramAPIError(description) raise exceptions.TelegramAPIError(f"{description} [{status_code}]") - @abc.abstractmethod - async def close(self): - pass +async def make_request(session, token, method, data=None, files=None, **kwargs): + log.debug(f"Make request: '{method}' with data: {data} and files {files}") + url = Methods.api_url(token=token, method=method) -class AiohttpConnector(AbstractConnector): - def __init__(self, loop: Optional[AbstractEventLoop] = None, - proxy: Optional[str] = None, proxy_auth: Optional[aiohttp.BasicAuth] = None, - connections_limit: Optional[int] = None, *args, **kwargs): - super(AiohttpConnector, self).__init__(loop, *args, **kwargs) - - self.proxy = proxy - self.proxy_auth = proxy_auth - - # aiohttp main session - ssl_context = ssl.create_default_context(cafile=certifi.where()) - - if isinstance(proxy, str) and proxy.startswith('socks5://'): - from aiosocksy.connector import ProxyClientRequest, ProxyConnector - connector = ProxyConnector(limit=connections_limit, ssl_context=ssl_context, - loop=self.loop) - request_class = ProxyClientRequest - else: - connector = aiohttp.TCPConnector(limit=connections_limit, ssl_context=ssl_context, - loop=self.loop) - request_class = aiohttp.ClientRequest - - self.session = aiohttp.ClientSession(connector=connector, request_class=request_class, - loop=self.loop, json_serialize=json.dumps) - - async def request(self, url, data=None, files=None, **kwargs): - """ - Make request to API - - That make request with Content-Type: - application/x-www-form-urlencoded - For simple request - and multipart/form-data - for files uploading - - https://core.telegram.org/bots/api#making-requests - - :param url: requested URL - :type url: :obj:`str` - :param data: request payload - :type data: :obj:`dict` - :param files: files - :type files: :obj:`dict` - :return: result - :rtype :obj:`bool` or :obj:`dict` - """ - req = compose_data(data, files) - kwargs.update(proxy=self.proxy, proxy_auth=self.proxy_auth) - try: - async with self.session.post(url, data=req, **kwargs) as response: - return response.content_type, response.status, await response.text() - except aiohttp.ClientError as e: - raise exceptions.NetworkError(f"aiohttp client throws an error: {e.__class__.__name__}: {e}") - - async def close(self): - if self.session and not self.session.closed: - await self.session.close() + req = compose_data(data, files) + try: + async with session.post(url, data=req, **kwargs) as response: + return await check_result(method, response.content_type, response.status, await response.text()) + except aiohttp.ClientError as e: + raise exceptions.NetworkError(f"aiohttp client throws an error: {e.__class__.__name__}: {e}") def guess_filename(obj): diff --git a/aiogram/bot/base.py b/aiogram/bot/base.py index 25c1d568..137e0c2c 100644 --- a/aiogram/bot/base.py +++ b/aiogram/bot/base.py @@ -19,7 +19,6 @@ class BaseBot: def __init__(self, token: base.String, loop: Optional[Union[asyncio.BaseEventLoop, asyncio.AbstractEventLoop]] = None, - connector: Optional[api.AbstractConnector] = None, connections_limit: Optional[base.Integer] = None, proxy: Optional[base.String] = None, proxy_auth: Optional[aiohttp.BasicAuth] = None, validate_token: Optional[base.Boolean] = True, @@ -48,30 +47,46 @@ class BaseBot: api.check_token(token) self.__token = token - if connector and any((connections_limit, proxy, proxy_auth)): - raise ValueError('Connector instance can\'t be passed with connection settings in one time.') - elif connector: - self.connector = connector - else: - connector = api.AiohttpConnector(loop=loop, proxy=proxy, proxy_auth=proxy_auth, - connections_limit=connections_limit) - self.connector = connector + self.proxy = proxy + self.proxy_auth = proxy_auth # Asyncio loop instance if loop is None: loop = asyncio.get_event_loop() self.loop = loop - # Data stored in bot instance - self._data = {} + # aiohttp main session + ssl_context = ssl.create_default_context(cafile=certifi.where()) + + if isinstance(proxy, str) and (proxy.startswith('socks5://') or proxy.startswith('socks4://')): + from aiohttp_socks import SocksConnector + from aiohttp_socks.helpers import parse_socks_url + + socks_ver, host, port, username, password = parse_socks_url(proxy) + if proxy_auth and not username or password: + username = proxy_auth.login + password = proxy_auth.password + + connector = SocksConnector(socks_ver=socks_ver, host=host, port=port, + username=username, password=password, + limit=connections_limit, ssl_context=ssl_context, + loop=self.loop) + else: + connector = aiohttp.TCPConnector(limit=connections_limit, ssl_context=ssl_context, + loop=self.loop) + + self.session = aiohttp.ClientSession(connector=connector, loop=self.loop, json_serialize=json.dumps) self.parse_mode = parse_mode + # Data stored in bot instance + self._data = {} + async def close(self): """ Close all client sessions """ - await self.connector.close() + await self.session.close() async def request(self, method: base.String, data: Optional[Dict] = None, @@ -91,7 +106,8 @@ class BaseBot: :rtype: Union[List, Dict] :raise: :obj:`aiogram.exceptions.TelegramApiError` """ - return await self.connector.make_request(self.__token, method, data, files) + return await api.make_request(self.session, self.__token, method, data, files, + proxy=self.proxy, proxy_auth=self.proxy_auth) async def download_file(self, file_path: base.String, destination: Optional[base.InputFile] = None, @@ -118,7 +134,7 @@ class BaseBot: url = api.Methods.file_url(token=self.__token, path=file_path) dest = destination if isinstance(destination, io.IOBase) else open(destination, 'wb') - async with self.connector.session.get(url, timeout=timeout) as response: + async with self.session.get(url, timeout=timeout, proxy=self.proxy, proxy_auth=self.proxy_auth) as response: while True: chunk = await response.content.read(chunk_size) if not chunk: diff --git a/examples/proxy_and_emojize.py b/examples/proxy_and_emojize.py index d979243c..7e4452ee 100644 --- a/examples/proxy_and_emojize.py +++ b/examples/proxy_and_emojize.py @@ -18,7 +18,7 @@ PROXY_URL = 'http://PROXY_URL' # Or 'socks5://...' # PROXY_AUTH = aiohttp.BasicAuth(login='login', password='password') # And add `proxy_auth=PROXY_AUTH` argument in line 25, like this: # >>> bot = Bot(token=API_TOKEN, loop=loop, proxy=PROXY_URL, proxy_auth=PROXY_AUTH) -# Also you can use Socks5 proxy but you need manually install aiosocksy package. +# Also you can use Socks5 proxy but you need manually install aiohttp_socks package. # Get my ip URL GET_IP_URL = 'http://bot.whatismyipaddress.com/'