mirror of https://gitlab.com/litecord/litecord.git
99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
"""
|
|
|
|
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 <http://www.gnu.org/licenses/>.
|
|
|
|
"""
|
|
|
|
import hashlib
|
|
import os
|
|
|
|
from typing import Optional
|
|
from litecord.presence import BasePresence
|
|
|
|
|
|
def gen_session_id() -> str:
|
|
"""Generate a random session ID."""
|
|
return hashlib.sha1(os.urandom(128)).hexdigest()
|
|
|
|
|
|
class PayloadStore:
|
|
"""Store manager for payloads.
|
|
|
|
This will only store a maximum of MAX_STORE_SIZE,
|
|
dropping the older payloads when adding new ones.
|
|
"""
|
|
|
|
MAX_STORE_SIZE = 250
|
|
|
|
def __init__(self):
|
|
self.store = {}
|
|
|
|
def __getitem__(self, opcode: int):
|
|
return self.store[opcode]
|
|
|
|
def __setitem__(self, opcode: int, payload: dict):
|
|
if len(self.store) > 250:
|
|
# if more than 250, remove old keys until we get 250
|
|
opcodes = sorted(list(self.store.keys()))
|
|
to_remove = len(opcodes) - self.MAX_STORE_SIZE
|
|
|
|
for idx in range(to_remove):
|
|
opcode = opcodes[idx]
|
|
self.store.pop(opcode)
|
|
|
|
self.store[opcode] = payload
|
|
|
|
|
|
class GatewayState:
|
|
"""Main websocket state.
|
|
|
|
Used to store all information tied to the websocket's session.
|
|
"""
|
|
|
|
def __init__(self, **kwargs):
|
|
self.session_id: str = kwargs.get("session_id", gen_session_id())
|
|
|
|
#: last seq received by the client
|
|
self.seq: int = int(kwargs.get("seq") or 0)
|
|
|
|
#: last seq sent by gateway
|
|
self.last_seq: int = 0
|
|
|
|
#: shard information (id and total count)
|
|
shard = kwargs.get("shard") or [0, 1]
|
|
self.current_shard: int = int(shard[0])
|
|
self.shard_count: int = int(shard[1])
|
|
|
|
self.user_id: int = int(kwargs["user_id"])
|
|
self.bot: bool = bool(kwargs.get("bot") or False)
|
|
|
|
#: set by the gateway connection
|
|
# on OP STATUS_UPDATE
|
|
self.presence: Optional[BasePresence] = None
|
|
|
|
#: set by the backend once identify happens
|
|
self.ws = None
|
|
|
|
#: store of all payloads sent by the gateway (for recovery purposes)
|
|
self.store = PayloadStore()
|
|
|
|
self.compress: bool = kwargs.get("compress") or False
|
|
|
|
self.large: int = kwargs.get("large") or 50
|
|
|
|
def __repr__(self):
|
|
return f"GatewayState<seq={self.seq} shard={self.current_shard},{self.shard_count} uid={self.user_id}>"
|