divent/bot.py

@@ -1,6 +1,7 @@
 import json
 import logging
 from datetime import datetime, timedelta
+from functools import wraps
 from os import environ, path
 from typing import Optional
 
@@ -17,7 +18,6 @@ from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware  # type: ign
 
 load_dotenv()
 
-API_BASE_URL = environ.get("API_BASE_URL", "https://discordapp.com/api")
 DISCORD_TOKEN = environ.get("DISCORD_TOKEN")
 OAUTH2_CLIENT_ID = environ.get("OAUTH2_CLIENT_ID")
 OAUTH2_CLIENT_SECRET = environ.get("OAUTH2_CLIENT_SECRET")
@@ -36,6 +36,10 @@ SENTRY_DSN = environ.get("SENTRY_DSN")
 if SENTRY_DSN:
     sentry_sdk.init(SENTRY_DSN, integrations=[QuartIntegration()])
 
+API_BASE_URL = environ.get("API_BASE_URL", "https://discordapp.com/api")
+AUTHORIZATION_BASE_URL = f"{API_BASE_URL}/oauth2/authorize"
+TOKEN_URL = f"{API_BASE_URL}/oauth2/token"
+
 
 class Discord(Client):
     async def on_ready(self):
@@ -44,6 +48,7 @@ class Discord(Client):
 
 client = Discord()
 app = Quart(__name__)
+app.config["SECRET_KEY"] = OAUTH2_CLIENT_SECRET
 app.asgi_app = ProxyHeadersMiddleware(app.asgi_app, "*")  # type: ignore
 
 
@@ -69,22 +74,22 @@ async def not_found(error: Exception):
     return await render_template("error.html.j2", error=str(error)), 404
 
 
-def token_updater(token):
+def token_updater(token: str):
     session["oauth2_token"] = token
 
 
-def make_session(token=None, state=None, scope=None):
+def make_session(token=None, state=None, scope=None) -> OAuth2Session:
     return OAuth2Session(
         client_id=OAUTH2_CLIENT_ID,
         token=token,
         state=state,
-        scope=scope,
+        scope=["identify", "guilds"],
-        redirect_uri=url_for(".callback"),
+        redirect_uri=f"{request.host_url}callback",
         auto_refresh_kwargs={
             "client_id": OAUTH2_CLIENT_ID,
             "client_secret": OAUTH2_CLIENT_SECRET,
         },
-        auto_refresh_url=f"{API_BASE_URL}/oauth2/token",
+        auto_refresh_url=TOKEN_URL,
         token_updater=token_updater,
     )
 
@@ -113,36 +118,110 @@ def days_before_failure() -> int:
     return nextDelta.days
 
 
+def cdn_avatar_url(user_id: int, hash: str) -> str:
+    ext = "gif" if hash.startswith("a_") else "png"
+    return f"https://cdn.discordapp.com/avatars/{user_id}/{hash}.{ext}"
+
+
 @app.context_processor
 def context_processor():
-    return dict(_=i18n, client=client, days_before_failure=days_before_failure())
+    return dict(
+        _=i18n,
+        client=client,
+        cdn_avatar_url=cdn_avatar_url,
+        days_before_failure=days_before_failure(),
+    )
+
+
+def login_required(fn):
+    @wraps(fn)
+    async def wrapper(*args, **kwargs):
+        discord = make_session(token=session.get("oauth2_token"))
+        if discord:
+            return await fn(*args, **kwargs)
+        return redirect(url_for(".login"))
+
+    return wrapper
 
 
 @app.route("/")
 async def index():
+    return await render_template("index.html.j2")
+
+
+@app.route("/login")
+async def login():
+    discord = make_session()
+    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 errorhandler(request_values.get("error"))
+
+    discord = make_session(state=session.get("oauth2_state"))
+    token = discord.fetch_token(
+        TOKEN_URL,
+        client_secret=OAUTH2_CLIENT_SECRET,
+        authorization_response=request.url,
+    )
+    token_updater(token)
+    return redirect(url_for(".guilds"))
+
+
+@app.route("/guilds")
+@login_required
+async def guilds():
     guild_id = request.args.get("guild")
     guild = get_guild_by_id(guild_id)
 
     if guild:
-        return redirect(url_for(".subscribe", guild_id=guild_id))
+        return redirect(
+            url_for(".subscribe", guild_id=guild.vanity_url_code or guild.id)
+        )
 
-    return await render_template("index.html.j2")
+    discord = make_session(token=session.get("oauth2_token"))
+    user = discord.get(f"{API_BASE_URL}/users/@me").json()
+    user_guilds = discord.get(f"{API_BASE_URL}/users/@me/guilds").json()
+
+    common_guilds = []
+    for bot_guild in client.guilds:
+        for user_guild in user_guilds:
+            if str(bot_guild.id) == user_guild["id"]:
+                common_guilds.append(bot_guild)
+
+    return await render_template(
+        "guilds.html.j2", user=user, common_guilds=common_guilds
+    )
 
 
 @app.route("/subscribe/<guild_id>")
