ripe-cidr-collector/cidr_collector.py

import requests
import datetime
import ipaddress
import os
import argparse
import logging
import socket
import sys

from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

import db
from storage import StorageError, file_lock, load_json, save_json_atomic

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# Каталог данных: по умолчанию рядом с кодом, в контейнере - том (RIPE_DATA_DIR=/data)
DATA_DIR = os.environ.get("RIPE_DATA_DIR", BASE_DIR)
CONFIG_FILE = os.path.join(DATA_DIR, "config.json")
DB_FILE = os.path.join(DATA_DIR, "ripe.db")  # собранные адреса (SQLite)
# Старые JSON-хранилища: используются только для однократного импорта в DB_FILE (см. db.py)
DATA_FILE = os.path.join(DATA_DIR, "data.json")
FQDN_DATA_FILE = os.path.join(DATA_DIR, "fqdn_data.json")
STATUS_FILE = os.path.join(DATA_DIR, "status.json")  # пишет collector_daemon, читает API (/health)
COLLECT_REQUEST_FILE = os.path.join(DATA_DIR, "collect_request.json")  # API -> демон: POST /collect
# Резервные копии базы (задание backup в демоне); по умолчанию на том же томе, каталог можно вынести
BACKUP_DIR = os.environ.get("RIPE_BACKUP_DIR", os.path.join(DATA_DIR, "backups"))
RESTORE_FILE = os.path.join(DATA_DIR, "last_restore.json")  # след автовосстановления базы (читает /health)
RECREATED_FILE = os.path.join(DATA_DIR, "db_recreated.json")  # база пересоздана после порчи без копий (см. db.recreated_pending)
BASE_URL = "https://stat.ripe.net/data/announced-prefixes/data.json"
DEFAULT_SOURCEAPP = "ripe-cidr-collector"  # RIPEstat просит указывать приложение (ключ ripestat_sourceapp в config.json)
RIPESTAT_TIMEOUT = 10  # секунд на одну попытку
RIPESTAT_RETRIES = 3   # повторов при сбоях соединения и кодах 429/5xx
MAX_RETRY_AFTER = 30   # потолок паузы по заголовку Retry-After, секунд: сборщик не должен зависать на часы

DEFAULT_TTL_DAYS = 90
DEFAULT_BACKUP_KEEP = 7  # сколько копий базы хранить
DEFAULT_CHANGES_RETENTION_DAYS = 30  # срок хранения журнала изменений для /addresses/diff
# Демон сверяет расписание с config.json и пишет heartbeat в status.json с этим периодом;
# API считает демон живым, пока heartbeat не старше STATUS_STALE_AFTER секунд
SYNC_INTERVAL = 30
STATUS_STALE_AFTER = 120
COLLECT_TYPES = ("asn", "fqdn")
TRIGGER_POLL_INTERVAL = 5  # как часто демон проверяет запросы на немедленный сбор, секунды

logger = logging.getLogger(__name__)


def load_full_config():
    return load_json(CONFIG_FILE, {"asns": [], "fqdns": []})


def update_config(mutator):
    """Read-modify-write config.json под блокировкой с атомарной записью.

    mutator(config) изменяет словарь на месте и возвращает результат, который отдаётся вызывающему.
    Единственный путь записи конфигурации для CLI и API.
    """
    with file_lock(CONFIG_FILE):
        config = load_full_config()
        result = mutator(config)
        save_json_atomic(CONFIG_FILE, config)
        return result


def add_to_config_list(key, value):
    """Добавляет значение в список конфигурации ("asns"/"fqdns"). Возвращает (добавлено, актуальный список)."""
    def mutate(config):
        items = config.setdefault(key, [])
        if value in items:
            return False, items
        items.append(value)
        return True, items
    return update_config(mutate)


def remove_from_config_list(key, value):
    """Убирает значение из списка конфигурации. Возвращает (удалено, актуальный список)."""
    def mutate(config):
        items = config.setdefault(key, [])
        if value not in items:
            return False, items
        items.remove(value)
        return True, items
    return update_config(mutate)


def request_collection(types):
    """API: ставит запрос на немедленный сбор. Повторные запросы объединяются. Возвращает список типов."""
    with file_lock(COLLECT_REQUEST_FILE):
        pending = set(load_json(COLLECT_REQUEST_FILE, {}).get("types", []))
        merged = sorted(pending | set(types))
        save_json_atomic(COLLECT_REQUEST_FILE, {"types": merged, "requested_at": datetime.datetime.now().isoformat()})
        return merged


