Refactor bot setup and improve error handling; update command prefix handling and enhance session security in Flask app. Add profile and balance card generation features with image processing capabilities.

2026-05-31 12:07:12 +00:00
parent 1b91cbcb2f
commit 55b16529b5
4 changed files with 356 additions and 66 deletions
@@ -20,6 +20,13 @@ import json
 load_dotenv()

 app = Flask(__name__)
+app.config.update(
+    SESSION_COOKIE_HTTPONLY=True,
+    SESSION_COOKIE_SECURE=os.getenv("FLASK_ENV") == "production",
+    SESSION_COOKIE_SAMESITE="Lax",
+    PERMANENT_SESSION_LIFETIME=datetime.timedelta(hours=1),
+)
+
 # Ensure a secret key is set for session security. In production this MUST be
 # provided via the SECRET_KEY environment variable. In development we generate
 # a temporary one (not suitable for production).
@@ -39,6 +46,19 @@ DISCORD_REDIRECT_URI = os.getenv("DISCORD_REDIRECT_URI")
 DISCORD_API_BASE_URL = "https://discord.com/api"
 DISCORD_BOT_TOKEN = os.getenv("TOKEN")

+http = requests.Session()
+http.headers.update({"User-Agent": "DiscordBotWeb/1.0"})
+
+def discord_request(method, endpoint, headers=None, **kwargs):
+    url = f"{DISCORD_API_BASE_URL}{endpoint}"
+    try:
+        response = http.request(method, url, headers=headers, timeout=REQUEST_TIMEOUT, **kwargs)
+        response.raise_for_status()
+        return response
+    except requests.RequestException as exc:
+        app.logger.warning("Discord API request failed: %s %s %s", method, url, exc)
+        raise
+
 # Default request timeout (seconds) for external HTTP calls
 REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "10"))

@@ -84,18 +104,13 @@ def callback():
    }
    headers = {"Content-Type": "application/x-www-form-urlencoded"}
    try:
-        response = requests.post(
-            f"{DISCORD_API_BASE_URL}/oauth2/token",
+        response = discord_request(
+            "POST",
+            "/oauth2/token",
            data=data,
            headers=headers,
-            timeout=REQUEST_TIMEOUT,
        )
-        response.raise_for_status()
    except requests.RequestException:
-        flash("Failed to retrieve access token from Discord (network error).")
-        return redirect(url_for("home"))
-
-    if response.status_code != 200:
        flash("Failed to retrieve access token from Discord.")
        return redirect(url_for("home"))

@@ -104,66 +119,52 @@ def callback():
        flash("Access token missing in the response.")
        return redirect(url_for("home"))

-    # Store access token in session for later use
-    session["access_token"] = access_token
+    session.permanent = True

    # Fetch user info
    try:
-        user_resp = requests.get(
-            f"{DISCORD_API_BASE_URL}/users/@me",
+        user_resp = discord_request(
+            "GET",
+            "/users/@me",
            headers={"Authorization": f"Bearer {access_token}"},
-            timeout=REQUEST_TIMEOUT,
        )
-        user_resp.raise_for_status()
        user_data = user_resp.json()
    except requests.RequestException:
        flash("Failed to fetch user info from Discord.")
        return redirect(url_for("home"))

-    # Store user information in session
    session["user"] = user_data

    # Fetch guilds the user is in
    try:
-        guilds_response = requests.get(
-            f"{DISCORD_API_BASE_URL}/users/@me/guilds",
+        guilds_response = discord_request(
+            "GET",
+            "/users/@me/guilds",
            headers={"Authorization": f"Bearer {access_token}"},
-            timeout=REQUEST_TIMEOUT,
        )
-        guilds_response.raise_for_status()
-    except requests.RequestException:
-        flash("Failed to retrieve guilds from Discord.")
-        return redirect(url_for("home"))
-
-    if guilds_response.status_code == 200:
-        guilds = guilds_response.json()  # Store guilds in a variable
-        session["guilds"] = guilds  # Store guilds in session
+        guilds = guilds_response.json()
+        session["guilds"] = guilds

        # Fetch member count and permissions for each guild
        for guild in guilds:
            guild_id = guild["id"]
            try:
-                guild_info_response = requests.get(
-                    f"{DISCORD_API_BASE_URL}/guilds/{guild_id}",
+                guild_info_response = discord_request(
+                    "GET",
+                    f"/guilds/{guild_id}",
                    headers={"Authorization": f"Bearer {access_token}"},
-                    timeout=REQUEST_TIMEOUT,
                )
-                guild_info_response.raise_for_status()
                guild_info = guild_info_response.json()
                guild["approx_member_count"] = guild_info.get("member_count", "N/A")
-                guild["permissions"] = guild_info.get(
-                    "permissions", 0
-                )  # Store permissions
+                guild["permissions"] = guild_info.get("permissions", 0)
            except requests.RequestException:
                guild["approx_member_count"] = "N/A"
-            else:
-                guild["approx_member_count"] = "N/A"  # Fallback if unable to fetch

-        # Filter guilds where user can add bots
        guilds_with_bot_permission = [guild for guild in guilds if can_add_bot(guild)]
        session["guilds_with_bot_permission"] = guilds_with_bot_permission
-    else:
+    except requests.RequestException:
        flash("Failed to retrieve guilds from Discord.")
+        return redirect(url_for("home"))

    return redirect(url_for("home"))

@@ -352,12 +353,11 @@ def guild_settings(guild_id):

    bot_headers = {"Authorization": f"Bot {DISCORD_BOT_TOKEN}"}
    try:
-        guild_info_response = requests.get(
-            f"{DISCORD_API_BASE_URL}/guilds/{guild_id}?with_counts=true",
+        guild_info_response = discord_request(
+            "GET",
+            f"/guilds/{guild_id}?with_counts=true",
            headers=bot_headers,
-            timeout=REQUEST_TIMEOUT,
        )
-        guild_info_response.raise_for_status()
        guild_info = guild_info_response.json()
    except requests.RequestException:
        flash("Failed to retrieve guild information (bot may not be in this guild or network error).")
@@ -369,12 +369,11 @@ def guild_settings(guild_id):
    owner_name = "Unknown"
    if owner_id != "N/A":
        try:
-            owner_response = requests.get(
-                f"{DISCORD_API_BASE_URL}/users/{owner_id}",
+            owner_response = discord_request(
+                "GET",
+                f"/users/{owner_id}",
                headers=bot_headers,
-                timeout=REQUEST_TIMEOUT,
            )
-            owner_response.raise_for_status()
            owner_data = owner_response.json()
            owner_name = f"{owner_data.get('username', 'Unknown')}#{owner_data.get('discriminator', '0000')}"
        except requests.RequestException:
@@ -406,4 +405,14 @@ def profile():
 if __name__ == "__main__":
    # Running via the Flask development server for local testing.
    # For production, run under a WSGI server (gunicorn, waitress, etc.).
-    app.run(debug=True, port=5000, host="0.0.0.0", use_reloader=False)
+    if os.getenv("FLASK_ENV") == "production":
+        raise RuntimeError(
+            "Do not use the Flask development server in production. "
+            "Please serve this app with a WSGI server like gunicorn or waitress."
+        )
+    app.run(
+        debug=True,
+        port=int(os.getenv("PORT", "5000")),
+        host="0.0.0.0",
+        use_reloader=False,
+    )