BanGUI/backend/app/main.py

"""BanGUI FastAPI application factory.

Call :func:`create_app` to obtain a configured :class:`fastapi.FastAPI`
instance suitable for direct use with an ASGI server (e.g. ``uvicorn``) or
in tests via ``httpx.AsyncClient``.

The lifespan handler manages all shared resources — database connection, HTTP
session, and scheduler — so every component can rely on them being available
on ``app.state`` throughout the request lifecycle.
"""

from __future__ import annotations

import logging
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator, Awaitable, Callable

    from starlette.responses import Response as StarletteResponse

import aiohttp
import aiosqlite
import structlog
from apscheduler.schedulers.asyncio import AsyncIOScheduler  # type: ignore[import-untyped]
from fastapi import FastAPI, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from starlette.middleware.base import BaseHTTPMiddleware

from app.config import Settings, get_settings
from app.db import init_db
from app.routers import (
    auth,
    bans,
    blocklist,
    config,
    dashboard,
    file_config,
    geo,
    health,
    history,
    jails,
    server,
    setup,
)
from app.tasks import blocklist_import, geo_cache_flush, geo_re_resolve, health_check
from app.utils.fail2ban_client import Fail2BanConnectionError, Fail2BanProtocolError

# ---------------------------------------------------------------------------
# Ensure the bundled fail2ban package is importable from fail2ban-master/
#
# The directory layout differs between local dev and the Docker image:
#   Local:  <repo-root>/backend/app/main.py  → fail2ban-master at parents[2]
#   Docker: /app/app/main.py                 → fail2ban-master at parents[1]
# Walk up from this file until we find a "fail2ban-master" sibling directory
# so the path resolution is environment-agnostic.
# ---------------------------------------------------------------------------


def _find_fail2ban_master() -> Path | None:
    """Return the first ``fail2ban-master`` directory found while walking up.

    Returns:
        Absolute :class:`~pathlib.Path` to the ``fail2ban-master`` directory,
        or ``None`` if no such directory exists among the ancestors.
    """
    here = Path(__file__).resolve()
    for ancestor in here.parents:
        candidate = ancestor / "fail2ban-master"
        if candidate.is_dir():
            return candidate
    return None


_fail2ban_master: Path | None = _find_fail2ban_master()
if _fail2ban_master is not None and str(_fail2ban_master) not in sys.path:
    sys.path.insert(0, str(_fail2ban_master))

log: structlog.stdlib.BoundLogger = structlog.get_logger()


# ---------------------------------------------------------------------------
# Logging configuration
# ---------------------------------------------------------------------------


def _configure_logging(log_level: str) -> None:
    """Configure structlog for production JSON output.

    Args:
        log_level: One of ``debug``, ``info``, ``warning``, ``error``, ``critical``.
    """
    level: int = logging.getLevelName(log_level.upper())
    logging.basicConfig(level=level, stream=sys.stdout, format="%(message)s")
    structlog.configure(
        processors=[
            structlog.contextvars.merge_contextvars,
            structlog.stdlib.filter_by_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.stdlib.add_logger_name,
            structlog.stdlib.add_log_level,
            structlog.stdlib.PositionalArgumentsFormatter(),
            structlog.processors.StackInfoRenderer(),
            structlog.processors.format_exc_info,
            structlog.processors.UnicodeDecoder(),
            structlog.processors.JSONRenderer(),
        ],
        wrapper_class=structlog.stdlib.BoundLogger,
        context_class=dict,
        logger_factory=structlog.stdlib.LoggerFactory(),
        cache_logger_on_first_use=True,
    )


# ---------------------------------------------------------------------------
# Lifespan
# ---------------------------------------------------------------------------


