From bd62dfdd4b070d419b1ca42550676c351246a52b Mon Sep 17 00:00:00 2001
From: nex <me@nexy7574.co.uk>
Date: Thu, 29 Jun 2023 12:54:47 +0100
Subject: [PATCH] ENable webscoket

---
 cogs/events.py | 40 +++++++++++++++++++---------------------
 web/server.py  | 26 ++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 23 deletions(-)

diff --git a/cogs/events.py b/cogs/events.py
index b5cbe9a..f00665c 100644
--- a/cogs/events.py
+++ b/cogs/events.py
@@ -61,6 +61,7 @@ class Events(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
         self.http = httpx.AsyncClient()
+        self.bot.bridge_queue = asyncio.Queue()
         self.fetch_discord_atom_feed.start()
 
     def cog_unload(self):
@@ -366,27 +367,24 @@ class Events(commands.Cog):
             return
 
         if message.channel.name == "femboy-hole":
-            if Path("/tmp/jimmy-bridge").exists():
-                payload = {
-                    "author": str(message.user),
-                    "content": message.content,
-                    "attachments": [
-                        {
-                            "url": a.url,
-                            "filename": a.filename,
-                            "size": a.size,
-                            "width": a.width,
-                            "height": a.height,
-                            "content_type": a.content_type,
-                        }
-                        for a in message.attachments
-                    ]
-                }
-                dumped = json.dumps(payload, separators=(",", ":")).encode()
-                buffer_need = 8192 - len(dumped)
-                if buffer_need > 0:
-                    dumped += b"\0" * buffer_need
-                Path("/tmp/jimmy-bridge").write_bytes(dumped)
+            payload = {
+                "author": str(message.user),
+                "content": message.content,
+                "at": message.created_at.timestamp(),
+                "attachments": [
+                    {
+                        "url": a.url,
+                        "filename": a.filename,
+                        "size": a.size,
+                        "width": a.width,
+                        "height": a.height,
+                        "content_type": a.content_type,
+                    }
+                    for a in message.attachments
+                ]
+            }
+            # dumped = json.dumps(payload, separators=(",", ":"))
+            self.bot.bridge_queue.put_nowait(payload)
 
         if message.channel.name == "pinboard":
             if message.type == discord.MessageType.pins_add:
diff --git a/web/server.py b/web/server.py
index f2835f2..ad159e1 100644
--- a/web/server.py
+++ b/web/server.py
@@ -8,9 +8,12 @@ from pathlib import Path
 from datetime import datetime, timezone
 from hashlib import sha512
 
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Header
 from fastapi.responses import JSONResponse, RedirectResponse, HTMLResponse
 from http import HTTPStatus
+
+from starlette.websockets import WebSocket, WebSocketDisconnect
+
 from utils import Student, get_or_none, VerifyCode, console, BannedStudentID
 from utils.db import AccessTokens
 from config import guilds
@@ -283,7 +286,7 @@ async def verify(code: str):
     )
 
 
-@app.post("/bridge")
+@app.post("/bridge", include_in_schema=False)
 async def bridge(req: Request):
     body = await req.json()
     if body["secret"] != app.state.bot.http.token:
@@ -303,3 +306,22 @@ async def bridge(req: Request):
         f"**{body['sender']}**:\n>>> {body['message']}"
     )
     return {"status": "ok"}
+
+
+@app.websocket('/bridge/recv')
+def bridge_recv(ws: WebSocket, secret: str = Header(None)):
+    if secret != app.state.bot.http.token:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid secret."
+        )
+
+    await ws.accept()
+    while True:
+        data = await app.state.bot.bridge_queue.get()
+        try:
+            await ws.send_json(data)
+        except WebSocketDisconnect:
+            break
+        finally:
+            app.state.bot.bridge_queue.job_done()