From ae1699ee9fa4f6cd77b563a4a60f7f21198d78f5 Mon Sep 17 00:00:00 2001
From: Riley Winkler <rex8112@gmail.com>
Date: Tue, 11 Feb 2025 15:00:34 -0600
Subject: [PATCH] Add discord auth

---
 HeliosBackend/game.py       |  4 +++
 HeliosBackend/websockets.py | 65 +++++++++++++++++++++++++++++++++++--
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/HeliosBackend/game.py b/HeliosBackend/game.py
index 29c0fbc..e381ed4 100644
--- a/HeliosBackend/game.py
+++ b/HeliosBackend/game.py
@@ -1,4 +1,8 @@
+from pymitter import EventEmitter
+
 class Game:
+    ee = EventEmitter()
+
     def __init__(self):
         self.characters = []
         self.locations = {}
\ No newline at end of file
diff --git a/HeliosBackend/websockets.py b/HeliosBackend/websockets.py
index cf7a9d3..e0a5203 100644
--- a/HeliosBackend/websockets.py
+++ b/HeliosBackend/websockets.py
@@ -1,25 +1,86 @@
 import asyncio
+import os
+
+import aiohttp
+import json
 
 import websockets
 
+from .events import ee
+
 class WebSocketConnection:
     def __init__(self, websocket: websockets.ServerConnection):
         self.websocket = websocket
+        self.code = None
+        self.token = None
+
+    async def send(self, data):
+        await self.websocket.send(data)
+
+    async def recv(self):
+        return await self.websocket.recv()
 
 
 class WebSocketServer:
     def __init__(self, port):
         self.port = port
+        self.client = HTTPClient()
 
         self.connections: set[websockets.ServerConnection] = set()
 
     async def handler(self, websocket: websockets.ServerConnection):
         self.connections.add(websocket)
         try:
-            pass # Do something with the websocket
+            connection = WebSocketConnection(websocket)
+            if await self.authorization(connection):
+                await ee.emit_async('websocket_authorized', connection)
         finally:
             self.connections.remove(websocket)
 
+    async def authorization(self, websocket: WebSocketConnection):
+        msg = await websocket.recv()
+
+        data = json.loads(msg)
+        if data['type'] != 'authorization':
+            return False
+        if 'code' not in data:
+            return False
+        websocket.code = data['code']
+
+        token = await self.client.get_auth_token(websocket.code)
+        if token is None:
+            return False
+
+        websocket.token = token
+        await websocket.send(json.dumps({'type': 'authorization', 'status': 200, 'token': token}))
+        return True
+
     async def start(self):
         async with websockets.serve(self.handler, "", self.port):
-            await asyncio.Future()
\ No newline at end of file
+            await asyncio.Future()
+
+class HTTPClient:
+    def __init__(self):
+        self.session = aiohttp.ClientSession()
+
+    async def get_auth_token(self, code):
+        """Get the access token from the Discord API."""
+        async with self.session.post(
+            'https://discord.com/api/oauth2/token',
+            data=aiohttp.FormData(fields={
+                'client_id': os.environ['DISCORD_CLIENT_ID'],
+                'client_secret': os.environ['DISCORD_CLIENT_SECRET'],
+                'grant_type': 'authorization_code',
+                'code': code,
+            }),
+            headers={
+                'Content-Type': 'application/x-www-form-urlencoded',
+            },
+        ) as response:
+            if response.status != 200:
+                return None
+
+            data = await response.json()
+            if 'access_token' not in data:
+                return None
+            return data['access_token']
\ No newline at end of file