Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import threading | |
| from urllib.parse import urlparse, parse_qs | |
| import discord | |
| from discord.ext import commands | |
| import secrets | |
| import asyncio | |
| # 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 = {} | |
| async def on_ready(): | |
| print(f'Logged in as {bot.user}') | |
| def generate_unique_string(length=6): | |
| return secrets.token_hex(length // 2) | |
| async def sendlink(ctx, user: discord.User): | |
| if ctx.author.id == 811235357663297546: | |
| 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}") | |
| # Gradio ------------------------------------------------------------------------------------------------------------ | |
| def validate_link(request: gr.Request) -> bool: | |
| referer = request.headers.get('Referer', '') | |
| url_str = str(referer) | |
| 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 False | |
| if int(user_id) not in user_tokens or user_tokens[int(user_id)] != token: | |
| return False | |
| return True | |
| def hello(profile: gr.OAuthProfile | None, request: gr.Request) -> str: | |
| if not validate_link(request): | |
| return "❌ Invalid link. Please use the link provided by the bot." | |
| referer = request.headers.get('Referer', '') | |
| url_str = str(referer) | |
| query_params = parse_qs(urlparse(url_str).query) | |
| user_id = query_params.get('user_id', [None])[0] | |
| 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: | |
| gr.Markdown("# Gradio OAuth Space") | |
| login_button = gr.LoginButton() | |
| m1 = gr.Markdown() | |
| def check_link_and_load(): | |
| if validate_link(gr.Request()): | |
| demo.load(hello, inputs=None, outputs=m1) | |
| else: | |
| m1.update("❌ Invalid link. Please use the link provided by the bot.") | |
| check_link_and_load() | |
| def check_login_status(): | |
| try: | |
| return login_button.get_session().get("oauth_info", None) | |
| except AttributeError: | |
| return None | |
| def check_login_wrapper(): | |
| session = check_login_status() | |
| if session is None: | |
| return "Not logged in." | |
| else: | |
| return f"Logged in as {session.get('username', 'Unknown')}" | |
| login_button.click(check_login_wrapper, inputs=None, outputs=m1) | |
| def run_gradio(): | |
| demo.launch() | |
| # Run the Gradio app in a separate thread | |
| gradio_thread = threading.Thread(target=run_gradio) | |
| gradio_thread.start() | |
| # Run the Discord bot on the main thread | |
| if DISCORD_TOKEN: | |
| bot.run(DISCORD_TOKEN) | |
| else: | |
| print("Discord token not found. Please set the DISCORD_TOKEN environment variable.") | |