Update auth.py

auth.py

@@ -25,6 +25,7 @@ class HuggingFaceAuth:
         self.db_manager = db_manager
         self.hf_api = HfApi()
         self.admin_username = os.environ.get("ADMIN_USERNAME", "Quazim0t0")
+        self.running_in_space = 'SPACE_ID' in os.environ
     
     def login_user(self, token):
         """Log in a user with their HuggingFace token.
@@ -83,7 +84,21 @@ class HuggingFaceAuth:
         if not request:
             return None
         
-        # Get token from cookies
+        # First, check if we're in a HuggingFace Space with OAuth
+        if self.running_in_space:
+            # Check for HF-User header from Space OAuth
+            username = request.headers.get("HF-User")
+            if username:
+                # Check if user exists in our database, create if not
+                user = self.db_manager.get_user_by_username(username)
+                if not user:
+                    # Create a new user 
+                    is_admin = (username == self.admin_username)
+                    user_id = self.db_manager.add_user(username, username, is_admin)
+                    user = self.db_manager.get_user_by_username(username)
+                return user
+        
+        # Fallback to token-based auth for local development
         token = request.cookies.get("hf_token")
         
         if not token:
@@ -214,59 +229,57 @@ def create_login_ui():
     """Create the login UI components.
     
     Returns:
-        tuple: (login_button, logout_button, token_input, user_info)
+        tuple: (login_button, logout_button, user_info)
     """
     with gr.Row():
         with gr.Column(scale=3):
-            token_input = gr.Textbox(
-                placeholder="Enter your HuggingFace token",
-                label="HuggingFace Token",
-                type="password",
-                visible=True,
-                info="Your token is only stored temporarily in browser session cookies and is never saved permanently"
-            )
-            login_button = gr.Button("Login")
-            logout_button = gr.Button("Logout", visible=False)
+            # If running in a HuggingFace Space, use their OAuth
+            if 'SPACE_ID' in os.environ:
+                login_button = gr.Button("Login with HuggingFace", visible=False)
+                logout_button = gr.Button("Logout", visible=False)
+            else:
+                # For local development, use token-based login
+                login_button = gr.Button("Login with HuggingFace Token")
+                logout_button = gr.Button("Logout", visible=False)
         
         with gr.Column(scale=2):
-            user_info = gr.Markdown("Not logged in")
+            user_info = gr.Markdown("Checking login status...")
     
-    return login_button, logout_button, token_input, user_info
+    return login_button, logout_button, user_info
 
-def login_handler(token, auth_manager):
+def login_handler(auth_manager):
     """Handle login button click.
     
     Args:
-        token: HuggingFace token
         auth_manager: Authentication manager instance
         
     Returns:
-        tuple: Updated UI components visibility and user info
+        tuple: JS to redirect to login and updated UI visibility
     """
-    if not token:
-        return gr.update(visible=True), gr.update(visible=False), "Please enter your HuggingFace token"
-    
-    user = auth_manager.login_user(token)
-    
-    if user:
-        # Set cookie in JavaScript with session-only flag (no persistent storage)
-        # Cookie will expire when browser is closed
-        js = f"""
-        document.cookie = "hf_token={token}; path=/; SameSite=Strict";
+    # This is only used for local development
+    # For HuggingFace Spaces, the built-in OAuth is used
+    return (
+        gr.update(visible=False),  # Hide login button
+        gr.update(visible=True),   # Show logout button
+        "Redirecting to login...",
         """
+        <script>
+        // Open a popup window for token entry
+        function promptForToken() {
+            const token = prompt("Enter your HuggingFace token:");
+            if (token) {
+                // Set the token as a cookie
+                document.cookie = "hf_token=" + token + "; path=/; SameSite=Strict";
+                // Reload the page to apply the token
+                window.location.reload();
+            }
+        }
         
-        # Return updated UI components
-        return (
-            gr.update(visible=False),  # Hide token input
-            gr.update(visible=True),   # Show logout button
-            f"Logged in as {user['username']}"  # Update user info
-        )
-    else:
-        return (
-            gr.update(visible=True),   # Keep token input visible
-            gr.update(visible=False),  # Hide logout button
-            "Login failed. Please check your token and try again."  # Update user info
-        )
+        // Call the function
+        promptForToken();
+        </script>
+        """
+    )
 
 def logout_handler():
     """Handle logout button click.
@@ -274,36 +287,40 @@ def logout_handler():
     Returns:
         tuple: Updated UI components visibility and user info
     """
-    # Clear cookie in JavaScript
-    js = """
-    document.cookie = "hf_token=; path=/; max-age=0; SameSite=Strict";
-    """
-    
-    # Return updated UI components
+    # Clear token cookie in JavaScript
     return (
-        gr.update(visible=True),   # Show token input
+        gr.update(visible=True),   # Show login button
         gr.update(visible=False),  # Hide logout button
-        "Logged out"               # Update user info
+        "Logged out",
+        """
+        <script>
+        // Clear the token cookie
+        document.cookie = "hf_token=; path=/; max-age=0; SameSite=Strict";
+        // Reload the page
+        window.location.reload();
+        </script>
+        """
     )
 
-def setup_auth_handlers(login_button, logout_button, token_input, user_info, auth_manager):
+def setup_auth_handlers(login_button, logout_button, user_info, auth_manager):
     """Set up event handlers for authentication UI components.
     
     Args:
         login_button: Login button component
         logout_button: Logout button component
-        token_input: Token input component
         user_info: User info component
         auth_manager: Authentication manager instance
     """
-    login_button.click(
-        fn=lambda token: login_handler(token, auth_manager),
-        inputs=[token_input],
-        outputs=[token_input, logout_button, user_info]
-    )
-    
-    logout_button.click(
-        fn=logout_handler,
-        inputs=[],
-        outputs=[token_input, logout_button, user_info]
-    )
+    # Only add event handlers if not running in a HuggingFace Space
+    if 'SPACE_ID' not in os.environ:
+        login_button.click(
+            fn=lambda: login_handler(auth_manager),
+            inputs=[],
+            outputs=[login_button, logout_button, user_info, gr.HTML()]
+        )
+        
+        logout_button.click(
+            fn=logout_handler,
+            inputs=[],
+            outputs=[login_button, logout_button, user_info, gr.HTML()]
+        )