def pop_collection_requests():
    """Демон: забирает и удаляет запрос на немедленный сбор. Возвращает список типов (возможно пустой)."""
    if not os.path.exists(COLLECT_REQUEST_FILE):
        return []
    with file_lock(COLLECT_REQUEST_FILE):
        types = load_json(COLLECT_REQUEST_FILE, {}).get("types", [])
        if os.path.exists(COLLECT_REQUEST_FILE):
            os.remove(COLLECT_REQUEST_FILE)
    return [t for t in types if t in COLLECT_TYPES]


class _BoundedRetry(Retry):
    """Retry, у которого пауза по Retry-After ограничена MAX_RETRY_AFTER."""

    def get_retry_after(self, response):
        value = super().get_retry_after(response)
        return None if value is None else min(value, MAX_RETRY_AFTER)


def make_session():
    """Сессия для RIPEstat: повторы с нарастающей паузой при сбоях соединения и кодах 429/500/502/503/504."""
    retry = _BoundedRetry(total=RIPESTAT_RETRIES, backoff_factor=1, status_forcelist=(429, 500, 502, 503, 504),
                          allowed_methods=("GET",))
    session = requests.Session()
    session.mount("https://", HTTPAdapter(max_retries=retry))
    session.mount("http://", HTTPAdapter(max_retries=retry))
    return session


def _is_global(address):
    """Публичный адрес (не loopback, не частный, не 0.0.0.0, не link-local и т. п.)."""
    try:
        return ipaddress.ip_address(address).is_global
    except ValueError:
        return False


class CIDRCollector:
    def __init__(self):
        self.config = load_full_config()
        self.asns = self.config.get("asns", [])
        self.ttl_days = self.config.get("ttl_days", DEFAULT_TTL_DAYS)
        self.sourceapp = self.config.get("ripestat_sourceapp", DEFAULT_SOURCEAPP)
        self.session = None  # создаётся при первом запросе, закрывается в конце сбора

    def add_asn(self, asn):
        added, self.asns = add_to_config_list("asns", asn)
        logger.info("ASN %s added." if added else "ASN %s already in list.", asn)

    def remove_asn(self, asn):
        removed, self.asns = remove_from_config_list("asns", asn)
        logger.info("ASN %s removed." if removed else "ASN %s not found in list.", asn)

    def list_asns(self):
        print("Current ASNs:", self.asns)

    def fetch_prefixes(self, asn):
        params = {'resource': f'AS{asn}', 'sourceapp': self.sourceapp}
        try:
            if self.session is None:
                self.session = make_session()
            response = self.session.get(BASE_URL, params=params, timeout=RIPESTAT_TIMEOUT)
            response.raise_for_status()
            data = response.json()

            prefixes = []
            if 'data' in data and 'prefixes' in data['data']:
                for item in data['data']['prefixes']:
                    if 'prefix' in item:
                        prefixes.append(item['prefix'])
            return prefixes
        except Exception as e:
            logger.error("Error fetching data for AS%s: %s", asn, e)
            return None

    def run_collection(self):
        logger.info("Starting ASN CIDR collection...")
        # Сетевые запросы выполняем до транзакции, чтобы держать блокировку записи минимально
        fetched = {}
        try:
            for asn in self.asns:
                logger.info("Processing AS%s...", asn)
                prefixes = self.fetch_prefixes(asn)
                if prefixes is not None:
                    fetched[str(asn)] = set(prefixes)
        finally:
            if self.session is not None:
                self.session.close()
                self.session = None

        now = datetime.datetime.now()
        with db.session() as conn, db.transaction(conn):
            # Конфиг перечитываем внутри транзакции: источник, удалённый во время сбора, не воскресает
            configured = {str(a) for a in load_full_config().get("asns", [])}
            for str_asn, prefixes in fetched.items():
                if str_asn not in configured:
                    continue
                added, removed = db.merge_source(conn, "asn", str_asn, prefixes, now, self.ttl_days)
                logger.info("AS%s: +%d / -%d prefixes", str_asn, len(added), len(removed))
            swept = db.sweep_unconfigured(conn, "asn", configured, now, self.ttl_days)
            if swept:
                logger.info("Expired %d prefixes of unconfigured ASNs", swept)
            db.prune_changes(conn, now, self.config.get("changes_retention_days", DEFAULT_CHANGES_RETENTION_DAYS))
        logger.info("CIDR data saved to %s", DB_FILE)
        with db.session() as conn:
            db.settle_recreated(conn)


