mirror of
https://github.com/morpheus65535/bazarr
synced 2024-12-25 09:12:38 +00:00
585 lines
24 KiB
Python
585 lines
24 KiB
Python
import asyncio
|
|
import ssl
|
|
|
|
try:
|
|
import aiohttp
|
|
except ImportError: # pragma: no cover
|
|
aiohttp = None
|
|
import six
|
|
|
|
from . import client
|
|
from . import exceptions
|
|
from . import packet
|
|
from . import payload
|
|
|
|
|
|
class AsyncClient(client.Client):
|
|
"""An Engine.IO client for asyncio.
|
|
|
|
This class implements a fully compliant Engine.IO web client with support
|
|
for websocket and long-polling transports, compatible with the asyncio
|
|
framework on Python 3.5 or newer.
|
|
|
|
:param logger: To enable logging set to ``True`` or pass a logger object to
|
|
use. To disable logging set to ``False``. The default is
|
|
``False``.
|
|
:param json: An alternative json module to use for encoding and decoding
|
|
packets. Custom json modules must have ``dumps`` and ``loads``
|
|
functions that are compatible with the standard library
|
|
versions.
|
|
:param request_timeout: A timeout in seconds for requests. The default is
|
|
5 seconds.
|
|
:param ssl_verify: ``True`` to verify SSL certificates, or ``False`` to
|
|
skip SSL certificate verification, allowing
|
|
connections to servers with self signed certificates.
|
|
The default is ``True``.
|
|
"""
|
|
def is_asyncio_based(self):
|
|
return True
|
|
|
|
async def connect(self, url, headers={}, transports=None,
|
|
engineio_path='engine.io'):
|
|
"""Connect to an Engine.IO server.
|
|
|
|
:param url: The URL of the Engine.IO server. It can include custom
|
|
query string parameters if required by the server.
|
|
:param headers: A dictionary with custom headers to send with the
|
|
connection request.
|
|
:param transports: The list of allowed transports. Valid transports
|
|
are ``'polling'`` and ``'websocket'``. If not
|
|
given, the polling transport is connected first,
|
|
then an upgrade to websocket is attempted.
|
|
:param engineio_path: The endpoint where the Engine.IO server is
|
|
installed. The default value is appropriate for
|
|
most cases.
|
|
|
|
Note: this method is a coroutine.
|
|
|
|
Example usage::
|
|
|
|
eio = engineio.Client()
|
|
await eio.connect('http://localhost:5000')
|
|
"""
|
|
if self.state != 'disconnected':
|
|
raise ValueError('Client is not in a disconnected state')
|
|
valid_transports = ['polling', 'websocket']
|
|
if transports is not None:
|
|
if isinstance(transports, six.text_type):
|
|
transports = [transports]
|
|
transports = [transport for transport in transports
|
|
if transport in valid_transports]
|
|
if not transports:
|
|
raise ValueError('No valid transports provided')
|
|
self.transports = transports or valid_transports
|
|
self.queue = self.create_queue()
|
|
return await getattr(self, '_connect_' + self.transports[0])(
|
|
url, headers, engineio_path)
|
|
|
|
async def wait(self):
|
|
"""Wait until the connection with the server ends.
|
|
|
|
Client applications can use this function to block the main thread
|
|
during the life of the connection.
|
|
|
|
Note: this method is a coroutine.
|
|
"""
|
|
if self.read_loop_task:
|
|
await self.read_loop_task
|
|
|
|
async def send(self, data, binary=None):
|
|
"""Send a message to a client.
|
|
|
|
:param data: The data to send to the client. Data can be of type
|
|
``str``, ``bytes``, ``list`` or ``dict``. If a ``list``
|
|
or ``dict``, the data will be serialized as JSON.
|
|
:param binary: ``True`` to send packet as binary, ``False`` to send
|
|
as text. If not given, unicode (Python 2) and str
|
|
(Python 3) are sent as text, and str (Python 2) and
|
|
bytes (Python 3) are sent as binary.
|
|
|
|
Note: this method is a coroutine.
|
|
"""
|
|
await self._send_packet(packet.Packet(packet.MESSAGE, data=data,
|
|
binary=binary))
|
|
|
|
async def disconnect(self, abort=False):
|
|
"""Disconnect from the server.
|
|
|
|
:param abort: If set to ``True``, do not wait for background tasks
|
|
associated with the connection to end.
|
|
|
|
Note: this method is a coroutine.
|
|
"""
|
|
if self.state == 'connected':
|
|
await self._send_packet(packet.Packet(packet.CLOSE))
|
|
await self.queue.put(None)
|
|
self.state = 'disconnecting'
|
|
await self._trigger_event('disconnect', run_async=False)
|
|
if self.current_transport == 'websocket':
|
|
await self.ws.close()
|
|
if not abort:
|
|
await self.read_loop_task
|
|
self.state = 'disconnected'
|
|
try:
|
|
client.connected_clients.remove(self)
|
|
except ValueError: # pragma: no cover
|
|
pass
|
|
self._reset()
|
|
|
|
def start_background_task(self, target, *args, **kwargs):
|
|
"""Start a background task.
|
|
|
|
This is a utility function that applications can use to start a
|
|
background task.
|
|
|
|
:param target: the target function to execute.
|
|
:param args: arguments to pass to the function.
|
|
:param kwargs: keyword arguments to pass to the function.
|
|
|
|
This function returns an object compatible with the `Thread` class in
|
|
the Python standard library. The `start()` method on this object is
|
|
already called by this function.
|
|
|
|
Note: this method is a coroutine.
|
|
"""
|
|
return asyncio.ensure_future(target(*args, **kwargs))
|
|
|
|
async def sleep(self, seconds=0):
|
|
"""Sleep for the requested amount of time.
|
|
|
|
Note: this method is a coroutine.
|
|
"""
|
|
return await asyncio.sleep(seconds)
|
|
|
|
def create_queue(self):
|
|
"""Create a queue object."""
|
|
q = asyncio.Queue()
|
|
q.Empty = asyncio.QueueEmpty
|
|
return q
|
|
|
|
def create_event(self):
|
|
"""Create an event object."""
|
|
return asyncio.Event()
|
|
|
|
def _reset(self):
|
|
if self.http: # pragma: no cover
|
|
asyncio.ensure_future(self.http.close())
|
|
super()._reset()
|
|
|
|
async def _connect_polling(self, url, headers, engineio_path):
|
|
"""Establish a long-polling connection to the Engine.IO server."""
|
|
if aiohttp is None: # pragma: no cover
|
|
self.logger.error('aiohttp not installed -- cannot make HTTP '
|
|
'requests!')
|
|
return
|
|
self.base_url = self._get_engineio_url(url, engineio_path, 'polling')
|
|
self.logger.info('Attempting polling connection to ' + self.base_url)
|
|
r = await self._send_request(
|
|
'GET', self.base_url + self._get_url_timestamp(), headers=headers,
|
|
timeout=self.request_timeout)
|
|
if r is None:
|
|
self._reset()
|
|
raise exceptions.ConnectionError(
|
|
'Connection refused by the server')
|
|
if r.status < 200 or r.status >= 300:
|
|
raise exceptions.ConnectionError(
|
|
'Unexpected status code {} in server response'.format(
|
|
r.status))
|
|
try:
|
|
p = payload.Payload(encoded_payload=await r.read())
|
|
except ValueError:
|
|
six.raise_from(exceptions.ConnectionError(
|
|
'Unexpected response from server'), None)
|
|
open_packet = p.packets[0]
|
|
if open_packet.packet_type != packet.OPEN:
|
|
raise exceptions.ConnectionError(
|
|
'OPEN packet not returned by server')
|
|
self.logger.info(
|
|
'Polling connection accepted with ' + str(open_packet.data))
|
|
self.sid = open_packet.data['sid']
|
|
self.upgrades = open_packet.data['upgrades']
|
|
self.ping_interval = open_packet.data['pingInterval'] / 1000.0
|
|
self.ping_timeout = open_packet.data['pingTimeout'] / 1000.0
|
|
self.current_transport = 'polling'
|
|
self.base_url += '&sid=' + self.sid
|
|
|
|
self.state = 'connected'
|
|
client.connected_clients.append(self)
|
|
await self._trigger_event('connect', run_async=False)
|
|
|
|
for pkt in p.packets[1:]:
|
|
await self._receive_packet(pkt)
|
|
|
|
if 'websocket' in self.upgrades and 'websocket' in self.transports:
|
|
# attempt to upgrade to websocket
|
|
if await self._connect_websocket(url, headers, engineio_path):
|
|
# upgrade to websocket succeeded, we're done here
|
|
return
|
|
|
|
self.ping_loop_task = self.start_background_task(self._ping_loop)
|
|
self.write_loop_task = self.start_background_task(self._write_loop)
|
|
self.read_loop_task = self.start_background_task(
|
|
self._read_loop_polling)
|
|
|
|
async def _connect_websocket(self, url, headers, engineio_path):
|
|
"""Establish or upgrade to a WebSocket connection with the server."""
|
|
if aiohttp is None: # pragma: no cover
|
|
self.logger.error('aiohttp package not installed')
|
|
return False
|
|
websocket_url = self._get_engineio_url(url, engineio_path,
|
|
'websocket')
|
|
if self.sid:
|
|
self.logger.info(
|
|
'Attempting WebSocket upgrade to ' + websocket_url)
|
|
upgrade = True
|
|
websocket_url += '&sid=' + self.sid
|
|
else:
|
|
upgrade = False
|
|
self.base_url = websocket_url
|
|
self.logger.info(
|
|
'Attempting WebSocket connection to ' + websocket_url)
|
|
|
|
if self.http is None or self.http.closed: # pragma: no cover
|
|
self.http = aiohttp.ClientSession()
|
|
|
|
try:
|
|
if not self.ssl_verify:
|
|
ssl_context = ssl.create_default_context()
|
|
ssl_context.check_hostname = False
|
|
ssl_context.verify_mode = ssl.CERT_NONE
|
|
ws = await self.http.ws_connect(
|
|
websocket_url + self._get_url_timestamp(),
|
|
headers=headers, ssl=ssl_context)
|
|
else:
|
|
ws = await self.http.ws_connect(
|
|
websocket_url + self._get_url_timestamp(),
|
|
headers=headers)
|
|
except (aiohttp.client_exceptions.WSServerHandshakeError,
|
|
aiohttp.client_exceptions.ServerConnectionError):
|
|
if upgrade:
|
|
self.logger.warning(
|
|
'WebSocket upgrade failed: connection error')
|
|
return False
|
|
else:
|
|
raise exceptions.ConnectionError('Connection error')
|
|
if upgrade:
|
|
p = packet.Packet(packet.PING, data='probe').encode(
|
|
always_bytes=False)
|
|
try:
|
|
await ws.send_str(p)
|
|
except Exception as e: # pragma: no cover
|
|
self.logger.warning(
|
|
'WebSocket upgrade failed: unexpected send exception: %s',
|
|
str(e))
|
|
return False
|
|
try:
|
|
p = (await ws.receive()).data
|
|
except Exception as e: # pragma: no cover
|
|
self.logger.warning(
|
|
'WebSocket upgrade failed: unexpected recv exception: %s',
|
|
str(e))
|
|
return False
|
|
pkt = packet.Packet(encoded_packet=p)
|
|
if pkt.packet_type != packet.PONG or pkt.data != 'probe':
|
|
self.logger.warning(
|
|
'WebSocket upgrade failed: no PONG packet')
|
|
return False
|
|
p = packet.Packet(packet.UPGRADE).encode(always_bytes=False)
|
|
try:
|
|
await ws.send_str(p)
|
|
except Exception as e: # pragma: no cover
|
|
self.logger.warning(
|
|
'WebSocket upgrade failed: unexpected send exception: %s',
|
|
str(e))
|
|
return False
|
|
self.current_transport = 'websocket'
|
|
self.logger.info('WebSocket upgrade was successful')
|
|
else:
|
|
try:
|
|
p = (await ws.receive()).data
|
|
except Exception as e: # pragma: no cover
|
|
raise exceptions.ConnectionError(
|
|
'Unexpected recv exception: ' + str(e))
|
|
open_packet = packet.Packet(encoded_packet=p)
|
|
if open_packet.packet_type != packet.OPEN:
|
|
raise exceptions.ConnectionError('no OPEN packet')
|
|
self.logger.info(
|
|
'WebSocket connection accepted with ' + str(open_packet.data))
|
|
self.sid = open_packet.data['sid']
|
|
self.upgrades = open_packet.data['upgrades']
|
|
self.ping_interval = open_packet.data['pingInterval'] / 1000.0
|
|
self.ping_timeout = open_packet.data['pingTimeout'] / 1000.0
|
|
self.current_transport = 'websocket'
|
|
|
|
self.state = 'connected'
|
|
client.connected_clients.append(self)
|
|
await self._trigger_event('connect', run_async=False)
|
|
|
|
self.ws = ws
|
|
self.ping_loop_task = self.start_background_task(self._ping_loop)
|
|
self.write_loop_task = self.start_background_task(self._write_loop)
|
|
self.read_loop_task = self.start_background_task(
|
|
self._read_loop_websocket)
|
|
return True
|
|
|
|
async def _receive_packet(self, pkt):
|
|
"""Handle incoming packets from the server."""
|
|
packet_name = packet.packet_names[pkt.packet_type] \
|
|
if pkt.packet_type < len(packet.packet_names) else 'UNKNOWN'
|
|
self.logger.info(
|
|
'Received packet %s data %s', packet_name,
|
|
pkt.data if not isinstance(pkt.data, bytes) else '<binary>')
|
|
if pkt.packet_type == packet.MESSAGE:
|
|
await self._trigger_event('message', pkt.data, run_async=True)
|
|
elif pkt.packet_type == packet.PONG:
|
|
self.pong_received = True
|
|
elif pkt.packet_type == packet.CLOSE:
|
|
await self.disconnect(abort=True)
|
|
elif pkt.packet_type == packet.NOOP:
|
|
pass
|
|
else:
|
|
self.logger.error('Received unexpected packet of type %s',
|
|
pkt.packet_type)
|
|
|
|
async def _send_packet(self, pkt):
|
|
"""Queue a packet to be sent to the server."""
|
|
if self.state != 'connected':
|
|
return
|
|
await self.queue.put(pkt)
|
|
self.logger.info(
|
|
'Sending packet %s data %s',
|
|
packet.packet_names[pkt.packet_type],
|
|
pkt.data if not isinstance(pkt.data, bytes) else '<binary>')
|
|
|
|
async def _send_request(
|
|
self, method, url, headers=None, body=None,
|
|
timeout=None): # pragma: no cover
|
|
if self.http is None or self.http.closed:
|
|
self.http = aiohttp.ClientSession()
|
|
http_method = getattr(self.http, method.lower())
|
|
|
|
try:
|
|
if not self.ssl_verify:
|
|
return await http_method(
|
|
url, headers=headers, data=body,
|
|
timeout=aiohttp.ClientTimeout(total=timeout), ssl=False)
|
|
else:
|
|
return await http_method(
|
|
url, headers=headers, data=body,
|
|
timeout=aiohttp.ClientTimeout(total=timeout))
|
|
|
|
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
|
self.logger.info('HTTP %s request to %s failed with error %s.',
|
|
method, url, exc)
|
|
|
|
async def _trigger_event(self, event, *args, **kwargs):
|
|
"""Invoke an event handler."""
|
|
run_async = kwargs.pop('run_async', False)
|
|
ret = None
|
|
if event in self.handlers:
|
|
if asyncio.iscoroutinefunction(self.handlers[event]) is True:
|
|
if run_async:
|
|
return self.start_background_task(self.handlers[event],
|
|
*args)
|
|
else:
|
|
try:
|
|
ret = await self.handlers[event](*args)
|
|
except asyncio.CancelledError: # pragma: no cover
|
|
pass
|
|
except:
|
|
self.logger.exception(event + ' async handler error')
|
|
if event == 'connect':
|
|
# if connect handler raised error we reject the
|
|
# connection
|
|
return False
|
|
else:
|
|
if run_async:
|
|
async def async_handler():
|
|
return self.handlers[event](*args)
|
|
|
|
return self.start_background_task(async_handler)
|
|
else:
|
|
try:
|
|
ret = self.handlers[event](*args)
|
|
except:
|
|
self.logger.exception(event + ' handler error')
|
|
if event == 'connect':
|
|
# if connect handler raised error we reject the
|
|
# connection
|
|
return False
|
|
return ret
|
|
|
|
async def _ping_loop(self):
|
|
"""This background task sends a PING to the server at the requested
|
|
interval.
|
|
"""
|
|
self.pong_received = True
|
|
if self.ping_loop_event is None:
|
|
self.ping_loop_event = self.create_event()
|
|
else:
|
|
self.ping_loop_event.clear()
|
|
while self.state == 'connected':
|
|
if not self.pong_received:
|
|
self.logger.info(
|
|
'PONG response has not been received, aborting')
|
|
if self.ws:
|
|
await self.ws.close()
|
|
await self.queue.put(None)
|
|
break
|
|
self.pong_received = False
|
|
await self._send_packet(packet.Packet(packet.PING))
|
|
try:
|
|
await asyncio.wait_for(self.ping_loop_event.wait(),
|
|
self.ping_interval)
|
|
except (asyncio.TimeoutError,
|
|
asyncio.CancelledError): # pragma: no cover
|
|
pass
|
|
self.logger.info('Exiting ping task')
|
|
|
|
async def _read_loop_polling(self):
|
|
"""Read packets by polling the Engine.IO server."""
|
|
while self.state == 'connected':
|
|
self.logger.info(
|
|
'Sending polling GET request to ' + self.base_url)
|
|
r = await self._send_request(
|
|
'GET', self.base_url + self._get_url_timestamp(),
|
|
timeout=max(self.ping_interval, self.ping_timeout) + 5)
|
|
if r is None:
|
|
self.logger.warning(
|
|
'Connection refused by the server, aborting')
|
|
await self.queue.put(None)
|
|
break
|
|
if r.status < 200 or r.status >= 300:
|
|
self.logger.warning('Unexpected status code %s in server '
|
|
'response, aborting', r.status)
|
|
await self.queue.put(None)
|
|
break
|
|
try:
|
|
p = payload.Payload(encoded_payload=await r.read())
|
|
except ValueError:
|
|
self.logger.warning(
|
|
'Unexpected packet from server, aborting')
|
|
await self.queue.put(None)
|
|
break
|
|
for pkt in p.packets:
|
|
await self._receive_packet(pkt)
|
|
|
|
self.logger.info('Waiting for write loop task to end')
|
|
await self.write_loop_task
|
|
self.logger.info('Waiting for ping loop task to end')
|
|
if self.ping_loop_event: # pragma: no cover
|
|
self.ping_loop_event.set()
|
|
await self.ping_loop_task
|
|
if self.state == 'connected':
|
|
await self._trigger_event('disconnect', run_async=False)
|
|
try:
|
|
client.connected_clients.remove(self)
|
|
except ValueError: # pragma: no cover
|
|
pass
|
|
self._reset()
|
|
self.logger.info('Exiting read loop task')
|
|
|
|
async def _read_loop_websocket(self):
|
|
"""Read packets from the Engine.IO WebSocket connection."""
|
|
while self.state == 'connected':
|
|
p = None
|
|
try:
|
|
p = (await self.ws.receive()).data
|
|
if p is None: # pragma: no cover
|
|
raise RuntimeError('WebSocket read returned None')
|
|
except aiohttp.client_exceptions.ServerDisconnectedError:
|
|
self.logger.info(
|
|
'Read loop: WebSocket connection was closed, aborting')
|
|
await self.queue.put(None)
|
|
break
|
|
except Exception as e:
|
|
self.logger.info(
|
|
'Unexpected error "%s", aborting', str(e))
|
|
await self.queue.put(None)
|
|
break
|
|
if isinstance(p, six.text_type): # pragma: no cover
|
|
p = p.encode('utf-8')
|
|
pkt = packet.Packet(encoded_packet=p)
|
|
await self._receive_packet(pkt)
|
|
|
|
self.logger.info('Waiting for write loop task to end')
|
|
await self.write_loop_task
|
|
self.logger.info('Waiting for ping loop task to end')
|
|
if self.ping_loop_event: # pragma: no cover
|
|
self.ping_loop_event.set()
|
|
await self.ping_loop_task
|
|
if self.state == 'connected':
|
|
await self._trigger_event('disconnect', run_async=False)
|
|
try:
|
|
client.connected_clients.remove(self)
|
|
except ValueError: # pragma: no cover
|
|
pass
|
|
self._reset()
|
|
self.logger.info('Exiting read loop task')
|
|
|
|
async def _write_loop(self):
|
|
"""This background task sends packages to the server as they are
|
|
pushed to the send queue.
|
|
"""
|
|
while self.state == 'connected':
|
|
# to simplify the timeout handling, use the maximum of the
|
|
# ping interval and ping timeout as timeout, with an extra 5
|
|
# seconds grace period
|
|
timeout = max(self.ping_interval, self.ping_timeout) + 5
|
|
packets = None
|
|
try:
|
|
packets = [await asyncio.wait_for(self.queue.get(), timeout)]
|
|
except (self.queue.Empty, asyncio.TimeoutError,
|
|
asyncio.CancelledError):
|
|
self.logger.error('packet queue is empty, aborting')
|
|
break
|
|
if packets == [None]:
|
|
self.queue.task_done()
|
|
packets = []
|
|
else:
|
|
while True:
|
|
try:
|
|
packets.append(self.queue.get_nowait())
|
|
except self.queue.Empty:
|
|
break
|
|
if packets[-1] is None:
|
|
packets = packets[:-1]
|
|
self.queue.task_done()
|
|
break
|
|
if not packets:
|
|
# empty packet list returned -> connection closed
|
|
break
|
|
if self.current_transport == 'polling':
|
|
p = payload.Payload(packets=packets)
|
|
r = await self._send_request(
|
|
'POST', self.base_url, body=p.encode(),
|
|
headers={'Content-Type': 'application/octet-stream'},
|
|
timeout=self.request_timeout)
|
|
for pkt in packets:
|
|
self.queue.task_done()
|
|
if r is None:
|
|
self.logger.warning(
|
|
'Connection refused by the server, aborting')
|
|
break
|
|
if r.status < 200 or r.status >= 300:
|
|
self.logger.warning('Unexpected status code %s in server '
|
|
'response, aborting', r.status)
|
|
self._reset()
|
|
break
|
|
else:
|
|
# websocket
|
|
try:
|
|
for pkt in packets:
|
|
if pkt.binary:
|
|
await self.ws.send_bytes(pkt.encode(
|
|
always_bytes=False))
|
|
else:
|
|
await self.ws.send_str(pkt.encode(
|
|
always_bytes=False))
|
|
self.queue.task_done()
|
|
except aiohttp.client_exceptions.ServerDisconnectedError:
|
|
self.logger.info(
|
|
'Write loop: WebSocket connection was closed, '
|
|
'aborting')
|
|
break
|
|
self.logger.info('Exiting write loop task')
|