users/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.cache import create_redis_client
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.presentation.routing import me_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'Users service instance started with id {instance_id}')

    app.state.redis = create_redis_client()

    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
    yield
    await app.state.redis.aclose()
    logger.info(f'Users service instance ended with id {instance_id}')


app: FastAPI = FastAPI(
    redoc_url=None,
    docs_url=None,
    lifespan=lifespan,
    title='Elcsa. Users Service',
    version='1.0.0',
    description='',
    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(me_router)
# app.include_router(me_devices_router)
# app.include_router(me_deals_router)


# Added middleware
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='.*',
    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',
    }