Remove Connectors mechanism and use aiohttp_socks instead of aiosocksy for Socks4/Socks5 proxies

This commit is contained in:
Alex Root Junior 2018-08-13 23:25:17 +03:00
parent 4d7555b1c3
commit d5290647c5
3 changed files with 41 additions and 96 deletions

View file

@ -40,29 +40,7 @@ def check_token(token: str) -> bool:
return True return True
class AbstractConnector(abc.ABC): async def check_result(method_name: str, content_type: str, status_code: int, body: str):
"""
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):
""" """
Checks whether `result` is a valid API response. Checks whether `result` is a valid API response.
A result is considered invalid if: A result is considered invalid if:
@ -113,67 +91,18 @@ class AbstractConnector(abc.ABC):
raise exceptions.TelegramAPIError(description) raise exceptions.TelegramAPIError(description)
raise exceptions.TelegramAPIError(f"{description} [{status_code}]") 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) req = compose_data(data, files)
kwargs.update(proxy=self.proxy, proxy_auth=self.proxy_auth)
try: try:
async with self.session.post(url, data=req, **kwargs) as response: async with session.post(url, data=req, **kwargs) as response:
return response.content_type, response.status, await response.text() return await check_result(method, response.content_type, response.status, await response.text())
except aiohttp.ClientError as e: except aiohttp.ClientError as e:
raise exceptions.NetworkError(f"aiohttp client throws an error: {e.__class__.__name__}: {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()
def guess_filename(obj): def guess_filename(obj):
""" """

View file

@ -19,7 +19,6 @@ class BaseBot:
def __init__(self, token: base.String, def __init__(self, token: base.String,
loop: Optional[Union[asyncio.BaseEventLoop, asyncio.AbstractEventLoop]] = None, loop: Optional[Union[asyncio.BaseEventLoop, asyncio.AbstractEventLoop]] = None,
connector: Optional[api.AbstractConnector] = None,
connections_limit: Optional[base.Integer] = None, connections_limit: Optional[base.Integer] = None,
proxy: Optional[base.String] = None, proxy_auth: Optional[aiohttp.BasicAuth] = None, proxy: Optional[base.String] = None, proxy_auth: Optional[aiohttp.BasicAuth] = None,
validate_token: Optional[base.Boolean] = True, validate_token: Optional[base.Boolean] = True,
@ -48,30 +47,46 @@ class BaseBot:
api.check_token(token) api.check_token(token)
self.__token = token self.__token = token
if connector and any((connections_limit, proxy, proxy_auth)): self.proxy = proxy
raise ValueError('Connector instance can\'t be passed with connection settings in one time.') self.proxy_auth = proxy_auth
elif connector:
self.connector = connector
else:
connector = api.AiohttpConnector(loop=loop, proxy=proxy, proxy_auth=proxy_auth,
connections_limit=connections_limit)
self.connector = connector
# Asyncio loop instance # Asyncio loop instance
if loop is None: if loop is None:
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
self.loop = loop self.loop = loop
# Data stored in bot instance # aiohttp main session
self._data = {} 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 self.parse_mode = parse_mode
# Data stored in bot instance
self._data = {}
async def close(self): async def close(self):
""" """
Close all client sessions Close all client sessions
""" """
await self.connector.close() await self.session.close()
async def request(self, method: base.String, async def request(self, method: base.String,
data: Optional[Dict] = None, data: Optional[Dict] = None,
@ -91,7 +106,8 @@ class BaseBot:
:rtype: Union[List, Dict] :rtype: Union[List, Dict]
:raise: :obj:`aiogram.exceptions.TelegramApiError` :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, async def download_file(self, file_path: base.String,
destination: Optional[base.InputFile] = None, destination: Optional[base.InputFile] = None,
@ -118,7 +134,7 @@ class BaseBot:
url = api.Methods.file_url(token=self.__token, path=file_path) url = api.Methods.file_url(token=self.__token, path=file_path)
dest = destination if isinstance(destination, io.IOBase) else open(destination, 'wb') 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: while True:
chunk = await response.content.read(chunk_size) chunk = await response.content.read(chunk_size)
if not chunk: if not chunk:

View file

@ -18,7 +18,7 @@ PROXY_URL = 'http://PROXY_URL' # Or 'socks5://...'
# PROXY_AUTH = aiohttp.BasicAuth(login='login', password='password') # PROXY_AUTH = aiohttp.BasicAuth(login='login', password='password')
# And add `proxy_auth=PROXY_AUTH` argument in line 25, like this: # 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) # >>> 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 my ip URL
GET_IP_URL = 'http://bot.whatismyipaddress.com/' GET_IP_URL = 'http://bot.whatismyipaddress.com/'