mirror of
https://github.com/nexy7574/LCC-bot.git
synced 2024-09-19 18:16:34 +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
|
import config
|
||||||
from datetime import datetime, timezone, timedelta
|
from datetime import datetime, timezone, timedelta
|
||||||
from utils import registry, console, get_or_none, JimmyBans
|
from utils import registry, console, get_or_none, JimmyBans
|
||||||
|
from web.server import app
|
||||||
|
import uvicorn
|
||||||
|
|
||||||
|
|
||||||
intents = discord.Intents.default()
|
intents = discord.Intents.default()
|
||||||
|
@ -13,14 +15,6 @@ intents += discord.Intents.members
|
||||||
intents += discord.Intents.presences
|
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 = [
|
extensions = [
|
||||||
"jishaku",
|
"jishaku",
|
||||||
"cogs.verify",
|
"cogs.verify",
|
||||||
|
@ -32,13 +26,35 @@ extensions = [
|
||||||
"cogs.starboard",
|
"cogs.starboard",
|
||||||
"cogs.uptime",
|
"cogs.uptime",
|
||||||
]
|
]
|
||||||
for ext in extensions:
|
|
||||||
try:
|
|
||||||
bot.load_extension(ext)
|
class Bot(commands.Bot):
|
||||||
except discord.ExtensionFailed as e:
|
def __init__(self):
|
||||||
console.log(f"[red]Failed to load extension {ext}: {e}")
|
super().__init__(
|
||||||
else:
|
command_prefix=self.get_prefix,
|
||||||
console.log(f"Loaded extension [green]{ext}")
|
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())
|
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__":
|
if __name__ == "__main__":
|
||||||
console.log("Starting...")
|
console.log("Starting...")
|
||||||
|
bot.started_at = discord.utils.utcnow()
|
||||||
bot.run(config.token)
|
bot.run(config.token)
|
||||||
|
|
|
@ -11,3 +11,6 @@ chromedriver==2.24.1
|
||||||
dnspython==2.2.1
|
dnspython==2.2.1
|
||||||
aiofiles==22.1.0
|
aiofiles==22.1.0
|
||||||
httpx==0.23.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),
|
"id": orm.String(min_length=7, max_length=7, unique=True),
|
||||||
"user_id": orm.BigInteger(unique=True),
|
"user_id": orm.BigInteger(unique=True),
|
||||||
"name": orm.String(min_length=2, max_length=32),
|
"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:
|
if TYPE_CHECKING:
|
||||||
entry_id: uuid.UUID
|
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