Remove API Keys -> Authenticate with JWT

1 files changed, 36 insertions, 16 deletions

diff --git a/api/main.py b/api/main.py
index 6ede8ba..ac7b927 100644
--- a/api/main.py
+++ b/api/main.py
@@ -1,13 +1,20 @@
-from fastapi import FastAPI, Depends, HTTPException, Security
+from fastapi import FastAPI, Depends, HTTPException, Security, status
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse
+from datetime import timedelta
+from typing import Annotated
+from fastapi.security import OAuth2PasswordRequestForm
 import string
 import random
 
+from api.util.authentication import (
+    authenticate_user,
+    create_access_token,
+    get_current_user,
+)
 from api.routes.links_route import router as links_router
 from api.util.db_dependency import get_db
-from api.util.check_api_key import check_api_key
-from models import User
+from api.schemas.auth_schemas import User, Token
 
 
 metadata_tags = [
@@ -37,22 +44,35 @@ app.add_middleware(
 # Import routes
 app.include_router(links_router)
 
-# Regenerate the API key for the user
-@app.post("/regenerate")
-async def regenerate(api_key: str = Security(check_api_key), db = Depends(get_db)):
-    """Regenerate the API key for the user. Requires the current API key."""
-    user = db.query(User).filter(User.api_key == api_key['value']).first()
-    if not user:
-        raise HTTPException(status_code=401, detail="Invalid API key")
 
-    # Generate a new API key
-    new_api_key = ''.join(random.choices(string.ascii_letters + string.digits, k=20))
-    user.api_key = new_api_key
-    db.commit()
+"""
+Authentication
+"""
+
+
+@app.post("/token")
+async def login_for_access_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+    db=Depends(get_db),
+) -> Token:
+    """
+    Return an access token for the user, if the given authentication details are correct
+    """
+    user = authenticate_user(db, 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=30)
+    access_token = create_access_token(
+        data={"sub": user.username}, expires_delta=access_token_expires
+    )
+    return Token(access_token=access_token, token_type="bearer")
 
-    return {"status": "success", "new_api_key": new_api_key}
 
 # Redirect /api -> /api/docs
 @app.get("/")
 async def redirect_to_docs():
-    return RedirectResponse(url="/api/docs")
\ No newline at end of file
+    return RedirectResponse(url="/api/docs")