Divent/divent/bot.py

from asyncio import new_event_loop
from datetime import datetime, timedelta
from os import environ
from typing import List, Optional

from disnake import Client, Guild
from dotenv import load_dotenv
from ics import Calendar, Event  # type: ignore
from ics.alarm.display import DisplayAlarm  # type: ignore
from quart import Quart, redirect, render_template, request, session, url_for
from requests_oauthlib import OAuth2Session  # type: ignore

load_dotenv()

DISCORD_TOKEN = environ.get("DISCORD_TOKEN")
OAUTH2_CLIENT_ID = environ.get("OAUTH2_CLIENT_ID")
OAUTH2_CLIENT_SECRET = environ.get("OAUTH2_CLIENT_SECRET")
OAUTH2_REDIRECT_URI = environ.get("OAUTH2_REDIRECT_URI", "callback")
if not DISCORD_TOKEN or not OAUTH2_CLIENT_ID or not OAUTH2_CLIENT_SECRET:
    raise Exception(
        "Missing some env variables "
        "(could be DISCORD_TOKEN, OAUTH2_CLIENT_ID or OAUTH2_CLIENT_SECRET)"
    )

API_BASE_URL = environ.get("API_BASE_URL", "https://discordapp.com/api")
AUTHORIZATION_BASE_URL = API_BASE_URL + "/oauth2/authorize"
TOKEN_URL = API_BASE_URL + "/oauth2/token"


class Discord(Client):
    async def on_ready(self):
        print(f"Logged on as {self.user}!")


client = Discord()
app = Quart(__name__)
app.config["SECRET_KEY"] = OAUTH2_CLIENT_SECRET


def get_guild_by_id(guild_id: str) -> Optional[Guild]:
    for guild in client.guilds:
        if str(guild.id) == guild_id or guild.vanity_url_code == guild_id:
            return guild
    return None


def token_updater(token: str) -> None:
    session["oauth2_token"] = token


def make_oauth_session(
    host_url: str,
    token: Optional[str] = None,
    state: Optional[str] = None,
    scope: Optional[List[str]] = None,
) -> OAuth2Session:
    return OAuth2Session(
        client_id=OAUTH2_CLIENT_ID,
        token=token,
        state=state,
        scope=scope,
        redirect_uri=host_url + OAUTH2_REDIRECT_URI,
        auto_refresh_kwargs={
            "client_id": OAUTH2_CLIENT_ID,
            "client_secret": OAUTH2_CLIENT_SECRET,
        },
        auto_refresh_url=TOKEN_URL,
        token_updater=token_updater,
    )


@app.context_processor
def days_before_failure():
    nextYear = datetime.today().year + 5 - ((datetime.today().year + 5) % 5)
    nextDate = datetime(year=nextYear, month=6, day=3)
    nextDelta = nextDate - datetime.now()
    return dict(days_before_failure=nextDelta.days)


@app.template_filter("guild_icon_ext")
def guild_icon_ext(icon_hash: str):
    return "gif" if icon_hash.startswith("a_") else "png"


@app.route("/")
async def index():
    return redirect(url_for(".login"))


@app.route("/login")
async def login():
    token = session.get("oauth2_token")
    if token is not None:
        return redirect(url_for(".guilds"))

    discord = make_oauth_session(host_url=request.host_url, scope=["guilds"])
    authorization_url, state = discord.authorization_url(AUTHORIZATION_BASE_URL)
    session["oauth2_state"] = state
    return redirect(authorization_url)


@app.route("/callback")
async def callback():
    request_values = await request.values
    if request_values.get("error"):
        return request_values["error"]
    discord = make_oauth_session(
        host_url=request.host_url, state=session.get("oauth2_state")
    )
    token = discord.fetch_token(
        TOKEN_URL,
        client_secret=OAUTH2_CLIENT_SECRET,
        authorization_response=request.url,
    )
    session["oauth2_token"] = token
    return redirect(url_for(".guilds"))


@app.route("/guilds")
async def guilds():
    token = session.get("oauth2_token")
    if token is None:
        return redirect(url_for(".login"))

    discord = make_oauth_session(host_url=request.host_url, token=token)
    guilds = discord.get(API_BASE_URL + "/users/@me/guilds").json()

    return await render_template("guilds.html.j2", guilds=guilds)


@app.route("/<guild_id>.ics")
async def ical(guild_id):
    guild = get_guild_by_id(guild_id)
    if guild is None:
        return redirect(url_for(".login"))

    calendar = Calendar()

    for scheduled_event in guild.scheduled_events:
        event = Event()
        event.name = scheduled_event.name
        event.begin = scheduled_event.scheduled_start_time
        event.end = (
            scheduled_event.scheduled_end_time
            or scheduled_event.scheduled_start_time + timedelta(hours=2)
        )
        event.uid = str(scheduled_event.id)
        event.description = scheduled_event.description
        event.url = f"https://discord.com/events/{guild_id}/{scheduled_event.id}"
        event.location = (
            scheduled_event.entity_metadata.location
            if scheduled_event.entity_metadata
            else None
        )

        alarm = DisplayAlarm()
        alarm.trigger = timedelta(hours=1)
        event.alarms.append(alarm)

        calendar.events.add(event)

    return str(calendar)


def run():
    loop = new_event_loop()
    discord_task = loop.create_task(client.start(DISCORD_TOKEN))
    discord_task.add_done_callback(lambda f: loop.stop())
    quart_task = loop.create_task(app.run_task())
    quart_task.add_done_callback(lambda f: loop.stop())
    loop.run_forever()