b2b/src/main.py

from __future__ import annotations
from contextlib import asynccontextmanager
import secrets
from typing import AsyncGenerator
from fastapi import Depends, FastAPI
from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
from fastapi.responses import HTMLResponse
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from starlette.middleware.cors import CORSMiddleware
from starlette.exceptions import HTTPException
from fastapi.exceptions import RequestValidationError
from src.application.domain.exceptions import ApplicationException, UnauthorizedException
from src.infrastructure.vault import JwtKeyStore, start_jwt_keys_scheduler
from src.infrastructure.utils import generate_instance_id
from src.infrastructure.logger import logger
from src.infrastructure.config import settings
from src.presentation.handlers import (
    application_exception_handler,
    http_exception_handler,
    unhandled_exception_handler,
    validation_exception_handler,
)
from src.presentation.middleware import TraceIDMiddleware, SecurityHeadersMiddleware
from src.infrastructure.crypto.wallet_crypto import load_master_key_from_vault
from src.presentation.routing import purchase_requests_router, v1_router

security = HTTPBasic()


async def verify_credentials(credentials: HTTPBasicCredentials = Depends(security)) -> HTTPBasicCredentials:
    user_ok = secrets.compare_digest(credentials.username, settings.DOCS_USERNAME)
    pass_ok = secrets.compare_digest(credentials.password, settings.DOCS_PASSWORD)
    if not (user_ok and pass_ok):
        raise UnauthorizedException(
            message='Unauthorized',
            headers={'WWW-Authenticate': 'Basic'},
        )
    return credentials


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:

    instance_id = generate_instance_id()
    logger.set_instance_id(instance_id)
    logger.info(f'B2B service instance started with id {instance_id}')

    if not settings.VAULT_ROLE_ID.strip() or not settings.VAULT_SECRET_ID.strip():
        raise RuntimeError('VAULT_ROLE_ID and VAULT_SECRET_ID must be set')

    jwt_store = JwtKeyStore(
        vault_addr=settings.VAULT_ADDR,
        vault_role_id=settings.VAULT_ROLE_ID,
        vault_secret_id=settings.VAULT_SECRET_ID,
        vault_namespace=settings.VAULT_NAMESPACE,
        mount_point=settings.VAULT_MOUNT_POINT,
        kid_path=settings.VAULT_JWT_KID_PATH,
        kids_prefix=settings.VAULT_JWT_KIDS_PREFIX,
    )

    await jwt_store.refresh()

    jwt_scheduler = start_jwt_keys_scheduler(jwt_store, refresh_seconds=settings.JWT_KEYS_REFRESH_SECONDS)

    app.state.jwt_key_store = jwt_store
    app.state.jwt_keys_scheduler = jwt_scheduler

    load_master_key_from_vault(
        vault_addr=settings.VAULT_ADDR,
        vault_role_id=settings.VAULT_ROLE_ID,
        vault_secret_id=settings.VAULT_SECRET_ID,
        vault_namespace=settings.VAULT_NAMESPACE,
        mount_point=settings.VAULT_MOUNT_POINT,
        path=settings.VAULT_CRYPTO_MASTER_KEY_PATH,
    )
    logger.info('Crypto master key loaded from Vault')
    try:
        yield
    finally:
        sched = getattr(app.state,'jwt_keys_scheduler',None)
        if sched:
            sched.shutdown(wait=False)
        logger.info(f'B2B service instance ended with id {instance_id}')


app: FastAPI = FastAPI(
    redoc_url=None,
    docs_url=None,
    lifespan=lifespan,
    title='Elcsa. B2B Service',
    version='1.0.0',
    description='Purchase requests API for legal entity client users.',
    license_info={
        'name': 'MIT',
        'url': 'https://opensource.org/licenses/MIT',
    },
)

app.add_exception_handler(RequestValidationError, validation_exception_handler)
app.add_exception_handler(HTTPException, http_exception_handler)
app.add_exception_handler(ApplicationException, application_exception_handler)
app.add_exception_handler(Exception, unhandled_exception_handler)

app.include_router(v1_router)
app.include_router(purchase_requests_router)

app.add_middleware(TraceIDMiddleware, logger=logger)
app.add_middleware(
    SecurityHeadersMiddleware,
    hsts=True,
    hsts_preload=False,
    frame_options='DENY',
    referrer_policy='strict-origin-when-cross-origin',
    content_security_policy="default-src 'self'; frame-ancestors 'none'; base-uri 'self'; object-src 'none'",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=[],
    allow_origin_regex=settings.CORS_ALLOW_ORIGIN_REGEX,
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=['*'],
)


@app.get('/docs', include_in_schema=False)
async def custom_swagger_ui_html(_credentials: HTTPBasicCredentials = Depends(verify_credentials)) -> HTMLResponse:
    '''Custom Swagger documentation, optionally protected with basic authentication.'''
    return get_swagger_ui_html(
        openapi_url=getattr(app, 'openapi_url', '/openapi.json'),
        title=getattr(app, 'title', 'FastAPI') + ' - Swagger UI',
        oauth2_redirect_url=getattr(app, 'swagger_ui_oauth2_redirect_url', None),
        swagger_js_url='https://unpkg.com/swagger-ui-dist@5/swagger-ui-bundle.js',
        swagger_css_url='https://unpkg.com/swagger-ui-dist@5/swagger-ui.css',
    )


@app.get('/redoc', include_in_schema=False)
async def custom_redoc_html(_credentials: HTTPBasicCredentials = Depends(verify_credentials)) -> HTMLResponse:
    '''Custom ReDoc documentation, optionally protected with basic authentication.'''
    return get_redoc_html(
        openapi_url=getattr(app, 'openapi_url', '/openapi.json'),
        title=getattr(app, 'title', 'FastAPI') + ' - ReDoc',
        redoc_js_url='https://cdn.jsdelivr.net/npm/redoc@next/bundles/redoc.standalone.js',
    )


@app.post('/ping')
async def ping() -> dict[str, str]:
    return {
        'message': 'pong',
        'status': 'ok',
    }

@app.get('/health')
async def health() -> dict[str, str]:
    return {
        'status': 'ok',
    }