Update app.py

app.py

@@ -1,110 +1,59 @@
-import os
-from flask import Flask, jsonify, request
-from flask_jwt_extended import (
-    JWTManager,
+from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from datetime import datetime, timedelta
+from auth import (
+    authenticate_user,
     create_access_token,
-    create_refresh_token,
-    jwt_required,
-    get_jwt_identity,
-    get_jwt
+    register_user,
+    ACCESS_TOKEN_EXPIRE_MINUTES,
+    SECRET_KEY,
 )
-from passlib.hash import pbkdf2_sha256
-from dotenv import load_dotenv
-
-# Load environment variables
-load_dotenv()
-
-app = Flask(__name__)
-
-# Configure JWT
-app.config["JWT_SECRET_KEY"] = os.getenv("JWT_SECRET_KEY", "super-secret")  # Change this in production!
-app.config["JWT_ACCESS_TOKEN_EXPIRES"] = 3600  # 1 hour
-app.config["JWT_REFRESH_TOKEN_EXPIRES"] = 86400  # 24 hours
-app.config["JWT_TOKEN_LOCATION"] = ["headers", "cookies", "json", "query_string"]
-app.config["JWT_COOKIE_SECURE"] = True  # Only send over HTTPS
-app.config["JWT_COOKIE_CSRF_PROTECT"] = True  # Enable CSRF protection for cookies
-app.config["JWT_CSRF_CHECK_FORM"] = True  # Check CSRF token in form data
-
-jwt = JWTManager(app)
-
-# In-memory user storage (replace with database in production)
-users = {
-    "testuser": {
-        "username": "testuser",
-        "password": pbkdf2_sha256.hash("testpassword"),
-        "role": "user"
-    }
-}
-
-# In-memory token blocklist (replace with database in production)
-blacklist = set()
-
-@jwt.token_in_blocklist_loader
-def check_if_token_in_blacklist(jwt_header, jwt_payload):
-    jti = jwt_payload["jti"]
-    return jti in blacklist
-
-@app.route("/register", methods=["POST"])
-def register():
-    data = request.get_json()
-    username = data.get("username")
-    password = data.get("password")
-    
-    if not username or not password:
-        return jsonify({"msg": "Username and password required"}), 400
-    
-    if username in users:
-        return jsonify({"msg": "Username already exists"}), 400
-    
-    users[username] = {
-        "username": username,
-        "password": pbkdf2_sha256.hash(password),
-        "role": "user"
-    }
-    
-    return jsonify({"msg": "User created successfully"}), 201
-
-@app.route("/login", methods=["POST"])
-def login():
-    data = request.get_json()
-    username = data.get("username")
-    password = data.get("password")
-    
-    if not username or not password:
-        return jsonify({"msg": "Username and password required"}), 400
-    
-    user = users.get(username)
-    if not user or not pbkdf2_sha256.verify(password, user["password"]):
-        return jsonify({"msg": "Invalid credentials"}), 401
-    
-    access_token = create_access_token(identity=username)
-    refresh_token = create_refresh_token(identity=username)
-    
-    return jsonify({
-        "access_token": access_token,
-        "refresh_token": refresh_token,
-        "username": username
-    })
-
-@app.route("/refresh", methods=["POST"])
-@jwt_required(refresh=True)
-def refresh():
-    current_user = get_jwt_identity()
-    new_token = create_access_token(identity=current_user)
-    return jsonify({"access_token": new_token})
-
-@app.route("/protected", methods=["GET"])
-@jwt_required()
-def protected():
-    current_user = get_jwt_identity()
-    return jsonify(logged_in_as=current_user), 200
-
-@app.route("/logout", methods=["DELETE"])
-@jwt_required()
-def logout():
-    jti = get_jwt()["jti"]
-    blacklist.add(jti)
-    return jsonify({"msg": "Successfully logged out"}), 200
-
-if __name__ == "__main__":
-    app.run(ssl_context="adhoc")  # Using adhoc SSL for development
+import jwt
+from pydantic import BaseModel
+
+app = FastAPI()
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+
+class User(BaseModel):
+    username: str
+
+@app.post("/token", response_model=Token)
+async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
+    user = authenticate_user(form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    access_token = create_access_token(
+        data={"sub": form_data.username}, expires_delta=access_token_expires
+    )
+    return {"access_token": access_token, "token_type": "bearer"}
+
+@app.post("/register")
+async def register(username: str, password: str):
+    return register_user(username, password)
+
+@app.get("/users/me", response_model=User)
+async def read_users_me(token: str = Depends(oauth2_scheme)):
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
+        username = payload.get("sub")
+        if username is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid authentication credentials",
+            )
+        return {"username": username}
+    except jwt.PyJWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication credentials",
+        )