Update app.py

app.py

@@ -1,20 +1,19 @@
+import os
+import json
+import bcrypt
 import gradio as gr
-from flask import Flask, jsonify, request
+from flask import Flask, request, jsonify
 from flask_jwt_extended import (
     JWTManager, create_access_token,
     jwt_required, get_jwt_identity
 )
-import bcrypt
-import json
-import os
-import requests
 
-# Initialize Flask
+# Flask setup
 flask_app = Flask(__name__)
 flask_app.config["JWT_SECRET_KEY"] = os.getenv("JWT_SECRET_KEY", "super-secret-key")
 jwt = JWTManager(flask_app)
 
-# Simple JSON-based user storage
+# File-based user storage
 USERS_FILE = "users.json"
 
 def load_users():
@@ -28,27 +27,28 @@ def save_users(users):
     with open(USERS_FILE, "w") as f:
         json.dump(users, f)
 
-# Flask Routes
+# ============================
+# Flask API routes
+# ============================
+
 @flask_app.route("/register", methods=["POST"])
 def register():
     data = request.get_json()
     username = data.get("username")
     password = data.get("password")
 
-    users = load_users()
-
     if not username or not password:
         return jsonify({"error": "Username and password required"}), 400
 
+    users = load_users()
     if username in users:
         return jsonify({"error": "Username already exists"}), 400
 
-    # Hash the password using bcrypt
-    hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
-    users[username] = {"password": hashed_password}
+    hashed = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
+    users[username] = {"password": hashed}
     save_users(users)
 
-    return jsonify({"message": "User created successfully"}), 201
+    return jsonify({"message": "User registered successfully"}), 201
 
 @flask_app.route("/login", methods=["POST"])
 def login():
@@ -59,15 +59,11 @@ def login():
     users = load_users()
     user = users.get(username)
 
-    if not user:
-        return jsonify({"error": "Invalid credentials"}), 401
-
-    stored_hashed = user["password"].encode('utf-8')
-    if not bcrypt.checkpw(password.encode('utf-8'), stored_hashed):
+    if not user or not bcrypt.checkpw(password.encode(), user["password"].encode()):
         return jsonify({"error": "Invalid credentials"}), 401
 
     access_token = create_access_token(identity=username)
-    return jsonify({"access_token": access_token, "username": username})
+    return jsonify({"access_token": access_token, "username": username}), 200
 
 @flask_app.route("/protected", methods=["GET"])
 @jwt_required()
@@ -75,52 +71,74 @@ def protected():
     current_user = get_jwt_identity()
     return jsonify(logged_in_as=current_user), 200
 
-# Gradio Interface Functions
+# ============================
+# Gradio UI
+# ============================
+
+# For local development, use local Flask server
+LOCAL_FLASK_URL = "http://localhost:5000"
+
 def register_interface(username, password):
-    response = flask_app.test_client().post(
-        "/register", json={"username": username, "password": password}
-    )
-    return response.get_json().get("message", response.get_json().get("error", "Unknown error"))
+    try:
+        response = requests.post(
+            f"{LOCAL_FLASK_URL}/register",
+            json={"username": username, "password": password}
+        )
+        return response.json().get("message", response.json().get("error", "Unknown error"))
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 def login_interface(username, password):
-    response = flask_app.test_client().post(
-        "/login", json={"username": username, "password": password}
-    )
-    if response.status_code == 200:
-        token = response.get_json().get("access_token")
-        return f"Login successful! Token: {token[:15]}... (truncated)"
-    return response.get_json().get("error", "Login failed")
+    try:
+        response = requests.post(
+            f"{LOCAL_FLASK_URL}/login",
+            json={"username": username, "password": password}
+        )
+        if response.status_code == 200:
+            token = response.json().get("access_token")
+            return f"Login successful! Token: {token[:15]}... (truncated)"
+        return response.json().get("error", "Login failed")
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔐 JWT Auth with bcrypt (Gradio + Flask)")
-
+    gr.Markdown("# 🔐 Secure Auth System (Gradio + Flask + JWT + bcrypt)")
+    
     with gr.Tab("Register"):
         reg_username = gr.Textbox(label="Username")
         reg_password = gr.Textbox(label="Password", type="password")
         reg_output = gr.Textbox(label="Output")
         reg_button = gr.Button("Register")
+        reg_button.click(
+            register_interface,
+            inputs=[reg_username, reg_password],
+            outputs=reg_output
+        )
 
     with gr.Tab("Login"):
         login_username = gr.Textbox(label="Username")
         login_password = gr.Textbox(label="Password", type="password")
         login_output = gr.Textbox(label="Output")
         login_button = gr.Button("Login")
+        login_button.click(
+            login_interface,
+            inputs=[login_username, login_password],
+            outputs=login_output
+        )
 
-    reg_button.click(
-        register_interface,
-        inputs=[reg_username, reg_password],
-        outputs=reg_output
-    )
-
-    login_button.click(
-        login_interface,
-        inputs=[login_username, login_password],
-        outputs=login_output
-    )
-
-# Mount Gradio onto Flask for Hugging Face Spaces
-app = gr.mount_gradio_app(flask_app, demo, path="/")
-
+# You need to run Flask and Gradio separately
 if __name__ == "__main__":
-    app.run()
+    # Run this in one terminal: python your_script.py --flask
+    # And in another: python your_script.py --gradio
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--flask', action='store_true')
+    parser.add_argument('--gradio', action='store_true')
+    args = parser.parse_args()
+    
+    if args.flask:
+        flask_app.run(port=5000)
+    elif args.gradio:
+        demo.launch(port=7860)
+    else:
+        print("Please specify --flask or --gradio")