Add debug endpoints and implement token auth

Co-authored-by: Joe Banks <[email protected]>

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all install relock lock lockci lint lintdeps precommit test retest
+.PHONY: all install relock lock lockci lint lintdeps precommit test retest seed
 
 all: install precommit lint
 
@@ -28,5 +28,8 @@ test:
 retest:
 	pytest -n 4 --lf
 
+seed:
+	cd thallium-backend && poetry run python -m scripts.seed
+
 revision:
 	cd thallium-backend && poetry run alembic revision --autogenerate -m CHANGEME

docker-compose.yml

@@ -20,7 +20,7 @@ services:
       cache_from:
         - type=registry,ref=ghcr.io/owl-corp/thallium-backend:latest
     restart: unless-stopped
-    command: ["alembic upgrade head && uvicorn src.app:fastapi_app --host 0.0.0.0 --port 8000 --reload"]
+    command: ["alembic upgrade head && uvicorn src.app:fastapi_app --host 0.0.0.0 --port 8000 --reload --no-server-header"]
     volumes:
       - ./thallium-backend/src:/thallium/src:ro
     env_file:
@@ -29,6 +29,7 @@ services:
       BACKEND_DATABASE_URL: postgresql+psycopg_async://thallium:thallium@postgres:5432/thallium
       BACKEND_TOKEN: suitable-for-development-only
       BACKEND_APP_PREFIX: /api
+      BACKEND_SIGNING_KEY: super-secure-key
     ports:
       - "8000:8000"
     depends_on:

pyproject.toml

@@ -16,6 +16,7 @@ sqlalchemy = {version = "^2.0.32", extras = ["asyncio"]}
 psycopg = {version = "^3.2.1", extras = ["binary"]}
 pydantic = "^2.8.2"
 pydantic-settings = "^2.4.0"
+pyjwt = "^2.9.0"
 uvicorn = "^0.30.6"
 
 [tool.poetry.group.linting.dependencies]

thallium-backend/Dockerfile

@@ -19,4 +19,4 @@ COPY . .
 HEALTHCHECK --start-period=5s --interval=30s --timeout=1s CMD curl http://localhost/heartbeat || exit 1
 
 ENTRYPOINT ["/bin/bash", "-c"]
-CMD ["alembic upgrade head && uvicorn src.app:fastapi_app --host 0.0.0.0 --port 8000"]
+CMD ["alembic upgrade head && uvicorn src.app:fastapi_app --host 0.0.0.0 --port 8000 --no-server-header"]

thallium-backend/migrations/versions/1724025217-bd897d0f21e1_add_products_users_vouchers.py

