From b662831193fe9b38e8377edc3bc7506b5c39b443 Mon Sep 17 00:00:00 2001
From: nex <me@nexy7574.co.uk>
Date: Thu, 23 Feb 2023 10:29:30 +0000
Subject: [PATCH] Web server

---
 main.py          | 14 +++-----------
 requirements.txt |  2 +-
 utils/client.py  | 32 ++++++++++++++++++++++++++++++++
 web/server.py    | 12 +++++++++---
 4 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/main.py b/main.py
index 147f9b1..54eb432 100644
--- a/main.py
+++ b/main.py
@@ -1,4 +1,5 @@
 import asyncio
+import sys
 
 import discord
 from discord.ext import commands
@@ -54,6 +55,7 @@ async def on_ready():
     if getattr(config, "CONNECT_MODE", None) == 1:
         console.log("Bot is now ready and exit target 1 is set, shutting down.")
         await bot.close()
+        sys.exit(0)
 
 
 @bot.slash_command()
@@ -92,13 +94,12 @@ if __name__ == "__main__":
     if getattr(config, "WEB_SERVER", True):
         from web.server import app
         import uvicorn
+        app.state.bot = bot
 
         http_config = uvicorn.Config(
             app,
             host=getattr(config, "HTTP_HOST", "127.0.0.1"),
             port=getattr(config, "HTTP_PORT", 3762),
-            lifespan="off",
-            access_log=False,
             **getattr(config, "UVICORN_CONFIG", {})
         )
         server = uvicorn.Server(http_config)
@@ -112,12 +113,3 @@ if __name__ == "__main__":
         }
 
     bot.run(config.token)
-    if hasattr(bot, "web"):
-        console.log("Cancelling web task...")
-        bot.web["task"].cancel()
-        console.log("Shutting down web server...")
-        try:
-            bot.web["task"].result()
-        except asyncio.CancelledError:
-            pass
-        console.log("Web server closed.")
diff --git a/requirements.txt b/requirements.txt
index 80bc4e4..69cb603 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-py-cord==2.3.2
+py-cord==2.4.0
 aiosmtplib==1.1.7
 orm[sqlite]==0.3.1
 httpx==0.23.0
diff --git a/utils/client.py b/utils/client.py
index 533fe6a..e229793 100644
--- a/utils/client.py
+++ b/utils/client.py
@@ -1,8 +1,15 @@
+import asyncio
+import sys
+
 import discord
 import config
 from asyncio import Lock
 from discord.ext import commands
 from datetime import datetime, timezone
+from typing import Optional, Dict, TYPE_CHECKING, Union
+if TYPE_CHECKING:
+    from uvicorn import Server, Config
+    from asyncio import Task
 
 
 __all__ = ("Bot", 'bot')
@@ -10,6 +17,9 @@ __all__ = ("Bot", 'bot')
 
 # noinspection PyAbstractClass
 class Bot(commands.Bot):
+    if TYPE_CHECKING:
+        web: Optional[Dict[str, Union[Server, Config, Task]]]
+
     def __init__(self, intents: discord.Intents, guilds: list[int], extensions: list[str]):
         from .db import JimmyBans, registry
         from .console import console
@@ -39,6 +49,28 @@ class Bot(commands.Bot):
             self.console.log("Exit target 2 reached, shutting down (not connecting to discord).")
             return
 
+    async def close(self) -> None:
+        if getattr(self, "web", None) is not None:
+            await self.http.close()
+            self.console.log("Closing web server...")
+            await self.web["server"].shutdown()
+            self.web["task"].cancel()
+            self.console.log("Web server closed.")
+            try:
+                await self.web["task"]
+            except asyncio.CancelledError:
+                pass
+            del self.web["server"]
+            del self.web["config"]
+            del self.web["task"]
+            del self.web
+        try:
+            await asyncio.wait_for(asyncio.create_task(super().close()), timeout=10)
+        except asyncio.TimeoutError:
+            self.console.log("Timed out while closing, forcing shutdown.")
+            sys.exit(1)
+        self.console.log("Finished shutting down.")
+
 
 try:
     from config import intents as _intents
diff --git a/web/server.py b/web/server.py
index 88cfbd1..674dc51 100644
--- a/web/server.py
+++ b/web/server.py
@@ -21,6 +21,12 @@ app.state.bot = None
 app.state.states = set()
 app.state.http = httpx.Client()
 
+try:
+    from utils.client import bot
+    app.state.bot = bot
+except ImportError:
+    bot = None
+
 
 @app.middleware("http")
 async def check_bot_instanced(request, call_next):
@@ -34,12 +40,12 @@ async def check_bot_instanced(request, call_next):
 
 @app.get("/ping")
 def ping():
-    bot_started = app.state.bot.started_at - datetime.now(tz=timezone.utc)
+    bot_started = datetime.now(tz=timezone.utc) - app.state.bot.started_at
     return {
         "ping": "pong", 
         "online": app.state.bot.is_ready(), 
-        "latency": app.state.bot.latency,
-        "uptime": bot_started.total_seconds()
+        "latency": max(round(app.state.bot.latency, 2), 0.01),
+        "uptime": max(round(bot_started.total_seconds(), 2), 1)
     }