ripe-cidr-collector/api_server.py

from datetime import datetime, timezone
from enum import Enum
import ipaddress
import logging
import os
import re
import secrets
import sqlite3
from typing import Literal, Optional

from apscheduler.triggers.cron import CronTrigger
from fastapi import Depends, FastAPI, Header, HTTPException, Path, Query, Response
from fastapi.responses import JSONResponse, PlainTextResponse
from pydantic import BaseModel, Field, field_validator

import cidr_collector as cc
import db
import formatters
from cidr_collector import load_full_config
from storage import StorageError, load_json

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

TOKEN_ENV = "RIPE_API_TOKEN"
MAX_CURSOR_DIGITS = 30  # курсор длиннее не бывает: такое значение - ошибка запроса

app = FastAPI(title="RIPE CIDR/FQDN API")


class AddressType(str, Enum):
    cidr = "cidr"
    fqdn = "fqdn"
    all_types = "all"


class OutputFormat(str, Enum):
    json = "json"
    nftables = "nftables"
    mikrotik = "mikrotik"
    bird = "bird"
    frr = "frr"


class IPVersion(str, Enum):
    all_versions = "all"
    v4 = "4"
    v6 = "6"


class CollectRequest(BaseModel):
    type: Literal["asn", "fqdn", "all"] = "all"


class ScheduleUpdate(BaseModel):
    type: Literal["asn", "fqdn", "backup"]
    cron: str

    model_config = {"json_schema_extra": {"example": {"type": "asn", "cron": "*/10 * * * *"}}}


ASN_MIN, ASN_MAX = 1, 4294967295
_FQDN_LABEL = re.compile(r"^[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?$")


def normalize_fqdn(value: str) -> str:
    """Нижний регистр, без завершающей точки; проверка длины и меток, IP-литералы отклоняются."""
    fqdn = value.strip().lower().rstrip(".")
    if not fqdn or len(fqdn) > 253:
        raise ValueError("FQDN must be 1-253 characters")
    try:
        ipaddress.ip_address(fqdn)
    except ValueError:
        pass
    else:
        raise ValueError("IP addresses are not allowed, use a domain name")
    if not all(_FQDN_LABEL.match(label) for label in fqdn.split(".")):
        raise ValueError("Invalid FQDN: labels must be 1-63 chars of [a-z0-9-], no leading/trailing hyphen")
    return fqdn


class ASNBody(BaseModel):
    asn: int = Field(ge=ASN_MIN, le=ASN_MAX)


class FQDNBody(BaseModel):
    fqdn: str

    @field_validator("fqdn")
    @classmethod
    def _normalize(cls, value):
        return normalize_fqdn(value)


@app.exception_handler(StorageError)
@app.exception_handler(sqlite3.Error)
def storage_error_handler(request, exc):
    return JSONResponse(status_code=503, content={"detail": "Data storage unavailable"})


def verify_token(x_api_key: Optional[str] = Header(None)):
    expected = os.environ.get(TOKEN_ENV)
    if not expected:
        # Fail closed: без заданного токена управление отключено
        raise HTTPException(status_code=503, detail=f"{TOKEN_ENV} is not configured; write access disabled.")
    # Сравнение по байтам: compare_digest для str принимает только ASCII и падает на других символах
    if not x_api_key or not secrets.compare_digest(x_api_key.encode(), expected.encode()):
        raise HTTPException(status_code=401, detail="Invalid or missing X-API-Key.")


def kinds_of(address_type):
    """Типы данных в базе для значения параметра type."""
    return {"asn" if t == AddressType.cidr else "fqdn"
            for t in (AddressType.cidr, AddressType.fqdn) if address_type in (t, AddressType.all_types)}


def ensure_data_ready(conn, kinds):
    """503, пока база пересоздана после порчи и данные не собраны заново: пустой список ввёл бы потребителя в заблуждение."""
    if db.recreated_pending(conn, kinds):
        raise HTTPException(status_code=503, headers={"Retry-After": "300"},
                            detail="The database was recreated after corruption; data is being collected again.")


@app.get("/addresses", response_model=None, responses={
    200: {"description": "JSON list for format=json, text/plain configuration script for other formats",
          "content": {"application/json": {"schema": {"type": "array", "items": {"type": "string"}}},
                      "text/plain": {"schema": {"type": "string"}}}}})
def get_addresses(
    type: AddressType = Query(AddressType.all_types, description="Filter by address type"),
    format: OutputFormat = Query(OutputFormat.json, description="Output format"),
    ip_version: IPVersion = Query(IPVersion.all_versions, description="Filter by IP version"),
    aggregate: bool = Query(False, description="Collapse overlapping/adjacent prefixes"),
    name: str = Query("ripe", pattern=r"^[A-Za-z][A-Za-z0-9_]{0,31}$",
                      description="List/set name used in generated configuration"),
    response: Response = None,
):
    # Одна сессия и один снимок: курсор в заголовке точно соответствует выданным данным
    kinds = kinds_of(type)
    results = set()
    with db.session() as conn, db.read_transaction(conn):
        ensure_data_ready(conn, kinds)
        headers = {"X-Changes-Cursor": str(db.journal_head(conn))}
        for kind in kinds:
            results.update(db.get_values(conn, kind))

    output = formatters.build_output(results, format.value, ip_version.value, aggregate, name)
    if format == OutputFormat.json:
        response.headers.update(headers)
        return output
    return PlainTextResponse(output, headers=headers)