@asynccontextmanager
async def _lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    """Manage the lifetime of all shared application resources.

    Resources are initialised in order on startup and released in reverse
    order on shutdown.  They are stored on ``app.state`` so they are
    accessible to dependency providers and tests.

    Args:
        app: The :class:`fastapi.FastAPI` instance being started.
    """
    settings: Settings = app.state.settings
    _configure_logging(settings.log_level)

    log.info("bangui_starting_up", database_path=settings.database_path)

    # --- Application database ---
    db: aiosqlite.Connection = await aiosqlite.connect(settings.database_path)
    db.row_factory = aiosqlite.Row
    await init_db(db)
    app.state.db = db

    # --- Shared HTTP client session ---
    http_session: aiohttp.ClientSession = aiohttp.ClientSession()
    app.state.http_session = http_session

    # --- Pre-warm geo cache from the persistent store ---
    from app.services import geo_service  # noqa: PLC0415

    geo_service.init_geoip(settings.geoip_db_path)
    await geo_service.load_cache_from_db(db)

    # Log unresolved geo entries so the operator can see the scope of the issue.
    async with db.execute(
        "SELECT COUNT(*) FROM geo_cache WHERE country_code IS NULL"
    ) as cur:
        row = await cur.fetchone()
        unresolved_count: int = int(row[0]) if row else 0
    if unresolved_count > 0:
        log.warning("geo_cache_unresolved_ips", unresolved=unresolved_count)

    # --- Background task scheduler ---
    scheduler: AsyncIOScheduler = AsyncIOScheduler(timezone="UTC")
    scheduler.start()
    app.state.scheduler = scheduler

    # --- Health-check background probe ---
    health_check.register(app)

    # --- Blocklist import scheduled task ---
    blocklist_import.register(app)

    # --- Periodic geo cache flush to SQLite ---
    geo_cache_flush.register(app)

    # --- Periodic re-resolve of NULL-country geo entries ---
    geo_re_resolve.register(app)

    log.info("bangui_started")

    try:
        yield
    finally:
        log.info("bangui_shutting_down")
        scheduler.shutdown(wait=False)
        await http_session.close()
        await db.close()
        log.info("bangui_shut_down")


# ---------------------------------------------------------------------------
# Exception handlers
# ---------------------------------------------------------------------------


async def _unhandled_exception_handler(
    request: Request,
    exc: Exception,
) -> JSONResponse:
    """Return a sanitised 500 JSON response for any unhandled exception.

    The exception is logged with full context before the response is sent.
    No stack trace is leaked to the client.

    Args:
        request: The incoming FastAPI request.
        exc: The unhandled exception.

    Returns:
        A :class:`fastapi.responses.JSONResponse` with status 500.
    """
    log.error(
        "unhandled_exception",
        path=request.url.path,
        method=request.method,
        exc_info=exc,
    )
    return JSONResponse(
        status_code=500,
        content={"detail": "An unexpected error occurred. Please try again later."},
    )


async def _fail2ban_connection_handler(
    request: Request,
    exc: Fail2BanConnectionError,
) -> JSONResponse:
    """Return a ``502 Bad Gateway`` response when fail2ban is unreachable.

    Args:
        request: The incoming FastAPI request.
        exc: The :class:`~app.utils.fail2ban_client.Fail2BanConnectionError`.

    Returns:
        A :class:`fastapi.responses.JSONResponse` with status 502.
    """
    log.warning(
        "fail2ban_connection_error",
        path=request.url.path,
        method=request.method,
        error=str(exc),
    )
    return JSONResponse(
        status_code=502,
        content={"detail": f"Cannot reach fail2ban: {exc}"},
    )


async def _fail2ban_protocol_handler(
    request: Request,
    exc: Fail2BanProtocolError,
) -> JSONResponse:
    """Return a ``502 Bad Gateway`` response for fail2ban protocol errors.

    Args:
        request: The incoming FastAPI request.
        exc: The :class:`~app.utils.fail2ban_client.Fail2BanProtocolError`.

    Returns:
        A :class:`fastapi.responses.JSONResponse` with status 502.
    """
    log.warning(
        "fail2ban_protocol_error",
        path=request.url.path,
        method=request.method,
        error=str(exc),
    )
    return JSONResponse(
        status_code=502,
        content={"detail": f"fail2ban protocol error: {exc}"},
    )


# ---------------------------------------------------------------------------
# Setup-redirect middleware
# ---------------------------------------------------------------------------

