Update app.py

app.py

@@ -1,68 +1,55 @@
-from fastapi import FastAPI, Depends, HTTPException, status
-from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
-from datetime import datetime, timedelta
-from jose import jwt
-from pydantic import BaseModel
-import os
-
-app = FastAPI(title="JWT Auth API")
-
-# Configuration
-SECRET_KEY = os.getenv("SECRET_KEY", "your-secret-key-keep-it-safe")
-ALGORITHM = "HS256"
-ACCESS_TOKEN_EXPIRE_MINUTES = 30
-
-# Mock database
-fake_users_db = {
-    "testuser": {
-        "username": "testuser",
-        "hashed_password": "$2b$12$EixZaYVK1fsbw1ZfbX3OXePaWxn96p36WQoeG6Lruj3vjPGga31lW",  # password = "secret"
-    }
-}
-
-class Token(BaseModel):
-    access_token: str
-    token_type: str
-
-class User(BaseModel):
-    username: str
-
-def create_access_token(data: dict, expires_delta: timedelta = None):
-    to_encode = data.copy()
-    if expires_delta:
-        expire = datetime.utcnow() + expires_delta
-    else:
-        expire = datetime.utcnow() + timedelta(minutes=15)
-    to_encode.update({"exp": expire})
-    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
-    return encoded_jwt
+# 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 jwt
+from datetime import timedelta
+
+app = Flask(__name__)
+
+@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
+
+@app.route("/token", 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
 
-@app.post("/token", response_model=Token)
-async def login(form_data: OAuth2PasswordRequestForm = Depends()):
-    user = fake_users_db.get(form_data.username)
-    if not user or form_data.password != "secret":  # In real app, verify hashed password
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Incorrect username or password",
-        )
     access_token = create_access_token(
-        data={"sub": form_data.username},
-        expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-    )  # <-- Correctly closed here
-    return {"access_token": access_token, "token_type": "bearer"}
-
-
-@app.get("/users/me", response_model=User)
-async def read_users_me(token: str = Depends(OAuth2PasswordBearer(tokenUrl="token"))):
+        data={"sub": username}, expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    )
+    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=[ALGORITHM])
+        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
         username = payload.get("sub")
         if username is None:
-            raise HTTPException(status_code=400, detail="Invalid token")
-        return {"username": username}
-    except jwt.JWTError:
-        raise HTTPException(status_code=400, detail="Invalid token")
-
-@app.get("/")
-async def root():
-    return {"message": "JWT Authentication API"}
+            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
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=8000)