@@ -0,0 +1,70 @@
+"""
+Add products, users and vouchers.
+
+Revision ID: bd897d0f21e1
+Revises:
+Create Date: 2024-08-18 23:53:37.211777+00:00
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "bd897d0f21e1"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Apply this migration."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "products",
+        sa.Column("product_id", sa.Integer(), nullable=False),
+        sa.Column("name", sa.String(), nullable=False),
+        sa.Column("description", sa.String(), nullable=False),
+        sa.Column("price", sa.Numeric(), nullable=False),
+        sa.Column("image", sa.LargeBinary(), nullable=False),
+        sa.Column("id", sa.Uuid(), server_default=sa.text("gen_random_uuid()"), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.PrimaryKeyConstraint("product_id", "id", name=op.f("products_pk")),
+    )
+    op.create_table(
+        "users",
+        sa.Column("permissions", sa.Integer(), nullable=False),
+        sa.Column("id", sa.Uuid(), server_default=sa.text("gen_random_uuid()"), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.PrimaryKeyConstraint("id", name=op.f("users_pk")),
+    )
+    op.create_table(
+        "vouchers",
+        sa.Column("voucher_code", sa.String(), nullable=False),
+        sa.Column("active", sa.Boolean(), nullable=False),
+        sa.Column("balance", sa.Numeric(), nullable=False),
+        sa.Column("id", sa.Uuid(), server_default=sa.text("gen_random_uuid()"), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.PrimaryKeyConstraint("id", name=op.f("vouchers_pk")),
+    )
+    op.create_index(
+        "ix_unique_active_voucher_code",
+        "vouchers",
+        ["voucher_code"],
+        unique=True,
+        postgresql_where=sa.text("active"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Revert this migration."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("vouchers_voucher_code_ix"), table_name="vouchers")
+    op.drop_index("ix_unique_active_voucher_code", table_name="vouchers", postgresql_where=sa.text("active"))
+    op.drop_table("vouchers")
+    op.drop_table("users")
+    op.drop_table("products")
+    # ### end Alembic commands ###

thallium-backend/requirements.txt

@@ -307,6 +307,9 @@ pydantic-settings==2.4.0 ; python_full_version >= "3.12.0" and python_full_versi
 pydantic==2.8.2 ; python_full_version >= "3.12.0" and python_full_version < "4.0.0" \
     --hash=sha256:6f62c13d067b0755ad1c21a34bdd06c0c12625a22b0fc09c6b149816604f7c2a \
     --hash=sha256:73ee9fddd406dc318b885c7a2eab8a6472b68b8fb5ba8150949fc3db939f23c8
+pyjwt==2.9.0 ; python_full_version >= "3.12.0" and python_full_version < "4.0.0" \
+    --hash=sha256:3b02fb0f44517787776cf48f2ae25d8e14f300e6d7545a4315cee571a415e850 \
+    --hash=sha256:7e1e5b56cc735432a7369cbfa0efe50fa113ebecdc04ae6922deba8b84582d0c
 python-dotenv==1.0.1 ; python_full_version >= "3.12.0" and python_full_version < "4.0.0" \
     --hash=sha256:e324ee90a023d808f1959c46bcbc04446a10ced277783dc6ee09987c37ec10ca \
     --hash=sha256:f7b63ef50f1b690dddf550d03497b66d609393b40b564ed0d674909a68ebf16a

thallium-backend/scripts/seed.py

@@ -0,0 +1,20 @@
+import asyncio
+
+from src.orm import Voucher
+from src.settings import Connections
+
+
+async def main() -> None:
+    """Seed the database with some test data."""
+    async with Connections.DB_SESSION_MAKER() as session, session.begin():
+        session.add_all(
+            [
+                Voucher(voucher_code="k1p", balance="13.37", active=False),
+                Voucher(voucher_code="k1p", balance="13.37", active=False),
+                Voucher(voucher_code="k1p", balance="13.37"),
+            ]
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

thallium-backend/src/app.py

@@ -1,6 +1,8 @@
 import logging
+import time
+from collections.abc import Awaitable, Callable
 
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Response
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
 
@@ -24,3 +26,19 @@ def pydantic_validation_error(request: Request, error: RequestValidationError) -
     """Raise a warning for pydantic validation errors, before returning."""
     log.warning("Error from %s: %s", request.url, error)
     return JSONResponse({"error": str(error)}, status_code=422)
+
+
+@fastapi_app.middleware("http")
+async def add_process_time_and_security_headers(
+    request: Request,
+    call_next: Callable[[Request], Awaitable[Response]],
+) -> Response:
+    """Add process time and some security headers before sending the response."""
+    start_time = time.perf_counter()
+    response = await call_next(request)
+    response.headers["X-Process-Time"] = str(time.perf_counter() - start_time)
+    response.headers["X-Frame-Options"] = "DENY"
+    response.headers["X-XSS-Protection"] = "1; mode=block"
+    response.headers["Strict-Transport-Security"] = "max-age=31536000"
+    response.headers["X-Content-Type-Options"] = "nosniff"
+    return response

thallium-backend/src/auth.py

@@ -0,0 +1,110 @@
+import logging
+import typing as t
+from datetime import UTC, datetime, timedelta
+from enum import IntFlag
+from uuid import uuid4
+
+import jwt
+from fastapi import HTTPException, Request
+from fastapi.security import HTTPAuthorizationCredentials
+from fastapi.security.http import HTTPBase
+
+from src.dto.users import User, UserPermission
+from src.settings import CONFIG
+
+log = logging.getLogger(__name__)
+
+
+class UserTypes(IntFlag):
+    """All types of users."""
+
+    VOUCHER_USER = 2**0
+    REGULAR_USER = 2**1
+
+
+class TokenAuth(HTTPBase):
+    """Ensure all requests with this auth enabled include an auth header with the expected token."""
+
+    def __init__(
+        self,
+        *,
+        auto_error: bool = True,
+        allow_vouchers: bool = False,
+        allow_regular_users: bool = False,
+    ) -> None:
+        super().__init__(scheme="token", auto_error=auto_error)
+        self.allow_vouchers = allow_vouchers
+        self.allow_regular_users = allow_regular_users
+
+    async def __call__(self, request: Request) -> HTTPAuthorizationCredentials:
+        """Parse the token in the auth header, and check it matches with the expected token."""
+        creds: HTTPAuthorizationCredentials = await super().__call__(request)
+        if creds.scheme.lower() != "token":
+            raise HTTPException(
+                status_code=401,
+                detail="Incorrect scheme passed",
+            )
+        if self.allow_regular_users and creds.credentials == CONFIG.super_admin_token.get_secret_value():
+            request.state.user = User(user_id=uuid4(), permissions=~UserPermission(0))
+            return
+
+        jwt_data = verify_jwt(
+            creds.credentials,
+            allow_vouchers=self.allow_vouchers,
+            allow_regular_users=self.allow_regular_users,
+        )
+        if not jwt_data:
+            raise HTTPException(
+                status_code=403,
+                detail="Invalid authentication credentials",
+            )
+        if jwt_data["iss"] == "thallium:user":
+            request.state.user_id = jwt_data["sub"]
+        else:
+            request.state.voucher_id = jwt_data["sub"]
+
+
+def build_jwt(
+    identifier: str,
+    user_type: t.Literal["voucher", "user"],
+) -> str:
+    """Build & sign a jwt."""
+    return jwt.encode(
+        payload={
+            "sub": identifier,
+            "iss": f"thallium:{user_type}",
+            "exp": datetime.now(tz=UTC) + timedelta(minutes=30),
+            "nbf": datetime.now(tz=UTC) - timedelta(minutes=1),
+        },
+        key=CONFIG.signing_key.get_secret_value(),
+    )
+
+
+def verify_jwt(
+    jwt_data: str,
+    *,
+    allow_vouchers: bool,
+    allow_regular_users: bool,
+) -> dict | None:
+    """Return and verify the given JWT."""
+    issuers = []
+    if allow_vouchers:
+        issuers.append("thallium:voucher")
+    if allow_regular_users:
+        issuers.append("thallium:user")
+    try:
+        return jwt.decode(
+            jwt_data,
+            key=CONFIG.signing_key.get_secret_value(),
+            issuer=issuers,
+            algorithms=("HS256",),
+            options={"require": ["exp", "iss", "sub", "nbf"]},
+        )
+    except jwt.InvalidIssuerError as e:
+        raise HTTPException(403, "Your user type does not have access to this resource") from e
+    except jwt.InvalidSignatureError as e:
+        raise HTTPException(401, "Invalid JWT signature") from e
+    except (jwt.DecodeError, jwt.MissingRequiredClaimError, jwt.InvalidAlgorithmError) as e:
+        raise HTTPException(401, "Invalid JWT passed") from e
+    except (jwt.ImmatureSignatureError, jwt.ExpiredSignatureError) as e:
+        raise HTTPException(401, "JWT not valid for current time") from e

thallium-backend/src/dto/__init__.py

@@ -0,0 +1,5 @@
+from .login import VoucherClaim, VoucherLogin
+from .users import User
+from .vouchers import Voucher
+
+__all__ = ("LoginData", "User", "Voucher", "VoucherClaim", "VoucherLogin")

thallium-backend/src/dto/login.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel
+
+
+class VoucherLogin(BaseModel):
+    """The data needed to login with a voucher."""
+
+    voucher_code: str
+
+
+class VoucherClaim(VoucherLogin):
+    """A JWT for a verified voucher."""
+
+    jwt: str