diff --git a/src/controller/lobbyClient.py b/src/controller/lobbyClient.py index e8637b3..6b52a87 100644 --- a/src/controller/lobbyClient.py +++ b/src/controller/lobbyClient.py @@ -12,6 +12,8 @@ from .auth import discord_process_code from controller import ClientController +from .session_storage import generate_cookie, session_storage + if TYPE_CHECKING: from controller import LobbyController @@ -90,6 +92,30 @@ class LobbyClientController: """Send an error string to the connected client""" await self.send(Message(type=MessageType.ERROR, data_type="error", data=error)) + async def login_user(self, user: dict): + self.discord_uid = user["id"] + self.username = user["username"] + self.discord_avatar = user["avatar"] + self.on_message = self.lobby_context + cookie = generate_cookie(user) + await self.send( + Message(type=MessageType.SET, data_type="session_cookie", data=cookie) + ) + await self.send( + Message( + type=MessageType.CREATE, + data_type="object", + data=self.lobby_controller.lobby.serialize(), + ) + ) + await self.send( + Message( + type=MessageType.CONTEXT, + data_type="game_context_id", + data="lobby", + ) + ) + @vibe_check(LoginMessage) async def login_context(self, message: LoginMessage) -> None: if message.type == MessageType.DISCORD_LOGIN: @@ -97,24 +123,10 @@ class LobbyClientController: if user is None: return if "id" in user: - self.discord_uid = user["id"] - self.username = user["username"] - self.discord_avatar = user["avatar"] - self.on_message = self.lobby_context - await self.send( - Message( - type=MessageType.CREATE, - data_type="object", - data=self.lobby_controller.lobby.serialize(), - ) - ) - await self.send( - Message( - type=MessageType.CONTEXT, - data_type="game_context_id", - data="lobby", - ) - ) + await self.login_user(user) + if message.type == MessageType.COOKIE_LOGIN: + if message.data in session_storage: + await self.login_user(session_storage[message.data]) else: print("user not connected") diff --git a/src/controller/session_storage.py b/src/controller/session_storage.py new file mode 100644 index 0000000..adb8a47 --- /dev/null +++ b/src/controller/session_storage.py @@ -0,0 +1,12 @@ +import secrets +from websockets.server import WebSocketServerProtocol + +from netcode.models import Message + +session_storage: dict[str, dict] = {} + + +def generate_cookie(user: dict): + cookie = secrets.token_urlsafe(32) + session_storage[cookie] = user + return cookie diff --git a/src/netcode/models.py b/src/netcode/models.py index 0ac7b99..ce37bca 100644 --- a/src/netcode/models.py +++ b/src/netcode/models.py @@ -10,6 +10,7 @@ class MessageType(str, Enum): LOGIN = "login" READY = "ready" DISCORD_LOGIN = "discord_login" + COOKIE_LOGIN = "cookie_login" JOIN = "join" CONTEXT = "context"