+@login_required
 async def subscribe(guild_id: str):
     guild = get_guild_by_id(guild_id)
     if guild is None:
-        return redirect(url_for(".index"))
+        return redirect(url_for(".login"))
 
-    return await render_template("subscribe.html.j2", guild=guild)
+    discord = make_session(token=session.get("oauth2_token"))
+    user_guilds = discord.get(f"{API_BASE_URL}/users/@me/guilds").json()
+
+    if not any(guild_id == user_guild.id for user_guild in user_guilds):
+        return redirect(url_for(".login"))
+
+    user = discord.get(f"{API_BASE_URL}/users/@me").json()
+
+    return await render_template("subscribe.html.j2", user=user, guild=guild)
 
 
 @app.route("/<guild_id>.ics")
 async def ical(guild_id: str):
     guild = get_guild_by_id(guild_id)
     if guild is None:
-        return redirect(url_for(".index"))
+        return redirect(url_for(".login"))
 
     calendar = Calendar()
 

divent/templates/guilds.html.j2

@@ -0,0 +1,55 @@
+{% extends "base.html.j2" %}
+{% block content %}
+    <form action="{{ url_for(".guilds") }}" method="get">
+        <div id="box">
+            <div id="avatars">
+                <img src="{{ client.user.display_avatar }}"
+                     alt="{{ _('Bot Logo') }}"
+                     width="80"
+                     height="80"/>
+                <span id="dots">…</span>
+                <img src="{{ cdn_avatar_url(user.id, user.avatar) }}"
+                    alt="{{ _('User Avatar') }}"
+                    width="80"
+                    height="80"/>
+            </div>
+            <h1>
+                <a href="{{ url_for(".index") }}">{{ client.user.display_name }}</a>
+            </h1>
+            <h3>{{ _('Choose a server:') }}</h3>
+            <select name="guild" class="black_input">
+                <option>
+                    &nbsp;
+                </option>
+                {% for guild in common_guilds %}
+                    <option value="{{ guild.vanity_url_code|default(guild.id, True) }}">
+                        {{ guild.name }}
+                    </option>
+                {% endfor %}
+            </select>
+            <div class="hr-sect">{{ _("OR") }}</div>
+            <a class="button"
+               target="_blank"
+               href="https://discord.com/api/oauth2/authorize?client_id={{ client.user.id }}&permissions=8589934592&scope=bot">
+                {{ _("Add the bot on your server") }}
+            </a>
+            <ul>
+                <li>
+                    {{ _("You must have") }}
+                    <strong>{{ _("Manage Server") }}</strong>
+                    {{ _("permission on this server to perform this action") }}
+                </li>
+                <li>
+                    {{ _("After adding the bot,") }}
+                    <a href="{{ url_for(".guilds") }}">
+                        <i class="fa fa-refresh"></i>
+                        <strong>{{ _("reload the page") }}</strong>
+                    </a>
+                </li>
+            </ul>
+        </div>
+        <div id="buttons">
+            <input type="submit" class="button" value="{{ _("Let's go!") }}"/>
+        </div>
+    </form>
+{% endblock content %}

divent/templates/index.html.j2

@@ -24,41 +24,9 @@
                     {{ _('Throwing you to a new isekai world') }}
                 </li>
             </ul>
-            <hr />
-            <h3>{{ _('Choose a server:') }}</h3>
-            <select name="guild" class="black_input">
-                <option>
-                    &nbsp;
-                </option>
-                {% for guild in client.guilds %}
-                    <option value="{{ guild.vanity_url_code|default(guild.id, True) }}">
-                        {{ guild.name }}
-                    </option>
-                {% endfor %}
-            </select>
-            <div class="hr-sect">{{ _("OR") }}</div>
-            <a class="button"
-               target="_blank"
-               href="https://discord.com/api/oauth2/authorize?client_id={{ client.user.id }}&permissions=536870912&scope=bot">
-                {{ _("Add the bot on your server") }}
-            </a>
-            <ul>
-                <li>
-                    {{ _("You must have") }}
-                    <strong>{{ _("Manage Server") }}</strong>
-                    {{ _("permission on this server to perform this action") }}
-                </li>
-                <li>
-                    {{ _("After adding the bot,") }}
-                    <a href="{{ url_for(".index") }}">
-                        <i class="fa fa-refresh"></i>
-                        <strong>{{ _("reload the page") }}</strong>
-                    </a>
-                </li>
-            </ul>
         </div>
         <div id="buttons">
-            <input type="submit" class="button" value="{{ _("Let's go!") }}"/>
+            <a class="button" href="{{ url_for(".login") }}">{{ _("Let's go!") }}</a>
         </div>
     </form>
 {% endblock content %}