From b27388a44af95b09471a405c04d200541b1d27bf Mon Sep 17 00:00:00 2001 From: EEKIM10 Date: Wed, 22 Feb 2023 01:33:30 +0000 Subject: [PATCH] begin http server integration --- main.py | 47 +++++++++++----- requirements.txt | 3 + utils/db.py | 3 + web/server.py | 139 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 177 insertions(+), 15 deletions(-) create mode 100644 web/server.py diff --git a/main.py b/main.py index af6e69e..2c24704 100644 --- a/main.py +++ b/main.py @@ -4,6 +4,8 @@ from asyncio import Lock import config from datetime import datetime, timezone, timedelta from utils import registry, console, get_or_none, JimmyBans +from web.server import app +import uvicorn intents = discord.Intents.default() @@ -13,14 +15,6 @@ intents += discord.Intents.members intents += discord.Intents.presences -bot = commands.Bot( - commands.when_mentioned_or("h!"), - debug_guilds=config.guilds, - allowed_mentions=discord.AllowedMentions.none(), - intents=intents, -) -bot.training_lock = Lock() - extensions = [ "jishaku", "cogs.verify", @@ -32,13 +26,35 @@ extensions = [ "cogs.starboard", "cogs.uptime", ] -for ext in extensions: - try: - bot.load_extension(ext) - except discord.ExtensionFailed as e: - console.log(f"[red]Failed to load extension {ext}: {e}") - else: - console.log(f"Loaded extension [green]{ext}") + + +class Bot(commands.Bot): + def __init__(self): + super().__init__( + command_prefix=self.get_prefix, + debug_guilds=config.guilds, + allowed_mentions=discord.AllowedMentions.none(), + intents=intents, + ) + self.training_lock = Lock() + self.started_at = datetime.now(tz=timezone.utc) + self.bans = JimmyBans() + for ext in extensions: + try: + bot.load_extension(ext) + except discord.ExtensionFailed as e: + console.log(f"[red]Failed to load extension {ext}: {e}") + else: + console.log(f"Loaded extension [green]{ext}") + + app.state.bot = self + config = uvicorn.Config( + app, + port=3762 + ) + + +bot = Bot() bot.loop.run_until_complete(registry.create_all()) @@ -118,4 +134,5 @@ async def check_not_banned(ctx: discord.ApplicationContext | commands.Context): if __name__ == "__main__": console.log("Starting...") + bot.started_at = discord.utils.utcnow() bot.run(config.token) diff --git a/requirements.txt b/requirements.txt index 71b670f..932f735 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,6 @@ chromedriver==2.24.1 dnspython==2.2.1 aiofiles==22.1.0 httpx==0.23.0 +fastapi==0.92.0 +uvicorn==0.20.0 +bcrypt==4.0.1 \ No newline at end of file diff --git a/utils/db.py b/utils/db.py index 5192e36..e34a7b9 100644 --- a/utils/db.py +++ b/utils/db.py @@ -74,6 +74,9 @@ class Student(orm.Model): "id": orm.String(min_length=7, max_length=7, unique=True), "user_id": orm.BigInteger(unique=True), "name": orm.String(min_length=2, max_length=32), + "access_token": orm.String(min_length=6, max_length=128, default=None, allow_null=True), + "ip_info": orm.JSON(default=None, allow_null=True), + "access_token_hash": orm.String(min_length=128, max_length=128, default=None, allow_null=True), } if TYPE_CHECKING: entry_id: uuid.UUID diff --git a/web/server.py b/web/server.py new file mode 100644 index 0000000..7b56036 --- /dev/null +++ b/web/server.py @@ -0,0 +1,139 @@ +import discord +import os +import httpx +from datetime import datetime, timedelta, timezone +from hashlib import sha512 +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, RedirectResponse +from utils.db import Student, get_or_none + +try: + from config import OAUTH_ID, OAUTH_SECRET, OAUTH_REDIRECT_URI +except ImportError: + OAUTH_ID = OAUTH_SECRET = OAUTH_REDIRECT_URI = None + +OAUTH_ENABLED = OAUTH_ID and OAUTH_SECRET and OAUTH_REDIRECT_URI + +app = FastAPI() +app.state.bot = None +app.state.states = set() +app.state.http = httpx.Client() + + +@app.middleware("http") +async def check_bot_instanced(request, call_next): + if not request.app.state.bot: + return JSONResponse( + status_code=503, + content={"message": "Not ready."} + ) + return await call_next(request) + + +@app.get("/ping") +def ping(): + bot_started = app.state.bot.started_at - datetime.now(tz=timezone.utc) + return { + "ping": "pong", + "online": app.state.bot.is_ready(), + "latency": app.state.bot.latency, + "uptime": bot_started + } + +@app.get("/auth") +async def authenticate(req: Request, code: str = None, state: str = None): + if not (code and state) or state not in app.state.states: + value = os.urandom(3).hex() + assert value not in app.state.states, "Generated a state that already exists." + app.state.states.add(value) + return RedirectResponse( + discord.utils.oauth_url( + OAUTH_ID, + redirect_uri=OAUTH_REDIRECT_URI, + scopes=('identify',) + ) + f"&state={value}", + status_code=301 + ) + else: + app.state.states.discard(state) + # First, we need to do the auth code flow + response = app.state.http.post( + "https://discord.com/api/oauth2/token", + data={ + "client_id": OAUTH_ID, + "client_secret": OAUTH_SECRET, + "grant_type": "authorization_code", + "code": code, + "redirect_uri": OAUTH_REDIRECT_URI, + } + ) + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail=response.text + ) + data = response.json() + access_token = data["access_token"] + + # Now we can generate a token + token = sha512(access_token.encode()).hexdigest() + + # Now we can get the user's info + response = app.state.http.get( + "https://discord.com/api/users/@me", + headers={ + "Authorization": "Bearer " + data["access_token"] + } + ) + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail=response.text + ) + + user = response.json() + + # Now we need to fetch the student from the database + student = await get_or_none(Student, discord_id=user["id"]) + if not student: + raise HTTPException( + status_code=404, + detail="Student not found. Please run /verify first." + ) + + # Now send a request to https://ip-api.com/json/{ip}?fields=17136 + response = app.state.http.get( + f"https://ip-api.com/json/{req.client.host}?fields=17136" + ) + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail=response.text + ) + data = response.json() + if data["status"] != "success": + raise HTTPException( + status_code=500, + detail="Failed to get IP data." + ) + + # Now we can update the student entry with this data + await student.update(ip_info=data, access_token_hash=token) + + # And set it as a cookie + response = RedirectResponse( + "/", + status_code=307, + headers={ + "Cache-Control": "max-age=86400" + } + ) + # set the cookie for at most 86400 seconds - expire after that + response.set_cookie( + "token", + token, + max_age=86400, + same_site="strict", + httponly=True, + ) + return response