Add barebones implementation for GUILD_SYNC

- state_manager: add StateManager.guild_states
 - add PresenceManager in the presence module
 - fix get_user_guilds not returning ints
 - gateway: add dummy handler for op 4
 - gateway: add hazmat implementation for op 14
 - run: keep websockets logger on INFO
 - run: add more headers on app_after_request
This commit is contained in:
Luna Mendes 2018-09-01 23:53:36 -03:00
parent ee6ad56604
commit d39783e666
6 changed files with 109 additions and 10 deletions

View File

@ -13,3 +13,4 @@ class OP:
HELLO = 10 HELLO = 10
HEARTBEAT_ACK = 11 HEARTBEAT_ACK = 11
GUILD_SYNC = 12 GUILD_SYNC = 12
UNKNOWN = 14

View File

@ -1,4 +1,4 @@
from typing import List from typing import List, Dict, Any
from collections import defaultdict from collections import defaultdict
from logbook import Logger from logbook import Logger
@ -51,3 +51,17 @@ class StateManager:
states.append(state) states.append(state)
return states return states
def guild_states(self, member_ids: List[int],
guild_id: int) -> List[GatewayState]:
states = []
for member_id in member_ids:
member_states = self.fetch_states(member_id, guild_id)
# for now, just get the first state
state = next(iter(member_states))
states.append(state)
return states

View File

