From fe9ccb68a878025b971590442b085202a1fe2a1a Mon Sep 17 00:00:00 2001 From: Luna Date: Wed, 30 Jun 2021 18:28:31 -0300 Subject: [PATCH] add gateway schema module --- litecord/gateway/schemas.py | 144 ++++++++++++++++++++++++++++++++++ litecord/gateway/websocket.py | 12 +-- litecord/schemas.py | 73 ----------------- 3 files changed, 148 insertions(+), 81 deletions(-) create mode 100644 litecord/gateway/schemas.py diff --git a/litecord/gateway/schemas.py b/litecord/gateway/schemas.py new file mode 100644 index 0000000..5b823d5 --- /dev/null +++ b/litecord/gateway/schemas.py @@ -0,0 +1,144 @@ +""" + +Litecord +Copyright (C) 2018-2019 Luna Mendes + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, version 3 of the License. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with this program. If not, see . + +""" + +from typing import Dict + +from logbook import Logger + + +from litecord.gateway.errors import DecodeError +from litecord.schemas import LitecordValidator + +log = Logger(__name__) + + +def validate( + reqjson: Dict, + schema: Dict, +) -> Dict: + validator = LitecordValidator(schema) + + try: + valid = validator.validate(reqjson) + except Exception: + log.exception("Error while validating") + raise DecodeError(f"Error while validating: {reqjson}") + + if not valid: + errs = validator.errors + log.warning("Error validating doc {!r}: {!r}", reqjson, errs) + raise DecodeError(f"Error validating message : {errs!r}") + + return validator.document + + +BASE = { + "op": {"type": "number", "required": True}, + "s": {"type": "number", "required": False}, +} + +IDENTIFY_SCHEMA = { + **BASE, + **{ + "d": { + "type": "dict", + "schema": { + "token": {"type": "string", "required": True}, + "compress": {"type": "boolean", "required": False}, + "large_threshold": {"type": "number", "required": False}, + "shard": {"type": "list", "required": False}, + "presence": {"type": "dict", "required": False}, + }, + } + }, +} + + +GW_ACTIVITY = { + "name": {"type": "string", "required": True}, + "type": {"type": "activity_type", "required": True}, + "url": {"type": "string", "required": False, "nullable": True}, + "timestamps": { + "type": "dict", + "required": False, + "schema": { + "start": {"type": "number", "required": False}, + "end": {"type": "number", "required": False}, + }, + }, + "application_id": {"type": "snowflake", "required": False, "nullable": False}, + "details": {"type": "string", "required": False, "nullable": True}, + "state": {"type": "string", "required": False, "nullable": True}, + "party": { + "type": "dict", + "required": False, + "schema": { + "id": {"type": "snowflake", "required": False}, + "size": {"type": "list", "required": False}, + }, + }, + "assets": { + "type": "dict", + "required": False, + "schema": { + "large_image": {"type": "snowflake", "required": False}, + "large_text": {"type": "string", "required": False}, + "small_image": {"type": "snowflake", "required": False}, + "small_text": {"type": "string", "required": False}, + }, + }, + "secrets": { + "type": "dict", + "required": False, + "schema": { + "join": {"type": "string", "required": False}, + "spectate": {"type": "string", "required": False}, + "match": {"type": "string", "required": False}, + }, + }, + "instance": {"type": "boolean", "required": False}, + "flags": {"type": "number", "required": False}, + "emoji": { + "type": "dict", + "required": False, + "nullable": True, + "schema": { + "animated": {"type": "boolean", "required": False, "default": False}, + "id": {"coerce": int, "nullable": True, "default": None}, + "name": {"type": "string", "required": True}, + }, + }, +} + +GW_STATUS_UPDATE = { + "status": {"type": "status_external", "required": False, "default": "online"}, + "activities": { + "type": "list", + "required": False, + "schema": {"type": "dict", "schema": GW_ACTIVITY}, + }, + "afk": {"type": "boolean", "required": False}, + "since": {"type": "number", "required": False, "nullable": True}, + "game": { + "type": "dict", + "required": False, + "nullable": True, + "schema": GW_ACTIVITY, + }, +} diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 681cc56..e91e77b 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -32,7 +32,6 @@ from quart import current_app as app from litecord.auth import raw_token_check from litecord.enums import RelationshipType, ChannelType, ActivityType -from litecord.schemas import validate, GW_STATUS_UPDATE from litecord.utils import ( task_wrapper, yield_chunks, @@ -59,6 +58,7 @@ from litecord.gateway.encoding import encode_json, decode_json, encode_etf, deco from litecord.gateway.utils import WebsocketFileHandler from litecord.pubsub.guild import GuildFlags from litecord.pubsub.channel import ChannelFlags +from litecord.gateway.schemas import validate, IDENTIFY_SCHEMA, GW_STATUS_UPDATE from litecord.storage import int_ @@ -651,13 +651,9 @@ class GatewayWebsocket: async def handle_2(self, payload: Dict[str, Any]): """Handle the OP 2 Identify packet.""" - try: - data = payload["d"] - token = data["token"] - except KeyError: - raise DecodeError("Invalid identify parameters") - - # TODO proper validation of this payload + payload = validate(payload, IDENTIFY_SCHEMA) + data = payload["d"] + token = data["token"] compress = data.get("compress", False) large = data.get("large_threshold", 50) diff --git a/litecord/schemas.py b/litecord/schemas.py index cbb4781..3ba1f4c 100644 --- a/litecord/schemas.py +++ b/litecord/schemas.py @@ -418,79 +418,6 @@ MESSAGE_CREATE = { } -GW_ACTIVITY = { - "name": {"type": "string", "required": True}, - "type": {"type": "activity_type", "required": True}, - "url": {"type": "string", "required": False, "nullable": True}, - "timestamps": { - "type": "dict", - "required": False, - "schema": { - "start": {"type": "number", "required": False}, - "end": {"type": "number", "required": False}, - }, - }, - "application_id": {"type": "snowflake", "required": False, "nullable": False}, - "details": {"type": "string", "required": False, "nullable": True}, - "state": {"type": "string", "required": False, "nullable": True}, - "party": { - "type": "dict", - "required": False, - "schema": { - "id": {"type": "snowflake", "required": False}, - "size": {"type": "list", "required": False}, - }, - }, - "assets": { - "type": "dict", - "required": False, - "schema": { - "large_image": {"type": "snowflake", "required": False}, - "large_text": {"type": "string", "required": False}, - "small_image": {"type": "snowflake", "required": False}, - "small_text": {"type": "string", "required": False}, - }, - }, - "secrets": { - "type": "dict", - "required": False, - "schema": { - "join": {"type": "string", "required": False}, - "spectate": {"type": "string", "required": False}, - "match": {"type": "string", "required": False}, - }, - }, - "instance": {"type": "boolean", "required": False}, - "flags": {"type": "number", "required": False}, - "emoji": { - "type": "dict", - "required": False, - "nullable": True, - "schema": { - "animated": {"type": "boolean", "required": False, "default": False}, - "id": {"coerce": int, "nullable": True, "default": None}, - "name": {"type": "string", "required": True}, - }, - }, -} - -GW_STATUS_UPDATE = { - "status": {"type": "status_external", "required": False, "default": "online"}, - "activities": { - "type": "list", - "required": False, - "schema": {"type": "dict", "schema": GW_ACTIVITY}, - }, - "afk": {"type": "boolean", "required": False}, - "since": {"type": "number", "required": False, "nullable": True}, - "game": { - "type": "dict", - "required": False, - "nullable": True, - "schema": GW_ACTIVITY, - }, -} - INVITE = { # max_age in seconds # 0 for infinite