Add discord auth

This commit is contained in:
Riley Winkler 2025-02-11 15:00:34 -06:00
parent 899b3654d8
commit ae1699ee9f
2 changed files with 67 additions and 2 deletions

View File

@ -1,4 +1,8 @@
from pymitter import EventEmitter
class Game:
ee = EventEmitter()
def __init__(self):
self.characters = []
self.locations = {}

View File

@ -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()
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']