def parse_since(raw: str):
    """since: целый курсор или время ISO 8601 (без пояса - UTC). Возвращает (курсор, время UTC в формате журнала)."""
    invalid = HTTPException(status_code=400, detail="since must be a cursor (integer) or an ISO 8601 time")
    # Только ASCII-цифры и ограниченная длина: str.isdigit() принимает юникод-цифры, а int() ограничен по длине
    if raw.isascii() and raw.isdigit():
        if len(raw) > MAX_CURSOR_DIGITS:
            raise invalid
        return int(raw), None
    try:
        # «+» в адресной строке приходит пробелом; «Z» понимаем явно
        moment = datetime.fromisoformat(raw.strip().replace(" ", "+").replace("Z", "+00:00"))
        moment = moment.replace(tzinfo=timezone.utc) if moment.tzinfo is None else moment.astimezone(timezone.utc)
    except (ValueError, OverflowError):  # OverflowError: крайние даты с часовым поясом
        raise invalid
    return None, moment.strftime("%Y-%m-%dT%H:%M:%S.") + f"{moment.microsecond // 1000:03d}Z"


@app.get("/addresses/diff", responses={
    400: {"description": "Invalid since"},
    410: {"description": "since is older than the change journal: fetch the full /addresses list"}})
def get_addresses_diff(
    since: str = Query(..., description="Cursor from the previous response (recommended) or ISO 8601 time (UTC)"),
    type: AddressType = Query(AddressType.all_types, description="Filter by address type"),
    ip_version: IPVersion = Query(IPVersion.all_versions, description="Filter by IP version"),
):
    cursor, since_ts = parse_since(since)
    kinds = kinds_of(type)
    with db.session() as conn:
        ensure_data_ready(conn, kinds)
        changes = db.get_changes(conn, kinds, cursor, since_ts)
    if changes is None:
        raise HTTPException(status_code=410, detail="since is outside the change journal; fetch the full /addresses list")
    added, removed, head = changes
    return {
        "since": since,
        "now": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "cursor": head,
        "added": formatters.select_strings(added, ip_version.value),
        "removed": formatters.select_strings(removed, ip_version.value),
    }


def collector_state():
    """Состояние демона по status.json (задания + heartbeat): (жив ли, время heartbeat, задания)."""
    try:
        status = load_json(cc.STATUS_FILE, None)
    except StorageError:
        status = None

    alive, updated_at = False, None
    if status:
        updated_at = status.get("updated_at")
        try:
            age = (datetime.now() - datetime.fromisoformat(updated_at)).total_seconds()
            alive = age < cc.STATUS_STALE_AFTER
        except (TypeError, ValueError):
            pass
    return alive, updated_at, (status.get("jobs", {}) if status else {})


def source_health():
    """Сводка по источникам для /health: всего и только неисправные (сбоев подряд не меньше порога)."""
    try:
        config = load_full_config()
    except StorageError:
        return None, False
    threshold = cc.failure_threshold(config)
    with db.session() as conn:
        statuses = db.source_statuses(conn)
    summary, any_failing = {}, False
    for kind, key in (("asn", "asns"), ("fqdn", "fqdns")):
        configured = [str(source) for source in config.get(key, [])]
        failing = [{"source": source, "failures": statuses[kind, source]["failures"],
                    "last_success": statuses[kind, source]["last_success"],
                    "error_kind": statuses[kind, source]["error_kind"]}
                   for source in configured
                   if (kind, source) in statuses and statuses[kind, source]["failures"] >= threshold]
        summary[kind] = {"total": len(configured), "failing": failing}
        any_failing = any_failing or bool(failing)
    return summary, any_failing


@app.get("/health")
def health():
    """Состояние сборщика: демон пишет status.json (задания + heartbeat), API только читает."""
    alive, updated_at, jobs = collector_state()

    with db.session() as conn:
        counts = {"cidrs": db.count_values(conn, "asn"), "fqdn_ips": db.count_values(conn, "fqdn")}
        pending = db.recreated_pending(conn, ("asn", "fqdn"))

    try:
        restore = load_json(cc.RESTORE_FILE, None)  # след автовосстановления базы из копии
    except StorageError:
        restore = None
    # Наружу - только время и имя файла копии; каталоги и текст ошибки остаются в last_restore.json
    last_restore = ({"at": restore.get("at"), "backup": os.path.basename(str(restore.get("backup") or ""))}
                    if isinstance(restore, dict) else None)

    try:
        recreated = load_json(cc.RECREATED_FILE, None)  # база пересоздана после порчи без копий
    except StorageError:
        recreated = None

    sources, any_failing = source_health()
    healthy = alive and not pending and not any_failing and not any(j.get("last_error") for j in jobs.values())
    return {
        "status": "ok" if healthy else "degraded",
        "collector_alive": alive,
        "collector_updated_at": updated_at,
        "jobs": jobs,
        "counts": counts,
        "sources": sources,
        "last_restore": last_restore,
        # Только время и признак ожидания (пути карантина наружу не отдаём)
        "db_recreated": {"at": recreated.get("at"), "pending": pending} if recreated else None,
    }


