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 List, 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"

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"]
    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.")
    if not x_api_key or not secrets.compare_digest(x_api_key, expected):
        raise HTTPException(status_code=401, detail="Invalid or missing X-API-Key.")


def get_cidrs() -> List[str]:
    with db.session() as conn:
        return db.get_values(conn, "asn")


def get_fqdn_ips() -> List[str]:
    with db.session() as conn:
        return db.get_values(conn, "fqdn")


@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,
):
    # Курсор читаем до данных: изменения между чтением курсора и данных повторятся в diff, что безвредно
    with db.session() as conn:
        headers = {"X-Changes-Cursor": str(db.journal_head(conn))}
    results = set()

    if type in [AddressType.cidr, AddressType.all_types]:
        results.update(get_cidrs())

    if type in [AddressType.fqdn, AddressType.all_types]:
        results.update(get_fqdn_ips())

    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 в формате журнала)."""
    if raw.isdigit():
        return int(raw), None
    try:
        # «+» в адресной строке приходит пробелом; «Z» понимаем явно
        moment = datetime.fromisoformat(raw.strip().replace(" ", "+").replace("Z", "+00:00"))
    except ValueError:
        raise HTTPException(status_code=400, detail="since must be a cursor (integer) or an ISO 8601 time")
    moment = moment.replace(tzinfo=timezone.utc) if moment.tzinfo is None else moment.astimezone(timezone.utc)
    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 = {"asn" if t == AddressType.cidr else "fqdn"
             for t in (AddressType.cidr, AddressType.fqdn) if type in (t, AddressType.all_types)}
    with db.session() as conn:
        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 {})


@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")}

    healthy = alive 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,
    }


@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("/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)