Rapptz
9 years ago
10 changed files with 451 additions and 153 deletions
@ -0,0 +1,402 @@ |
|||||
|
# -*- coding: utf-8 -*- |
||||
|
|
||||
|
""" |
||||
|
The MIT License (MIT) |
||||
|
|
||||
|
Copyright (c) 2015-2016 Rapptz |
||||
|
|
||||
|
Permission is hereby granted, free of charge, to any person obtaining a |
||||
|
copy of this software and associated documentation files (the "Software"), |
||||
|
to deal in the Software without restriction, including without limitation |
||||
|
the rights to use, copy, modify, merge, publish, distribute, sublicense, |
||||
|
and/or sell copies of the Software, and to permit persons to whom the |
||||
|
Software is furnished to do so, subject to the following conditions: |
||||
|
|
||||
|
The above copyright notice and this permission notice shall be included in |
||||
|
all copies or substantial portions of the Software. |
||||
|
|
||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS |
||||
|
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING |
||||
|
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER |
||||
|
DEALINGS IN THE SOFTWARE. |
||||
|
""" |
||||
|
|
||||
|
import sys |
||||
|
import websockets |
||||
|
import asyncio |
||||
|
import aiohttp |
||||
|
from . import utils, endpoints, compat |
||||
|
from .enums import Status |
||||
|
from .game import Game |
||||
|
from .errors import GatewayNotFound, ConnectionClosed, InvalidArgument |
||||
|
import logging |
||||
|
import zlib, time, json |
||||
|
from collections import namedtuple |
||||
|
import threading |
||||
|
|
||||
|
log = logging.getLogger(__name__) |
||||
|
|
||||
|
__all__ = [ 'ReconnectWebSocket', 'get_gateway', 'DiscordWebSocket', |
||||
|
'KeepAliveHandler' ] |
||||
|
|
||||
|
class ReconnectWebSocket(Exception): |
||||
|
"""Signals to handle the RECONNECT opcode.""" |
||||
|
pass |
||||
|
|
||||
|
EventListener = namedtuple('EventListener', 'predicate event result future') |
||||
|
|
||||
|
class KeepAliveHandler(threading.Thread): |
||||
|
def __init__(self, *args, **kwargs): |
||||
|
ws = kwargs.pop('ws', None) |
||||
|
interval = kwargs.pop('interval', None) |
||||
|
threading.Thread.__init__(self, *args, **kwargs) |
||||
|
self.ws = ws |
||||
|
self.interval = interval |
||||
|
self.daemon = True |
||||
|
self._stop = threading.Event() |
||||
|
|
||||
|
def run(self): |
||||
|
while not self._stop.wait(self.interval): |
||||
|
data = self.get_payload() |
||||
|
msg = 'Keeping websocket alive with sequence {0[d]}'.format(data) |
||||
|
log.debug(msg) |
||||
|
coro = self.ws.send_as_json(data) |
||||
|
f = compat.run_coroutine_threadsafe(coro, loop=self.ws.loop) |
||||
|
try: |
||||
|
# block until sending is complete |
||||
|
f.result() |
||||
|
except Exception: |
||||
|
self.stop() |
||||
|
|
||||
|
def get_payload(self): |
||||
|
return { |
||||
|
'op': self.ws.HEARTBEAT, |
||||
|
'd': self.ws._connection.sequence |
||||
|
} |
||||
|
|
||||
|
def stop(self): |
||||
|
self._stop.set() |
||||
|
|
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def get_gateway(token, *, loop=None): |
||||
|
"""Returns the gateway URL for connecting to the WebSocket. |
||||
|
|
||||
|
Parameters |
||||
|
----------- |
||||
|
token : str |
||||
|
The discord authentication token. |
||||
|
loop |
||||
|
The event loop. |
||||
|
|
||||
|
Raises |
||||
|
------ |
||||
|
GatewayNotFound |
||||
|
When the gateway is not returned gracefully. |
||||
|
""" |
||||
|
headers = { |
||||
|
'authorization': token, |
||||
|
'content-type': 'application/json' |
||||
|
} |
||||
|
|
||||
|
with aiohttp.ClientSession(loop=loop) as session: |
||||
|
resp = yield from session.get(endpoints.GATEWAY, headers=headers) |
||||
|
if resp.status != 200: |
||||
|
yield from resp.release() |
||||
|
raise GatewayNotFound() |
||||
|
data = yield from resp.json() |
||||
|
return data.get('url') |
||||
|
|
||||
|
class DiscordWebSocket(websockets.client.WebSocketClientProtocol): |
||||
|
"""Implements a WebSocket for Discord's gateway v4. |
||||
|
|
||||
|
This is created through :func:`create_main_websocket`. Library |
||||
|
users should never create this manually. |
||||
|
|
||||
|
Attributes |
||||
|
----------- |
||||
|
DISPATCH |
||||
|
Receive only. Denotes an event to be sent to Discord, such as READY. |
||||
|
HEARTBEAT |
||||
|
When received tells Discord to keep the connection alive. |
||||
|
When sent asks if your connection is currently alive. |
||||
|
IDENTIFY |
||||
|
Send only. Starts a new session. |
||||
|
PRESENCE |
||||
|
Send only. Updates your presence. |
||||
|
VOICE_STATE |
||||
|
Send only. Starts a new connection to a voice server. |
||||
|
VOICE_PING |
||||
|
Send only. Checks ping time to a voice server, do not use. |
||||
|
RESUME |
||||
|
Send only. Resumes an existing connection. |
||||
|
RECONNECT |
||||
|
Receive only. Tells the client to reconnect to a new gateway. |
||||
|
REQUEST_MEMBERS |
||||
|
Send only. Asks for the full member list of a server. |
||||
|
INVALIDATE_SESSION |
||||
|
Receive only. Tells the client to invalidate the session and IDENTIFY |
||||
|
again. |
||||
|
gateway |
||||
|
The gateway we are currently connected to. |
||||
|
token |
||||
|
The authentication token for discord. |
||||
|
""" |
||||
|
|
||||
|
DISPATCH = 0 |
||||
|
HEARTBEAT = 1 |
||||
|
IDENTIFY = 2 |
||||
|
PRESENCE = 3 |
||||
|
VOICE_STATE = 4 |
||||
|
VOICE_PING = 5 |
||||
|
RESUME = 6 |
||||
|
RECONNECT = 7 |
||||
|
REQUEST_MEMBERS = 8 |
||||
|
INVALIDATE_SESSION = 9 |
||||
|
|
||||
|
def __init__(self, *args, **kwargs): |
||||
|
super().__init__(*args, max_size=None, **kwargs) |
||||
|
# an empty dispatcher to prevent crashes |
||||
|
self._dispatch = lambda *args: None |
||||
|
# generic event listeners |
||||
|
self._dispatch_listeners = [] |
||||
|
# the keep alive |
||||
|
self._keep_alive = None |
||||
|
|
||||
|
@classmethod |
||||
|
@asyncio.coroutine |
||||
|
def connect(cls, dispatch, *, token=None, connection=None, loop=None): |
||||
|
"""Creates a main websocket for Discord used for the client. |
||||
|
|
||||
|
Parameters |
||||
|
---------- |
||||
|
token : str |
||||
|
The token for Discord authentication. |
||||
|
connection |
||||
|
The ConnectionState for the client. |
||||
|
dispatch |
||||
|
The function that dispatches events. |
||||
|
loop |
||||
|
The event loop to use. |
||||
|
|
||||
|
Returns |
||||
|
------- |
||||
|
DiscordWebSocket |
||||
|
A websocket connected to Discord. |
||||
|
""" |
||||
|
|
||||
|
gateway = yield from get_gateway(token, loop=loop) |
||||
|
ws = yield from websockets.connect(gateway, loop=loop, klass=cls) |
||||
|
|
||||
|
# dynamically add attributes needed |
||||
|
ws.token = token |
||||
|
ws._connection = connection |
||||
|
ws._dispatch = dispatch |
||||
|
ws.gateway = gateway |
||||
|
|
||||
|
log.info('Created websocket connected to {}'.format(gateway)) |
||||
|
yield from ws.identify() |
||||
|
log.info('sent the identify payload to create the websocket') |
||||
|
return ws |
||||
|
|
||||
|
@classmethod |
||||
|
def from_client(cls, client): |
||||
|
"""Creates a main websocket for Discord from a :class:`Client`. |
||||
|
|
||||
|
This is for internal use only. |
||||
|
""" |
||||
|
return cls.connect(client.dispatch, token=client.token, |
||||
|
connection=client.connection, |
||||
|
loop=client.loop) |
||||
|
|
||||
|
def wait_for(self, event, predicate, result): |
||||
|
"""Waits for a DISPATCH'd event that meets the predicate. |
||||
|
|
||||
|
Parameters |
||||
|
----------- |
||||
|
event : str |
||||
|
The event name in all upper case to wait for. |
||||
|
predicate |
||||
|
A function that takes a data parameter to check for event |
||||
|
properties. The data parameter is the 'd' key in the JSON message. |
||||
|
result |
||||
|
A function that takes the same data parameter and executes to send |
||||
|
the result to the future. |
||||
|
|
||||
|
Returns |
||||
|
-------- |
||||
|
asyncio.Future |
||||
|
A future to wait for. |
||||
|
""" |
||||
|
|
||||
|
future = asyncio.Future(loop=self.loop) |
||||
|
entry = EventListener(event=event, predicate=predicate, result=result, future=future) |
||||
|
self._dispatch_listeners.append(entry) |
||||
|
return future |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def identify(self): |
||||
|
"""Sends the IDENTIFY packet.""" |
||||
|
payload = { |
||||
|
'op': self.IDENTIFY, |
||||
|
'd': { |
||||
|
'token': self.token, |
||||
|
'properties': { |
||||
|
'$os': sys.platform, |
||||
|
'$browser': 'discord.py', |
||||
|
'$device': 'discord.py', |
||||
|
'$referrer': '', |
||||
|
'$referring_domain': '' |
||||
|
}, |
||||
|
'compress': True, |
||||
|
'large_threshold': 250, |
||||
|
'v': 3 |
||||
|
} |
||||
|
} |
||||
|
yield from self.send_as_json(payload) |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def received_message(self, msg): |
||||
|
self._dispatch('socket_raw_receive', msg) |
||||
|
|
||||
|
if isinstance(msg, bytes): |
||||
|
msg = zlib.decompress(msg, 15, 10490000) # This is 10 MiB |
||||
|
msg = msg.decode('utf-8') |
||||
|
|
||||
|
msg = json.loads(msg) |
||||
|
|
||||
|
log.debug('WebSocket Event: {}'.format(msg)) |
||||
|
self._dispatch('socket_response', msg) |
||||
|
|
||||
|
op = msg.get('op') |
||||
|
data = msg.get('d') |
||||
|
|
||||
|
if 's' in msg: |
||||
|
self._connection.sequence = msg['s'] |
||||
|
|
||||
|
if op == self.RECONNECT: |
||||
|
# "reconnect" can only be handled by the Client |
||||
|
# so we terminate our connection and raise an |
||||
|
# internal exception signalling to reconnect. |
||||
|
yield from self.close() |
||||
|
raise ReconnectWebSocket() |
||||
|
|
||||
|
if op == self.INVALIDATE_SESSION: |
||||
|
self._connection.sequence = None |
||||
|
self._connection.session_id = None |
||||
|
return |
||||
|
|
||||
|
if op != self.DISPATCH: |
||||
|
log.info('Unhandled op {}'.format(op)) |
||||
|
return |
||||
|
|
||||
|
event = msg.get('t') |
||||
|
is_ready = event == 'READY' |
||||
|
|
||||
|
if is_ready: |
||||
|
self._connection.clear() |
||||
|
self._connection.sequence = msg['s'] |
||||
|
self._connection.session_id = data['session_id'] |
||||
|
|
||||
|
if is_ready or event == 'RESUMED': |
||||
|
interval = data['heartbeat_interval'] / 1000.0 |
||||
|
self._keep_alive = KeepAliveHandler(ws=self, interval=interval) |
||||
|
self._keep_alive.start() |
||||
|
|
||||
|
parser = 'parse_' + event.lower() |
||||
|
|
||||
|
try: |
||||
|
func = getattr(self._connection, parser) |
||||
|
except AttributeError: |
||||
|
log.info('Unhandled event {}'.format(event)) |
||||
|
else: |
||||
|
func(data) |
||||
|
|
||||
|
# remove the dispatched listeners |
||||
|
removed = [] |
||||
|
for index, entry in enumerate(self._dispatch_listeners): |
||||
|
if entry.event != event: |
||||
|
continue |
||||
|
|
||||
|
future = entry.future |
||||
|
if future.cancelled(): |
||||
|
removed.append(index) |
||||
|
|
||||
|
try: |
||||
|
valid = entry.predicate(data) |
||||
|
except Exception as e: |
||||
|
future.set_exception(e) |
||||
|
removed.append(index) |
||||
|
else: |
||||
|
if valid: |
||||
|
future.set_result(entry.result) |
||||
|
removed.append(index) |
||||
|
|
||||
|
for index in reversed(removed): |
||||
|
del self._dispatch_listeners[index] |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def poll_event(self): |
||||
|
"""Polls for a DISPATCH event and handles the general gateway loop. |
||||
|
|
||||
|
Raises |
||||
|
------ |
||||
|
ConnectionClosed |
||||
|
The websocket connection was terminated for unhandled reasons. |
||||
|
""" |
||||
|
try: |
||||
|
msg = yield from self.recv() |
||||
|
yield from self.received_message(msg) |
||||
|
except websockets.exceptions.ConnectionClosed as e: |
||||
|
if e.code in (4008, 4009) or e.code in range(1001, 1015): |
||||
|
raise ReconnectWebSocket() from e |
||||
|
else: |
||||
|
raise ConnectionClosed(e) from e |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def send(self, data): |
||||
|
self._dispatch('socket_raw_send', data) |
||||
|
yield from super().send(data) |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def send_as_json(self, data): |
||||
|
yield from super().send(utils.to_json(data)) |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def change_presence(self, *, game=None, idle=None): |
||||
|
if game is not None and not isinstance(game, Game): |
||||
|
raise InvalidArgument('game must be of Game or None') |
||||
|
|
||||
|
idle_since = None if idle == False else int(time.time() * 1000) |
||||
|
sent_game = game and {'name': game.name} |
||||
|
|
||||
|
payload = { |
||||
|
'op': self.PRESENCE, |
||||
|
'd': { |
||||
|
'game': sent_game, |
||||
|
'idle_since': idle_since |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
sent = utils.to_json(payload) |
||||
|
log.debug('Sending "{}" to change status'.format(sent)) |
||||
|
yield from self.send(sent) |
||||
|
|
||||
|
for server in self._connection.servers: |
||||
|
me = server.me |
||||
|
if me is None: |
||||
|
continue |
||||
|
|
||||
|
me.game = game |
||||
|
status = Status.idle if idle_since else Status.online |
||||
|
me.status = status |
||||
|
|
||||
|
@asyncio.coroutine |
||||
|
def close(self, code=1000, reason=''): |
||||
|
if self._keep_alive: |
||||
|
self._keep_alive.stop() |
||||
|
|
||||
|
yield from super().close(code, reason) |
||||
@ -1,3 +1,3 @@ |
|||||
aiohttp>=0.21.0,<0.22.0 |
aiohttp>=0.21.0,<0.22.0 |
||||
websockets==2.7 |
websockets==3.1 |
||||
PyNaCl==1.0.1 |
PyNaCl==1.0.1 |
||||
|
|||||
Loading…
Reference in new issue