class FQDNCollector:
    def __init__(self):
        self.config = load_full_config()
        self.fqdns = self.config.get("fqdns", [])
        self.ttl_days = self.config.get("ttl_days", DEFAULT_TTL_DAYS)
        # По умолчанию в список попадают только глобальные адреса; true - для внутренних имён
        self.allow_non_global_ips = bool(self.config.get("allow_non_global_ips", False))

    def add_fqdn(self, fqdn):
        added, self.fqdns = add_to_config_list("fqdns", fqdn)
        logger.info("FQDN %s added." if added else "FQDN %s already in list.", fqdn)

    def remove_fqdn(self, fqdn):
        removed, self.fqdns = remove_from_config_list("fqdns", fqdn)
        logger.info("FQDN %s removed." if removed else "FQDN %s not found in list.", fqdn)

    def list_fqdns(self):
        print("Current FQDNs:", self.fqdns)

    def resolve_fqdn(self, fqdn):
        try:
            # Family 0 - получаем и IPv4 (A), и IPv6 (AAAA)
            results = socket.getaddrinfo(fqdn, None)
            # result[4] - sockaddr, для IP-протоколов индекс 0 - строка с адресом; зона IPv6 (%eth0) не нужна
            addresses = {result[4][0].split("%")[0] for result in results}
        except socket.gaierror as e:
            logger.error("Error resolving %s: %s", fqdn, e)
            return []
        if self.allow_non_global_ips:
            return list(addresses)
        kept = {a for a in addresses if _is_global(a)}
        if kept != addresses:
            logger.warning("%s: ignoring non-global addresses: %s", fqdn, ", ".join(sorted(addresses - kept)))
        return list(kept)

    def run_collection(self):
        logger.info("Starting FQDN IP collection...")
        fetched = {}
        for fqdn in self.fqdns:
            logger.info("Processing %s...", fqdn)
            resolved_ips = self.resolve_fqdn(fqdn)
            if not resolved_ips:
                # Ошибка DNS: TTL не применяем, чтобы сбой не удалил данные
                logger.warning("No IPs resolved for %s", fqdn)
                continue
            fetched[fqdn] = set(resolved_ips)

        now = datetime.datetime.now()
        with db.session() as conn, db.transaction(conn):
            configured = set(load_full_config().get("fqdns", []))
            for fqdn, ips in fetched.items():
                if fqdn not in configured:
                    continue
                added, removed = db.merge_source(conn, "fqdn", fqdn, ips, now, self.ttl_days)
                logger.info("%s: +%d / -%d IPs", fqdn, len(added), len(removed))
            swept = db.sweep_unconfigured(conn, "fqdn", configured, now, self.ttl_days)
            if swept:
                logger.info("Expired %d IPs of unconfigured FQDNs", swept)
            db.prune_changes(conn, now, self.config.get("changes_retention_days", DEFAULT_CHANGES_RETENTION_DAYS))
        logger.info("FQDN data saved to %s", DB_FILE)
        with db.session() as conn:
            db.settle_recreated(conn)


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    parser = argparse.ArgumentParser(description="Collector for RIPE AS CIDRs and FQDN IPs")
    subparsers = parser.add_subparsers(dest="command")

    # Command: run (default)
    parser_run = subparsers.add_parser("run", help="Run the collection process")
    parser_run.add_argument("--mode", choices=["asn", "fqdn", "all"], default="all", help="Collection mode: asn, fqdn, or all (default)")

    # ASN Commands
    parser_add = subparsers.add_parser("add", help="Add an ASN")
    parser_add.add_argument("asn", type=int, help="ASN to add")

    parser_remove = subparsers.add_parser("remove", help="Remove an ASN")
    parser_remove.add_argument("asn", type=int, help="ASN to remove")

    # FQDN Commands
    parser_add_fqdn = subparsers.add_parser("add-fqdn", help="Add an FQDN")
    parser_add_fqdn.add_argument("fqdn", type=str, help="FQDN to add")

    parser_remove_fqdn = subparsers.add_parser("remove-fqdn", help="Remove an FQDN")
    parser_remove_fqdn.add_argument("fqdn", type=str, help="FQDN to remove")

    # Command: list
    subparsers.add_parser("list", help="List ASNs and FQDNs")

    args = parser.parse_args()

    try:
        asn_collector = CIDRCollector()
        fqdn_collector = FQDNCollector()

        if args.command == "add":
            asn_collector.add_asn(args.asn)
        elif args.command == "remove":
            asn_collector.remove_asn(args.asn)
        elif args.command == "add-fqdn":
            fqdn_collector.add_fqdn(args.fqdn)
        elif args.command == "remove-fqdn":
            fqdn_collector.remove_fqdn(args.fqdn)
        elif args.command == "list":
            asn_collector.list_asns()
            fqdn_collector.list_fqdns()
        elif args.command == "run":
            mode = args.mode
            if mode in ["asn", "all"]:
                asn_collector.run_collection()

            if mode == "all":
                print("-" * 20)

            if mode in ["fqdn", "all"]:
                fqdn_collector.run_collection()
        else:
            parser.print_help()
    except StorageError as e:
        logger.error("Aborted: %s", e)
        sys.exit(1)


if __name__ == "__main__":
    main()