From 62d00d5bd4af27a9ce0e39048e0783536a100438 Mon Sep 17 00:00:00 2001
From: Keith Stevens <fozziethebeat@gmail.com>
Date: Sat, 28 Jan 2023 16:56:34 +0900
Subject: [PATCH] Adding a minimal backend API route that decrypts the web's
 JWE and returns the email

---
 backend/oasst_backend/api/v1/auth.py | 54 ++++++++++++++++------------
 backend/oasst_backend/config.py      |  9 +++++
 backend/requirements.txt             |  3 +-
 3 files changed, 43 insertions(+), 23 deletions(-)

diff --git a/backend/oasst_backend/api/v1/auth.py b/backend/oasst_backend/api/v1/auth.py
index 9aa46150..888be410 100644
--- a/backend/oasst_backend/api/v1/auth.py
+++ b/backend/oasst_backend/api/v1/auth.py
@@ -1,39 +1,49 @@
+import json
 from typing import Union
 
-import jwt
-from fastapi import APIRouter, Depends, HTTPException, status
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from fastapi import APIRouter, Depends, Security
 from fastapi.security import APIKeyCookie
+from jose import jwe
+from oasst_backend.config import settings
 from pydantic import BaseModel
 
 router = APIRouter()
 
-oauth2_scheme = APIKeyCookie(name="next-auth.session-token")
-
-SECRET_KEY = "O/M2uIbGj+lDD2oyNa8ax4jEOJqCPJzO53UbWShmq98="
-ALGORITHM = "HS256"
-ACCESS_TOKEN_EXPIRE_MINUTES = 30
+oauth2_scheme = APIKeyCookie(name=settings.AUTH_COOKIE_NAME)
 
 
 class TokenData(BaseModel):
-    sub: Union[str, None] = None
+    """
+    A minimal re-creation of the web's token type.  To be expanded later.
+    """
+
+    email: Union[str, None] = None
 
 
-async def get_current_user(token: str = Depends(oauth2_scheme)):
-    print("get_current_user")
-    credentials_exception = HTTPException(
-        status_code=status.HTTP_401_UNAUTHORIZED,
-        detail="Could not validate credentials",
-        headers={"WWW-Authenticate": "Bearer"},
+async def get_current_user(token: str = Security(oauth2_scheme)):
+    """
+    Decrypts the user's JSON Web Token using HKDF encryption and returns the
+    TokenData.
+    """
+    # We first generate a key from the auth secret.
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=settings.AUTH_LENGTH,
+        salt=settings.AUTH_SALT,
+        info=settings.AUTH_INFO,
     )
-    print(token)
-    payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM], options={"verify_signature": False})
-    print(payload)
-    sub: str = payload.get("sub")
-    if sub is None:
-        raise credentials_exception
-    return TokenData(sub=sub)
+    key = hkdf.derive(settings.AUTH_SECRET)
+    # Next we decrypt the JWE token.
+    payload = jwe.decrypt(token, key)
+    # Finally we have the real token JSON payload and can do whatever we want.
+    content = json.loads(payload)
+    email = content["email"]
+    return TokenData(email=email)
 
 
 @router.get("/check", response_model=str)
 async def auth_check(token_data: TokenData = Depends(get_current_user)):
-    return token_data.sub
+    """Returns the user's email if it can be decrypted."""
+    return token_data.email
diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py
index 8ca2a413..a6ed1d60 100644
--- a/backend/oasst_backend/config.py
+++ b/backend/oasst_backend/config.py
@@ -63,6 +63,15 @@ class Settings(BaseSettings):
     API_V1_STR: str = "/api/v1"
     OFFICIAL_WEB_API_KEY: str = "1234"
 
+    # Encryption fields for handling the web generated JSON Web Tokens.
+    # These fields need to be shared with the web's auth settings in order to
+    # correctly decrypt the web tokens.
+    AUTH_INFO: bytes = b"NextAuth.js Generated Encryption Key"
+    AUTH_SALT: bytes = b""
+    AUTH_LENGTH: int = 32
+    AUTH_SECRET: bytes = b"O/M2uIbGj+lDD2oyNa8ax4jEOJqCPJzO53UbWShmq98="
+    AUTH_COOKIE_NAME: str = "next-auth.session-token"
+
     POSTGRES_HOST: str = "localhost"
     POSTGRES_PORT: str = "5432"
     POSTGRES_USER: str = "postgres"
diff --git a/backend/requirements.txt b/backend/requirements.txt
index dff8d14c..1f66fe09 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -1,4 +1,5 @@
 alembic==1.8.1
+cryptography==39.0.0
 fastapi==0.88.0
 fastapi-limiter==0.1.5
 fastapi-utils==0.2.1
@@ -6,8 +7,8 @@ loguru==0.6.0
 numpy==1.22.4
 psycopg2-binary==2.9.5
 pydantic==1.9.1
-pyjwt
 python-dotenv==0.21.0
+python-jose[cryptography]==3.3.0
 redis
 scipy==1.8.1
 SQLAlchemy==1.4.41