mirror of https://gitlab.com/litecord/litecord.git
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:
parent
ee6ad56604
commit
d39783e666
|
|
@ -13,3 +13,4 @@ class OP:
|
||||||
HELLO = 10
|
HELLO = 10
|
||||||
HEARTBEAT_ACK = 11
|
HEARTBEAT_ACK = 11
|
||||||
GUILD_SYNC = 12
|
GUILD_SYNC = 12
|
||||||
|
UNKNOWN = 14
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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."""
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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
14
run.py
|
|
@ -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)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue