Authentication with JWT, Part 1

florist/api/auth/__init__.py

@@ -0,0 +1 @@
+"""Module for handling authentication and token creation."""

florist/api/auth/token.py

@@ -0,0 +1,148 @@
+"""Module for handling token and user creation."""
+
+import hashlib
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+import bcrypt
+import jwt
+from motor.motor_asyncio import AsyncIOMotorDatabase
+from pydantic import BaseModel
+
+from florist.api.db.client_entities import UserDAO
+from florist.api.db.server_entities import User
+
+
+ENCRYPTION_ALGORITHM = "HS256"
+DEFAULT_USERNAME = "admin"
+DEFAULT_PASSWORD = "admin"
+TOKEN_EXPIRATION_TIMEDELTA = timedelta(days=7)
+
+
+class Token(BaseModel):
+    """Define the Token model."""
+
+    access_token: str
+    token_type: str
+
+    class Config:
+        """Config for the Token model."""
+
+        allow_population_by_field_name = True
+        schema_extra = {
+            "example": {
+                "access_token": "LQv3c1yqBWVHxkd0LHAkCOYz6T",
+                "token_type": "bearer",
+            },
+        }
+
+
+class AuthUser(BaseModel):
+    """Define the User model to be returned by the API."""
+
+    uuid: str
+    username: str
+
+    class Config:
+        """Config for the AuthUser model."""
+
+        allow_population_by_field_name = True
+        schema_extra = {
+            "example": {
+                "uuid": "LQv3c1yqBWVHxkd0LHAkCOYz6T",
+                "username": "admin",
+            },
+        }
+
+
+def _simple_hash(word: str) -> str:
+    """
+    Hash a word with sha256.
+
+    WARNING: This is not a secure hash function, it is only meant to obscure
+    plain text words. DO NOT use this for generating encrypted passwords for the
+    authentication users. For that, use the _password_hash function instead.
+
+    :param word: (str) the word to hash.
+    :return: (str) the word hashed as a sha256 hexadecimal string.
+    """
+    return hashlib.sha256(word.encode("utf-8")).hexdigest()
+
+
+def _password_hash(password: str) -> str:
+    """
+    Hash a password with bcrypt.
+
+    :param password: (str) the password to hash.
+    :return: (str) the hashed password.
+    """
+    password_bytes = password.encode("utf-8")
+    salt = bcrypt.gensalt()
+    hashed_password = bcrypt.hashpw(password=password_bytes, salt=salt)
+    return hashed_password.decode("utf-8")
+
+
+def verify_password(password: str, hashed_password: str) -> bool:
+    """
+    Verify if a password matches a hashed password.
+
+    :param password: (str) the password to verify.
+    :param hashed_password: (str) the hashed password to verify against.
+    :return: (bool) True if the password matches the hashed password, False otherwise.
+    """
+    return bcrypt.checkpw(password.encode("utf-8"), hashed_password.encode("utf-8"))
+
+
+async def make_default_server_user(database: AsyncIOMotorDatabase[Any]) -> User:
+    """
+    Make a default server user.
+
+    :param database: (AsyncIOMotorDatabase[Any]) the database to create the user in.
+    :return: (User) the default server user.
+    """
+    hashed_password = _password_hash(_simple_hash(DEFAULT_PASSWORD))
+    user = User(username=DEFAULT_USERNAME, hashed_password=hashed_password)
+    await user.create(database)
+    return user
+
+
+def make_default_client_user() -> UserDAO:
+    """
+    Make a default client user.
+
+    :return: (User) the default client user.
+    """
+    hashed_password = _password_hash(_simple_hash(DEFAULT_PASSWORD))
+    user = UserDAO(username=DEFAULT_USERNAME, hashed_password=hashed_password)
+    user.save()
+    return user
+
+
+def create_access_token(
+    data: dict[str, Any], secret_key: str, expiration_delta: timedelta = TOKEN_EXPIRATION_TIMEDELTA
+) -> str:
+    """
+    Create an access token.
+
+    :param data: (dict) the data to encode in the token.
+    :param secret_key: (str) the user's secret key to encode the token.
+    :param expiration_delta: (timedelta) the expiration time of the token.
+    :return: (str) the access token.
+    """
+    to_encode = data.copy()
+    expire = datetime.now(timezone.utc) + expiration_delta
+    to_encode.update({"exp": expire})
+    return jwt.encode(to_encode, secret_key, algorithm=ENCRYPTION_ALGORITHM)
+
+
+def decode_access_token(token: str, secret_key: str) -> dict[str, Any]:
+    """
+    Decode an access token.
+
+    :param token: (str) the token to decode.
+    :param secret_key: (str) the user's secret key to decode the token.
+    :return: (dict) the decoded token information.
+    """
+    data = jwt.decode(token, secret_key, algorithms=[ENCRYPTION_ALGORITHM])
+    assert isinstance(data, dict)
+    return data

florist/api/client.py

@@ -3,30 +3,46 @@
 import logging
 import os
 import signal
+from contextlib import asynccontextmanager
 from pathlib import Path
+from typing import Any, AsyncGenerator
 from uuid import uuid4
 
 import torch
-from fastapi import FastAPI
+from fastapi import Depends, FastAPI
 from fastapi.responses import JSONResponse
 from fl4health.utils.metrics import Accuracy
 
+from florist.api.auth.token import DEFAULT_USERNAME, make_default_client_user
 from florist.api.clients.clients import Client
 from florist.api.clients.optimizers import Optimizer
-from florist.api.db.client_entities import ClientDAO
+from florist.api.db.client_entities import ClientDAO, UserDAO
 from florist.api.launchers.local import launch_client
 from florist.api.models.models import Model
 from florist.api.monitoring.logs import get_client_log_file_path
 from florist.api.monitoring.metrics import RedisMetricsReporter, get_from_redis, get_host_and_port_from_address
+from florist.api.routes.client.auth import check_token
+from florist.api.routes.client.auth import router as auth_router
 
 
-app = FastAPI()
+LOGGER = logging.getLogger("uvicorn.error")
 
 
-LOGGER = logging.getLogger("uvicorn.error")
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncGenerator[Any, Any]:
+    """Set up function for app startup and shutdown."""
+    # Create default user if it does not exist
+    if not UserDAO.exists(DEFAULT_USERNAME):
+        make_default_client_user()
+
+    yield
+
 
+app = FastAPI(lifespan=lifespan)
+app.include_router(auth_router, tags=["auth"], prefix="/api/client/auth")
 
-@app.get("/api/client/connect")
+
+@app.get("/api/client/connect", dependencies=[Depends(check_token)])
 def connect() -> JSONResponse:
     """
     Confirm the client is up and ready to accept instructions.
@@ -36,7 +52,7 @@ def connect() -> JSONResponse:
     return JSONResponse({"status": "ok"})
 
 
-@app.get("/api/client/start")
+@app.get("/api/client/start", dependencies=[Depends(check_token)])
 def start(
     server_address: str,
     client: Client,
@@ -53,6 +69,7 @@ def start(
     :param client: (Client) the client to be used for training.
     :param data_path: (str) the path where the training data is located.
     :param redis_address: (str) the address for the Redis instance for metrics reporting.
+
     :return: (JSONResponse) If successful, returns 200 with a JSON containing the UUID for the client in the
         format below, which can be used to pull metrics from Redis.
             {
@@ -95,7 +112,7 @@ def start(
         return JSONResponse({"error": str(ex)}, status_code=500)
 
 
-@app.get("/api/client/check_status/{client_uuid}")
+@app.get("/api/client/check_status/{client_uuid}", dependencies=[Depends(check_token)])
 def check_status(client_uuid: str, redis_address: str) -> JSONResponse:
     """
     Retrieve value at key client_uuid in redis if it exists.
@@ -120,7 +137,7 @@ def check_status(client_uuid: str, redis_address: str) -> JSONResponse:
         return JSONResponse({"error": str(ex)}, status_code=500)
 
 
-@app.get("/api/client/get_log/{uuid}")
+@app.get("/api/client/get_log/{uuid}", dependencies=[Depends(check_token)])
 def get_log(uuid: str) -> JSONResponse:
     """
     Return the contents of the logs for the given client uuid.
@@ -147,12 +164,13 @@ def get_log(uuid: str) -> JSONResponse:
         return JSONResponse({"error": str(ex)}, status_code=500)
 
 
-@app.get("/api/client/stop/{uuid}")
+@app.get("/api/client/stop/{uuid}", dependencies=[Depends(check_token)])
 def stop(uuid: str) -> JSONResponse:
     """
     Stop the client with given UUID.
 
     :param uuid: (str) the UUID of the client to be stopped.
+
     :return: (JSONResponse) If successful, returns 200. If not successful, returns the appropriate
         error code with a JSON with the format below:
             {"error": <error message>}

florist/api/db/client_entities.py

@@ -1,6 +1,7 @@
 """Definitions for the SQLIte database entities (client database)."""
 
 import json
+import secrets
 import sqlite3
 from abc import ABC, abstractmethod
 from typing import Optional
@@ -55,7 +56,7 @@ def find(cls, uuid: str) -> Self:
         for result in results:
             return cls.from_json(result[1])
 
-        raise ValueError(f"Client with uuid '{uuid}' not found.")
+        raise ValueError(f"{cls.table_name} with uuid '{uuid}' not found.")
 
     @classmethod
     def exists(cls, uuid: str) -> bool:
@@ -167,3 +168,55 @@ def to_json(self) -> str:
                 "pid": self.pid,
             }
         )
+
+
+class UserDAO(EntityDAO):
+    """Data Access Object (DAO) for the User SQLite entity."""
+
+    table_name = "User"
+
+    def __init__(self, username: str, hashed_password: str):
+        """
+        Initialize a User entity.
+
+        :param uuid: (str) the UUID of the user.
+        :param username: (str) the username of the user.
+        :param hashed_password: (str) the hashed password of the user.
+        """
+        # The UUID for the user is the username
+        super().__init__(uuid=username)
+
+        self.username = username
+        self.hashed_password = hashed_password
+
+        # always create a new random secret key
+        self.secret_key = secrets.token_hex(32)
+
+    @classmethod
+    def from_json(cls, json_data: str) -> Self:
+        """
+        Convert from a JSON string into an instance of User.
+
+        :param json_data: the user's data as a JSON string.
+        :return: (Self) and instancxe of UserDAO populated with the JSON data.
+        """
+        data = json.loads(json_data)
+        user = cls(data["username"], data["hashed_password"])
+        user.uuid = data["uuid"]
+        user.secret_key = data["secret_key"]
+        return user
+
+    def to_json(self) -> str:
+        """
+        Convert the user data into a JSON string.
+
+        :return: (str) the user data as a JSON string.
+        """
+        return json.dumps(
+            {
+                "uuid": self.uuid,
+                "username": self.username,
+                "hashed_password": self.hashed_password,
+                "secret_key": self.secret_key,
+            }
+        )