From f6aca6b30f69440d1bc7f1f0bb9454463466d9f3 Mon Sep 17 00:00:00 2001
From: Alexis VIALARET <alexis.vialaret@artefact.com>
Date: Tue, 16 Jan 2024 14:10:40 +0100
Subject: [PATCH] upd: security improvements

---
 backend/authentication.py  | 48 --------------------------------------
 backend/main.py            | 10 +++++---
 backend/user_management.py | 27 +++++++++++++++------
 requirements.in            |  1 +
 4 files changed, 28 insertions(+), 58 deletions(-)
 delete mode 100644 backend/authentication.py

diff --git a/backend/authentication.py b/backend/authentication.py
deleted file mode 100644
index 80c3401..0000000
--- a/backend/authentication.py
+++ /dev/null
@@ -1,48 +0,0 @@
-import os
-from datetime import datetime, timedelta
-from typing import Optional
-
-from jose import jwt
-from pydantic import BaseModel
-
-from backend.database import Database
-
-SECRET_KEY = os.environ.get("SECRET_KEY", "default_unsecure_key")
-ALGORITHM = "HS256"
-
-
-class User(BaseModel):
-    email: str = None
-    password: str = None
-
-
-def create_user(user: User) -> None:
-    with Database() as connection:
-        connection.execute(
-            "INSERT INTO users (email, password) VALUES (?, ?)", (user.email, user.password)
-        )
-
-
-def get_user(email: str) -> User:
-    with Database() as connection:
-        user_row = connection.execute("SELECT * FROM users WHERE email = ?", (email,))
-        for row in user_row:
-            return User(**row)
-        raise Exception("User not found")
-
-
-def authenticate_user(username: str, password: str) -> Optional[User]:
-    user = get_user(username)
-    if not user or not password == user.password:
-        return False
-    return user
-
-
-def create_access_token(*, data: dict, expires_delta: Optional[timedelta] = None) -> str:
-    to_encode = data.copy()
-    if expires_delta:
-        expire = datetime.utcnow() + expires_delta
-    else:
-        expire = datetime.utcnow() + timedelta(minutes=15)
-    to_encode.update({"exp": expire})
-    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
diff --git a/backend/main.py b/backend/main.py
index 30e2af9..d390964 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -21,6 +21,7 @@
 from backend.user_management import (
     ALGORITHM,
     SECRET_KEY,
+    UnsecureUser,
     User,
     authenticate_user,
     create_access_token,
@@ -35,7 +36,7 @@
 with Database() as connection:
     connection.initialize_schema()
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/user/login")
 
 
 async def get_current_user(token: str = Depends(oauth2_scheme)) -> User:
@@ -211,7 +212,8 @@ async def feedback_thumbs_down(
 
 
 @app.post("/user/signup")
-async def signup(user: User) -> dict:
+async def signup(user: UnsecureUser) -> dict:
+    user = User.from_unsecure_user(user)
     if user_exists(user.email):
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST, detail=f"User {user.email} already registered"
@@ -248,7 +250,9 @@ async def login(form_data: OAuth2PasswordRequestForm = Depends()) -> dict:
             headers={"WWW-Authenticate": "Bearer"},
         )
     access_token_expires = timedelta(minutes=60)
-    access_token = create_access_token(data=user.model_dump(), expires_delta=access_token_expires)
+    user_data = user.model_dump()
+    del user_data["hashed_password"]
+    access_token = create_access_token(data=user_data, expires_delta=access_token_expires)
     return {"access_token": access_token, "token_type": "bearer"}
 
 
diff --git a/backend/user_management.py b/backend/user_management.py
index f2a8f9d..9d21cd3 100644
--- a/backend/user_management.py
+++ b/backend/user_management.py
@@ -4,6 +4,7 @@
 
 from jose import jwt
 from pydantic import BaseModel
+import argon2
 
 from backend.database import Database
 
@@ -11,15 +12,24 @@
 ALGORITHM = "HS256"
 
 
+class UnsecureUser(BaseModel):
+    email: str = None
+    password: bytes = None
+
 class User(BaseModel):
     email: str = None
-    password: str = None
+    hashed_password: bytes = None
+
+    @classmethod
+    def from_unsecure_user(cls, unsecure_user: UnsecureUser):
+        hashed_password = argon2.hash_password(unsecure_user.password)
+        return cls(email=unsecure_user.email, hashed_password=hashed_password)
 
 
 def create_user(user: User) -> None:
     with Database() as connection:
         connection.execute(
-            "INSERT INTO users (email, password) VALUES (?, ?)", (user.email, user.password)
+            "INSERT INTO users (email, password) VALUES (?, ?)", (user.email, user.hashed_password)
         )
 
 
@@ -33,7 +43,7 @@ def get_user(email: str) -> Optional[User]:
     with Database() as connection:
         user_row = connection.fetchone("SELECT * FROM users WHERE email = ?", (email,))
         if user_row:
-            return User(email=user_row[0], password=user_row[1])
+            return User(email=user_row[0], hashed_password=user_row[1])
         return None
 
 
@@ -42,12 +52,15 @@ def delete_user(email: str) -> None:
         connection.execute("DELETE FROM users WHERE email = ?", (email,))
 
 
-def authenticate_user(username: str, password: str) -> Optional[User]:
+def authenticate_user(username: str, password: bytes) -> bool | User:
     user = get_user(username)
-    if not user or not password == user.password:
+    if not user:
         return False
-    return user
-
+    
+    if argon2.verify_password(user.hashed_password, password.encode("utf-8")):
+        return user
+    
+    return False
 
 def create_access_token(*, data: dict, expires_delta: Optional[timedelta] = None) -> str:
     to_encode = data.copy()
diff --git a/requirements.in b/requirements.in
index 710f433..2c7be94 100644
--- a/requirements.in
+++ b/requirements.in
@@ -42,3 +42,4 @@ mkdocs
 mkdocs-techdocs-core
 pymdown-extensions
 mkdocs_monorepo_plugin
+argon2-cffi
\ No newline at end of file