Update app.py

app.py

@@ -1,19 +1,24 @@
 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,
-    register_user,
-    ACCESS_TOKEN_EXPIRE_MINUTES,
-    SECRET_KEY,
-)
-import jwt
+from jose import jwt
 from pydantic import BaseModel
+import os
 
-app = FastAPI()
+app = FastAPI(title="JWT Auth API")
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+# 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
@@ -22,38 +27,40 @@ class Token(BaseModel):
 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.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:
+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",
-            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
-    )
+        data={"sub": form_data.username},
+        expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
     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)):
+async def read_users_me(token: str = Depends(OAuth2PasswordBearer(tokenUrl="token"))):
     try:
-        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
         username = payload.get("sub")
         if username is None:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Invalid authentication credentials",
-            )
+            raise HTTPException(status_code=400, detail="Invalid token")
         return {"username": username}
-    except jwt.PyJWTError:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid authentication credentials",
-        )
+    except jwt.JWTError:
+        raise HTTPException(status_code=400, detail="Invalid token")
+
+@app.get("/")
+async def root():
+    return {"message": "JWT Authentication API"}