@ -22,7 +22,7 @@ WebsocketProperties = collections.namedtuple(
) )
WebsocketObjects = collections.namedtuple( WebsocketObjects = collections.namedtuple(
'WebsocketObjects', 'db state_manager storage loop dispatcher' 'WebsocketObjects', 'db state_manager storage loop dispatcher presence'
) )
@ -48,6 +48,7 @@ class GatewayWebsocket:
def __init__(self, ws, **kwargs): def __init__(self, ws, **kwargs):
self.ext = WebsocketObjects(*kwargs['prop']) self.ext = WebsocketObjects(*kwargs['prop'])
self.storage = self.ext.storage self.storage = self.ext.storage
self.presence = self.ext.presence
self.ws = ws self.ws = ws
self.wsp = WebsocketProperties(kwargs.get('v'), self.wsp = WebsocketProperties(kwargs.get('v'),
@ -130,11 +131,11 @@ class GatewayWebsocket:
return [ return [
{ {
**await self.storage.get_guild(row[0], user_id), **await self.storage.get_guild(guild_id, user_id),
**await self.storage.get_guild_extra(row[0], user_id, **await self.storage.get_guild_extra(guild_id, user_id,
self.state.large) self.state.large)
} }
for row in guild_ids for guild_id in guild_ids
] ]
async def guild_dispatch(self, unavailable_guilds: List[Dict[str, Any]]): async def guild_dispatch(self, unavailable_guilds: List[Dict[str, Any]]):
@ -307,6 +308,10 @@ class GatewayWebsocket:
"""Handle OP 3 Status Update.""" """Handle OP 3 Status Update."""
pass pass
async def handle_4(self, payload: Dict[str, Any]):
"""Handle OP 4 Voice Status Update."""
pass
async def handle_6(self, payload: Dict[str, Any]): async def handle_6(self, payload: Dict[str, Any]):
"""Handle OP 6 Resume.""" """Handle OP 6 Resume."""
data = payload['d'] data = payload['d']
@ -349,13 +354,54 @@ class GatewayWebsocket:
await self.dispatch('RESUMED', {}) await self.dispatch('RESUMED', {})
async def _guild_sync(self, guild_id: int):
members = await self.storage.get_member_data(guild_id)
member_ids = [int(m['user']['id']) for m in members]
log.debug(f'Syncing guild {guild_id} with {len(member_ids)} members')
presences = await self.presence.guild_presences(member_ids, guild_id)
await self.dispatch('GUILD_SYNC', {
'id': str(guild_id),
'presences': presences,
'members': members,
})
async def handle_12(self, payload: Dict[str, Any]): async def handle_12(self, payload: Dict[str, Any]):
"""Handle OP 12 Guild Sync.""" """Handle OP 12 Guild Sync."""
data = payload['d'] data = payload['d']
for _guild_id in data: gids = await self.storage.get_user_guilds(self.state.user_id)
for guild_id in data:
try:
guild_id = int(guild_id)
except (ValueError, TypeError):
continue
# check if user in guild # check if user in guild
pass if guild_id not in gids:
continue
await self._guild_sync(guild_id)
async def handle_14(self, payload: Dict[str, Any]):
# NOTE: put your HAZMAT suit on.
# OP 12 wasn't sent by the client, but OP 14 was,
# it contained a guild id, so i assume this is an
# evolution of OP 12.
# OP 14 is undocumented.
data = payload['d']
gids = await self.storage.get_user_guilds(self.state.user_id)
guild_id = int(data['guild_id'])
# make sure we are dealing with a sync to a guild
# the user is in.
if guild_id not in gids:
return
await self._guild_sync(guild_id)
async def process_message(self, payload): async def process_message(self, payload):
"""Process a single message coming in from the client.""" """Process a single message coming in from the client."""

28
litecord/presence.py Normal file
View File

@ -0,0 +1,28 @@
from typing import List, Dict, Any
class PresenceManager:
"""Presence related functions."""
def __init__(self, storage, state_manager):
self.storage = storage
self.state_manager = state_manager
async def guild_presences(self, member_ids: List[int],
guild_id: int) -> List[Dict[Any, str]]:
states = self.state_manager.guild_states(member_ids, guild_id)
presences = []
for state in states:
member = await self.storage.get_member_data_one(
guild_id, state.user_id)
presences.append({
'user': member['user'],
'roles': member['roles'],
'game': state.presence['game'],
'guild_id': guild_id,
'status': state.presence['status'],
})
return presences

View File

@ -80,7 +80,7 @@ class Storage:
WHERE user_id = $1 WHERE user_id = $1
""", user_id) """, user_id)
return guild_ids return [row['guild_id'] for row in guild_ids]
async def get_member_data_one(self, guild_id, member_id) -> Dict[str, any]: async def get_member_data_one(self, guild_id, member_id) -> Dict[str, any]:
basic = await self.db.fetchrow(""" basic = await self.db.fetchrow("""

14
run.py
View File

@ -3,6 +3,7 @@ import sys
import asyncpg import asyncpg
import logbook import logbook
import logging
import websockets import websockets
from quart import Quart, g, jsonify from quart import Quart, g, jsonify
from logbook import StreamHandler, Logger from logbook import StreamHandler, Logger
@ -16,6 +17,7 @@ from litecord.errors import LitecordError
from litecord.gateway.state_manager import StateManager from litecord.gateway.state_manager import StateManager
from litecord.storage import Storage from litecord.storage import Storage
from litecord.dispatcher import EventDispatcher from litecord.dispatcher import EventDispatcher
from litecord.presence import PresenceManager
# setup logbook # setup logbook
handler = StreamHandler(sys.stdout, level=logbook.INFO) handler = StreamHandler(sys.stdout, level=logbook.INFO)
@ -35,6 +37,9 @@ def make_app():
handler.level = logbook.DEBUG handler.level = logbook.DEBUG
app.logger.level = logbook.DEBUG app.logger.level = logbook.DEBUG
# always keep websockets on INFO
logging.getLogger('websockets').setLevel(logbook.INFO)
return app return app
@ -63,7 +68,10 @@ async def app_after_request(resp):
'X-Fingerprint, ' 'X-Fingerprint, '
'X-Context-Properties, ' 'X-Context-Properties, '
'X-Failed-Requests, ' 'X-Failed-Requests, '
'Content-Type') 'Content-Type, '
'Authorization, '
'Origin, '
'If-None-Match')
resp.headers['Access-Control-Allow-Methods'] = '*' resp.headers['Access-Control-Allow-Methods'] = '*'
return resp return resp
@ -80,6 +88,7 @@ async def app_before_serving():
app.state_manager = StateManager() app.state_manager = StateManager()
app.dispatcher = EventDispatcher(app.state_manager) app.dispatcher = EventDispatcher(app.state_manager)
app.storage = Storage(app.db) app.storage = Storage(app.db)
app.presence = PresenceManager(app.storage, app.state_manager)
# start the websocket, etc # start the websocket, etc
host, port = app.config['WS_HOST'], app.config['WS_PORT'] host, port = app.config['WS_HOST'], app.config['WS_PORT']
@ -89,7 +98,8 @@ async def app_before_serving():
# We wrap the main websocket_handler # We wrap the main websocket_handler
# so we can pass quart's app object. # so we can pass quart's app object.
await websocket_handler((app.db, app.state_manager, app.storage, await websocket_handler((app.db, app.state_manager, app.storage,
app.loop, app.dispatcher), ws, url) app.loop, app.dispatcher, app.presence),
ws, url)
ws_future = websockets.serve(_wrapper, host, port) ws_future = websockets.serve(_wrapper, host, port)