BanGUI/backend/app/routers/geo.py

"""Geo / IP lookup router.

Provides the IP enrichment endpoints:

* ``GET /api/geo/lookup/{ip}`` — ban status, ban history, and geo info for an IP
* ``POST /api/geo/re-resolve`` — retry all previously failed geo lookups
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Annotated

if TYPE_CHECKING:
    import aiohttp

import aiosqlite
from fastapi import APIRouter, Depends, HTTPException, Path, Request, status

from app.dependencies import AuthDep, get_db
from app.models.geo import GeoCacheStatsResponse, GeoDetail, IpLookupResponse
from app.services import geo_service, jail_service
from app.utils.fail2ban_client import Fail2BanConnectionError

router: APIRouter = APIRouter(prefix="/api/geo", tags=["Geo"])

_IpPath = Annotated[str, Path(description="IPv4 or IPv6 address to look up.")]


@router.get(
    "/lookup/{ip}",
    response_model=IpLookupResponse,
    summary="Look up ban status and geo information for an IP",
)
async def lookup_ip(
    request: Request,
    _auth: AuthDep,
    ip: _IpPath,
) -> IpLookupResponse:
    """Return current ban status, geo data, and network information for an IP.

    Checks every running fail2ban jail to determine whether the IP is
    currently banned, and enriches the result with country, ASN, and
    organisation data from ip-api.com.

    Args:
        request: Incoming request (used to access ``app.state``).
        _auth: Validated session — enforces authentication.
        ip: The IP address to look up.

    Returns:
        :class:`~app.models.geo.IpLookupResponse` with ban status and geo data.

    Raises:
        HTTPException: 400 when *ip* is not a valid IP address.
        HTTPException: 502 when fail2ban is unreachable.
    """
    socket_path: str = request.app.state.settings.fail2ban_socket
    http_session: aiohttp.ClientSession = request.app.state.http_session

    async def _enricher(addr: str) -> geo_service.GeoInfo | None:
        return await geo_service.lookup(addr, http_session)

    try:
        result = await jail_service.lookup_ip(
            socket_path,
            ip,
            geo_enricher=_enricher,
        )
    except ValueError as exc:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=str(exc),
        ) from exc
    except Fail2BanConnectionError as exc:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail=f"Cannot reach fail2ban: {exc}",
        ) from exc

    raw_geo = result.get("geo")
    geo_detail: GeoDetail | None = None
    if raw_geo is not None:
        geo_detail = GeoDetail(
            country_code=raw_geo.country_code,
            country_name=raw_geo.country_name,
            asn=raw_geo.asn,
            org=raw_geo.org,
        )

    return IpLookupResponse(
        ip=result["ip"],
        currently_banned_in=result["currently_banned_in"],
        geo=geo_detail,
    )


# ---------------------------------------------------------------------------
# POST /api/geo/re-resolve
# ---------------------------------------------------------------------------


# ---------------------------------------------------------------------------
# GET /api/geo/stats
# ---------------------------------------------------------------------------


@router.get(
    "/stats",
    response_model=GeoCacheStatsResponse,
    summary="Geo cache diagnostic counters",
)
async def geo_stats(
    _auth: AuthDep,
    db: Annotated[aiosqlite.Connection, Depends(get_db)],
) -> GeoCacheStatsResponse:
    """Return diagnostic counters for the geo cache subsystem.

    Useful for operators and the UI to gauge geo-resolution health.

    Args:
        _auth: Validated session — enforces authentication.
        db: BanGUI application database connection.

    Returns:
        :class:`~app.models.geo.GeoCacheStatsResponse` with current counters.
    """
    stats: dict[str, int] = await geo_service.cache_stats(db)
    return GeoCacheStatsResponse(**stats)


@router.post(
    "/re-resolve",
    summary="Re-resolve all IPs whose country could not be determined",
)
async def re_resolve_geo(
    request: Request,
    _auth: AuthDep,
    db: Annotated[aiosqlite.Connection, Depends(get_db)],
) -> dict[str, int]:
    """Retry geo resolution for every IP in ``geo_cache`` with a null country.

    Clears the in-memory negative cache first so that previously failing IPs
    are immediately eligible for a new API attempt.

    Args:
        request: Incoming request (used to access ``app.state.http_session``).
        _auth: Validated session — enforces authentication.
        db: BanGUI application database (for reading/writing ``geo_cache``).

    Returns:
        JSON object ``{"resolved": N, "total": M}`` where *N* is the number
        of IPs that gained a country code and *M* is the total number of IPs
        that were retried.
    """
    # Collect all IPs in geo_cache that still lack a country code.
    unresolved = await geo_service.get_unresolved_ips(db)

    if not unresolved:
        return {"resolved": 0, "total": 0}

    # Clear negative cache so these IPs bypass the TTL check.
    geo_service.clear_neg_cache()

    http_session: aiohttp.ClientSession = request.app.state.http_session
    geo_map = await geo_service.lookup_batch(unresolved, http_session, db=db)

    resolved_count = sum(
        1 for info in geo_map.values() if info.country_code is not None
    )
    return {"resolved": resolved_count, "total": len(unresolved)}