add unique string / token to prevent discord_id injection

app.py

@@ -4,34 +4,53 @@ import threading
 from urllib.parse import urlparse, parse_qs
 import discord
 from discord.ext import commands
+import secrets
 
-# discord bot -----------------------------------------------------------------------------------------------
+# Discord bot -----------------------------------------------------------------------------------------------
 intents = discord.Intents.all()
 bot = commands.Bot(command_prefix="!", intents=intents)
 GRADIO_APP_URL = "https://huggingface.co/spaces/lunarflu/gradio-oauth2"
 DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)
 
+# Dictionary to store user IDs and their corresponding unique strings
+user_tokens = {}
+
 @bot.event
 async def on_ready():
     print(f'Logged in as {bot.user}')
 
+def generate_unique_string(length=6):
+    return secrets.token_hex(length // 2)
+
 @bot.command()
 async def sendlink(ctx, user: discord.User):
     if ctx.author.id == 811235357663297546:
-        unique_link = f"{GRADIO_APP_URL}?user_id={user.id}"
+        unique_string = generate_unique_string()
+        user_tokens[user.id] = unique_string
+        unique_link = f"{GRADIO_APP_URL}?user_id={user.id}&token={unique_string}"
         await user.send(f"Click the link to sign in with Hugging Face: {unique_link}")
 
 def run_bot():
     bot.run(DISCORD_TOKEN)
+
 threading.Thread(target=run_bot).start()
 
-# gradio------------------------------------------------------------------------------------------------------------
+# Gradio ------------------------------------------------------------------------------------------------------------
 def hello(profile: gr.OAuthProfile | None, request: gr.Request) -> str:
     url_str = str(request.url)
     query_params = parse_qs(urlparse(url_str).query)
     user_id = query_params.get('user_id', [None])[0]
+    token = query_params.get('token', [None])[0]
+
+    if user_id is None or token is None:
+        return "❌ Invalid link. Missing user_id or token."
+
+    if int(user_id) not in user_tokens or user_tokens[int(user_id)] != token:
+        return "❌ Invalid or expired token."
+
     if profile is None:
         return f"❌ Not logged in. User ID: {user_id}"
+
     return f"✅ Successfully logged in as {profile.username}. User ID: {user_id}"
 
 with gr.Blocks() as demo: