mirror of
https://github.com/nexy7574/LCC-bot.git
synced 2024-09-19 10:03:40 +01:00
begin http server integration
This commit is contained in:
parent
3f2d02004d
commit
b27388a44a
4 changed files with 177 additions and 15 deletions
47
main.py
47
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)
|
||||
|
|
|
@ -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
|
|
@ -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
|
||||
|
|
139
web/server.py
Normal file
139
web/server.py
Normal file
|
@ -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
|
Loading…
Reference in a new issue