# Paths that are always reachable, even before setup is complete.
_ALWAYS_ALLOWED: frozenset[str] = frozenset(
    {"/api/setup", "/api/health"},
)


class SetupRedirectMiddleware(BaseHTTPMiddleware):
    """Redirect all API requests to ``/api/setup`` until setup is done.

    Once setup is complete this middleware is a no-op.  Paths listed in
    :data:`_ALWAYS_ALLOWED` are exempt so the setup endpoint itself is
    always reachable.
    """

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[StarletteResponse]],
    ) -> StarletteResponse:
        """Intercept requests before they reach the router.

        Args:
            request: The incoming HTTP request.
            call_next: The next middleware / router handler.

        Returns:
            Either a ``307 Temporary Redirect`` to ``/api/setup`` or the
            normal router response.
        """
        path: str = request.url.path.rstrip("/") or "/"

        # Allow requests that don't need setup guard.
        if any(path.startswith(allowed) for allowed in _ALWAYS_ALLOWED):
            return await call_next(request)

        # If setup is not complete, block all other API requests.
        # Fast path: setup completion is a one-way transition.  Once it is
        # True it is cached on app.state so all subsequent requests skip the
        # DB query entirely.  The flag is reset only when the app restarts.
        if path.startswith("/api") and not getattr(
            request.app.state, "_setup_complete_cached", False
        ):
            db: aiosqlite.Connection | None = getattr(request.app.state, "db", None)
            if db is not None:
                from app.services import setup_service  # noqa: PLC0415

                if await setup_service.is_setup_complete(db):
                    request.app.state._setup_complete_cached = True
                else:
                    return RedirectResponse(
                        url="/api/setup",
                        status_code=status.HTTP_307_TEMPORARY_REDIRECT,
                    )

        return await call_next(request)


# ---------------------------------------------------------------------------
# Application factory
# ---------------------------------------------------------------------------


def create_app(settings: Settings | None = None) -> FastAPI:
    """Create and configure the BanGUI FastAPI application.

    This factory is the single entry point for creating the application.
    Tests can pass a custom ``settings`` object to override defaults
    without touching environment variables.

    Args:
        settings: Optional pre-built :class:`~app.config.Settings` instance.
            If ``None``, settings are loaded from the environment via
            :func:`~app.config.get_settings`.

    Returns:
        A fully configured :class:`fastapi.FastAPI` application ready for use.
    """
    resolved_settings: Settings = settings if settings is not None else get_settings()

    app: FastAPI = FastAPI(
        title="BanGUI",
        description="Web interface for monitoring, managing, and configuring fail2ban.",
        version="0.1.0",
        lifespan=_lifespan,
    )

    # Store settings on app.state so the lifespan handler can access them.
    app.state.settings = resolved_settings

    # --- CORS ---
    # In production the frontend is served by the same origin.
    # CORS is intentionally permissive only in development.
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["http://localhost:5173"],  # Vite dev server
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # --- Middleware ---
    # Note: middleware is applied in reverse order of registration.
    # The setup-redirect must run *after* CORS, so it is added last.
    app.add_middleware(SetupRedirectMiddleware)

    # --- Exception handlers ---
    # Ordered from most specific to least specific.  FastAPI evaluates handlers
    # in the order they were registered, so fail2ban network errors get a 502
    # rather than falling through to the generic 500 handler.
    app.add_exception_handler(Fail2BanConnectionError, _fail2ban_connection_handler)  # type: ignore[arg-type]
    app.add_exception_handler(Fail2BanProtocolError, _fail2ban_protocol_handler)  # type: ignore[arg-type]
    app.add_exception_handler(Exception, _unhandled_exception_handler)

    # --- Routers ---
    app.include_router(health.router)
    app.include_router(setup.router)
    app.include_router(auth.router)
    app.include_router(dashboard.router)
    app.include_router(jails.router)
    app.include_router(bans.router)
    app.include_router(geo.router)
    app.include_router(config.router)
    app.include_router(file_config.router)
    app.include_router(server.router)
    app.include_router(history.router)
    app.include_router(blocklist.router)

    return app