Update app.py

app.py

@@ -1,55 +1,119 @@
-# app.py
-from flask import Flask, request, jsonify
-from auth import (
-    authenticate_user,
-    create_access_token,
-    register_user,
-    ACCESS_TOKEN_EXPIRE_MINUTES,
-    SECRET_KEY
+import gradio as gr
+from flask import Flask, jsonify, request
+from flask_jwt_extended import (
+    JWTManager, create_access_token, 
+    jwt_required, get_jwt_identity
 )
-import jwt
-from datetime import timedelta
+from passlib.hash import pbkdf2_sha256
+import json
+import os
 
-app = Flask(__name__)
+# Initialize Flask
+flask_app = Flask(__name__)
+flask_app.config["JWT_SECRET_KEY"] = os.getenv("JWT_SECRET_KEY", "super-secret-key")
+jwt = JWTManager(flask_app)
 
-@app.route("/register", methods=["POST"])
+# Simple JSON-based user storage
+USERS_FILE = "users.json"
+
+def load_users():
+    try:
+        with open(USERS_FILE, "r") as f:
+            return json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        return {}
+
+def save_users(users):
+    with open(USERS_FILE, "w") as f:
+        json.dump(users, f)
+
+# Flask API Routes
+@flask_app.route("/register", methods=["POST"])
 def register():
     data = request.get_json()
-    try:
-        result = register_user(data["username"], data["password"])
-        return jsonify(result), 201
-    except ValueError as e:
-        return jsonify({"detail": str(e)}), 400
+    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
+    
+    if username in users:
+        return jsonify({"error": "Username already exists"}), 400
+    
+    users[username] = {"password": pbkdf2_sha256.hash(password)}
+    save_users(users)
+    
+    return jsonify({"message": "User created successfully"}), 201
 
-@app.route("/token", methods=["POST"])
+@flask_app.route("/login", methods=["POST"])
 def login():
-    username = request.form.get("username")
-    password = request.form.get("password")
-    user = authenticate_user(username, password)
-    if not user:
-        return jsonify({"detail": "Invalid username or password"}), 401
-
-    access_token = create_access_token(
-        data={"sub": username}, expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    data = request.get_json()
+    username = data.get("username")
+    password = data.get("password")
+    
+    users = load_users()
+    user = users.get(username)
+    
+    if not user or not pbkdf2_sha256.verify(password, user["password"]):
+        return jsonify({"error": "Invalid credentials"}), 401
+    
+    access_token = create_access_token(identity=username)
+    return jsonify({"access_token": access_token, "username": username})
+
+@flask_app.route("/protected", methods=["GET"])
+@jwt_required()
+def protected():
+    current_user = get_jwt_identity()
+    return jsonify(logged_in_as=current_user), 200
+
+# Gradio Interface
+def register_interface(username, password):
+    response = requests.post(
+        "http://localhost:7860/register",
+        json={"username": username, "password": password}
     )
-    return jsonify({"access_token": access_token, "token_type": "bearer"})
-
-@app.route("/users/me", methods=["GET"])
-def get_user():
-    auth_header = request.headers.get("Authorization")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        return jsonify({"detail": "Missing or invalid token"}), 401
-    token = auth_header.split(" ")[1]
-    try:
-        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
-        username = payload.get("sub")
-        if username is None:
-            raise jwt.InvalidTokenError
-        return jsonify({"username": username})
-    except jwt.ExpiredSignatureError:
-        return jsonify({"detail": "Token expired"}), 401
-    except jwt.InvalidTokenError:
-        return jsonify({"detail": "Invalid token"}), 401
+    return response.json().get("message", response.json().get("error", "Unknown error"))
+
+def login_interface(username, password):
+    response = requests.post(
+        "http://localhost:7860/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")
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Secure Authentication System")
+    
+    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")
+    
+    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")
+    
+    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
+    )
+
+# Combined app for HuggingFace Spaces
+app = gr.mount_gradio_app(flask_app, demo, path="/")
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8000)
+    app.run()