From 271fb9de3a25cfbe9b54a306e7da1b7c6355fb5a Mon Sep 17 00:00:00 2001
From: Khalim Conn-Kowlessar <kconnkowlessar@gmail.com>
Date: Thu, 6 Jul 2023 11:52:38 +0100
Subject: [PATCH] Setting up jwt auth

---
 backend/.env.example            |  3 ++-
 backend/app/dependencies.py     | 48 ++++++++++++++++++++++++++++++++-
 backend/app/main.py             |  5 ++--
 backend/app/portfolio/router.py |  4 +--
 backend/requirements/base.txt   | 10 +++++++
 5 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/backend/.env.example b/backend/.env.example
index 2be99ae8..2d5dcdf9 100644
--- a/backend/.env.example
+++ b/backend/.env.example
@@ -1 +1,2 @@
-API_KEY = example-api-key
\ No newline at end of file
+API_KEY = example-api-key
+ENVIRONMENT = local
\ No newline at end of file
diff --git a/backend/app/dependencies.py b/backend/app/dependencies.py
index 4002cded..b4813859 100644
--- a/backend/app/dependencies.py
+++ b/backend/app/dependencies.py
@@ -1,9 +1,10 @@
 from fastapi import Depends, HTTPException, status
-from fastapi.security import APIKeyHeader
+from fastapi.security import APIKeyHeader, OAuth2PasswordBearer
 from app.config import get_settings
 
 
 api_key_header = APIKeyHeader(name=get_settings().API_KEY_NAME, auto_error=False)
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
 
 async def validate_api_key(api_key_header: str = Depends(api_key_header)):
@@ -12,3 +13,48 @@ async def validate_api_key(api_key_header: str = Depends(api_key_header)):
             status_code=status.HTTP_403_FORBIDDEN, detail="Could not validate credentials"
         )
     return api_key_header
+
+
+from jose import jwt, JWTError
+from fastapi import HTTPException, status
+from typing import Optional
+
+SECRET_KEY = "YOUR_SECRET_KEY"
+ALGORITHM = "HS256"
+
+def get_user(user_id: str):
+    # Define here how to fetch a user from your database
+    # using the user_id. Here's a simple placeholder implementation:
+    user = None
+    if user_id == "known_id":
+        user = {"id": user_id, "name": "Known User"}
+    return user
+
+
+def validate_jwt_token(token: str = Depends(oauth2_scheme)):
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        user_id: str = payload.get("sub")
+        if user_id is None:
+            raise credentials_exception
+        user = get_user(user_id=user_id)
+        if user is None:
+            raise credentials_exception
+        return user
+    except JWTError:
+        raise credentials_exception
+
+
+async def validate_token(token: str = Depends(oauth2_scheme)):
+    if get_settings().ENV != "local":
+        token_data = validate_jwt_token(token)
+        if not token_data:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN, detail="Could not validate credentials"
+            )
+    return token
diff --git a/backend/app/main.py b/backend/app/main.py
index 8e36c9ad..0c7ab751 100644
--- a/backend/app/main.py
+++ b/backend/app/main.py
@@ -1,8 +1,9 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Depends
 from app.portfolio import router as portfolio_router
+from app.dependencies import validate_api_key, validate_token
 
 
-app = FastAPI()
+app = FastAPI(dependencies=[Depends(validate_api_key), Depends(validate_token)])
 
 
 app.include_router(portfolio_router.router)
diff --git a/backend/app/portfolio/router.py b/backend/app/portfolio/router.py
index 27f32f7a..261bc188 100644
--- a/backend/app/portfolio/router.py
+++ b/backend/app/portfolio/router.py
@@ -1,9 +1,7 @@
-from fastapi import APIRouter, Depends
-from app.dependencies import validate_api_key
+from fastapi import APIRouter
 
 router = APIRouter(
     prefix="/portfolio",
-    dependencies=[Depends(validate_api_key)],
     tags=["portfolio"],
     responses={404: {"description": "Not found"}}
 )
diff --git a/backend/requirements/base.txt b/backend/requirements/base.txt
index d25af873..f8145eab 100644
--- a/backend/requirements/base.txt
+++ b/backend/requirements/base.txt
@@ -15,3 +15,13 @@ uvicorn==0.22.0
 uvloop==0.17.0
 watchfiles==0.19.0
 websockets==11.0.3
+pyjwt==2.7.0
+cffi==1.15.1
+cryptography==41.0.1
+ecdsa==0.18.0
+pyasn1==0.5.0
+pycparser==2.21
+python-jose==3.3.0
+rsa==4.9
+six==1.16.0
+