@app.get("/schedule")
def get_schedule():
    return load_full_config().get("schedule", {})


@app.post("/schedule", dependencies=[Depends(verify_token)])
def update_schedule(schedule_update: ScheduleUpdate):
    # Validate cron string by attempting to create trigger
    try:
        CronTrigger.from_crontab(schedule_update.cron)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid cron string: {e}")

    # Update config file (read-modify-write под блокировкой, запись атомарная)
    def mutate(config):
        config.setdefault("schedule", {})[schedule_update.type] = schedule_update.cron
    cc.update_config(mutate)

    # Демон подхватит новое расписание из config.json при ближайшей сверке
    logger.info("Schedule updated: %s -> %s", schedule_update.type, schedule_update.cron)

    return {"message": "Schedule updated", "type": schedule_update.type, "cron": schedule_update.cron,
            "applied_within_seconds": cc.SYNC_INTERVAL}


def _purge(kind, source):
    with db.session() as conn:
        return db.purge_source(conn, kind, source)


@app.get("/sources")
def list_sources():
    """Состояние всех настроенных источников: сколько адресов дал, когда опрашивался, сколько сбоев подряд."""
    config = load_full_config()
    with db.session() as conn:
        statuses, counts = db.source_statuses(conn), db.address_counts(conn)
    empty = {"last_attempt": None, "last_success": None, "failures": 0, "error_kind": None}
    return {"failure_threshold": cc.failure_threshold(config), "sources": [
        {"kind": kind, "source": str(source), "addresses": counts.get((kind, str(source)), 0),
         **statuses.get((kind, str(source)), empty)}
        for kind, key in (("asn", "asns"), ("fqdn", "fqdns")) for source in config.get(key, [])]}


@app.get("/asns")
def list_asns():
    return {"asns": load_full_config().get("asns", [])}


@app.post("/asns", dependencies=[Depends(verify_token)])
def add_asn(body: ASNBody, response: Response):
    added, _ = cc.add_to_config_list("asns", body.asn)
    response.status_code = 201 if added else 200
    logger.info("ASN %s %s", body.asn, "added" if added else "already present")
    return {"asn": body.asn, "added": added}


@app.delete("/asns/{asn}", dependencies=[Depends(verify_token)])
def remove_asn(asn: int = Path(ge=ASN_MIN, le=ASN_MAX),
               purge: bool = Query(False, description="Also delete collected prefixes immediately")):
    removed, _ = cc.remove_from_config_list("asns", asn)
    purged = _purge("asn", str(asn)) if purge else False
    if not removed and not purged:
        raise HTTPException(status_code=404, detail=f"AS{asn} not found")
    logger.info("ASN %s removed (purged=%s)", asn, purged)
    return {"asn": asn, "removed": removed, "purged": purged}


@app.get("/fqdns")
def list_fqdns():
    return {"fqdns": load_full_config().get("fqdns", [])}


@app.post("/fqdns", dependencies=[Depends(verify_token)])
def add_fqdn(body: FQDNBody, response: Response):
    added, _ = cc.add_to_config_list("fqdns", body.fqdn)
    response.status_code = 201 if added else 200
    logger.info("FQDN %s %s", body.fqdn, "added" if added else "already present")
    return {"fqdn": body.fqdn, "added": added}


@app.delete("/fqdns/{fqdn}", dependencies=[Depends(verify_token)])
def remove_fqdn(fqdn: str, purge: bool = Query(False, description="Also delete collected IPs immediately")):
    try:
        fqdn = normalize_fqdn(fqdn)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    removed, _ = cc.remove_from_config_list("fqdns", fqdn)
    purged = _purge("fqdn", fqdn) if purge else False
    if not removed and not purged:
        raise HTTPException(status_code=404, detail=f"{fqdn} not found")
    logger.info("FQDN %s removed (purged=%s)", fqdn, purged)
    return {"fqdn": fqdn, "removed": removed, "purged": purged}


@app.post("/collect", status_code=202, dependencies=[Depends(verify_token)])
def request_collect(body: Optional[CollectRequest] = None):
    """Немедленный сбор: запрос передаётся демону (он проверяет его каждые несколько секунд)."""
    alive, _, _ = collector_state()
    if not alive:
        raise HTTPException(status_code=503, detail="Collector daemon is not running; start it first.")
    kind = body.type if body else "all"
    requested = cc.request_collection(cc.COLLECT_TYPES if kind == "all" else [kind])
    logger.info("Manual collection requested: %s", requested)
    return {"requested": requested, "message": "Collection requested; follow progress in GET /health",
            "picked_up_within_seconds": cc.TRIGGER_POLL_INTERVAL}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)