ripe-cidr-collector/tests/test_db.py

import datetime
import glob
import json
import os
import sqlite3
import sys

import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import cidr_collector as cc
import db


@pytest.fixture
def files(tmp_path, monkeypatch):
    for attr, name in (("DB_FILE", "ripe.db"), ("DATA_FILE", "data.json"), ("FQDN_DATA_FILE", "fqdn_data.json")):
        monkeypatch.setattr(cc, attr, str(tmp_path / name))
    monkeypatch.setattr(cc, "BACKUP_DIR", str(tmp_path / "backups"))
    monkeypatch.setattr(cc, "RESTORE_FILE", str(tmp_path / "last_restore.json"))
    monkeypatch.setattr(cc, "RECREATED_FILE", str(tmp_path / "db_recreated.json"))
    monkeypatch.setattr(cc, "CONFIG_FILE", str(tmp_path / "config.json"))
    return tmp_path


def test_legacy_json_import(files):
    before = datetime.datetime.now().isoformat(timespec="seconds")
    (files / "data.json").write_text(json.dumps({
        # Новый формат: first_seen/last_seen сохраняются
        "1": {"last_updated": "2026-02-01T00:00:00", "prefixes": ["1.0.0.0/24"],
              "seen": {"1.0.0.0/24": {"first_seen": "2026-01-01T00:00:00", "last_seen": "2026-02-01T00:00:00"}}},
        # Старый формат без seen: first_seen = last_updated, last_seen = момент миграции
        "2": {"last_updated": "2025-12-01T00:00:00", "prefixes": ["2.0.0.0/24"]}}))
    (files / "fqdn_data.json").write_text(json.dumps(
        {"example.com": {"last_updated": "2025-12-01T00:00:00", "ips": ["9.9.9.9"]}}))

    with db.session() as conn:
        rows = {(k, s, v): (f, l) for k, s, v, f, l in conn.execute("SELECT * FROM addresses")}
    assert rows[("asn", "1", "1.0.0.0/24")] == ("2026-01-01T00:00:00", "2026-02-01T00:00:00")
    first, last = rows[("asn", "2", "2.0.0.0/24")]
    assert first == "2025-12-01T00:00:00" and last >= before
    assert ("fqdn", "example.com", "9.9.9.9") in rows

    # Оригиналы переименованы (резервная копия), повторное открытие ничего не импортирует
    assert not (files / "data.json").exists() and len(list(files.glob("data.json.migrated-*"))) == 1
    (files / "data.json").write_text(json.dumps({"3": {"prefixes": ["3.0.0.0/24"]}}))
    with db.session() as conn:
        assert db.count_values(conn) == 3


def test_change_journal(files):
    now = datetime.datetime.now()
    (files / "data.json").write_text(json.dumps({"9": {"prefixes": ["9.0.0.0/24"]}}))
    with db.session() as conn:
        # Импорт из JSON в журнал не попадает
        assert db.get_changes(conn, {"asn"}, cursor=0) == (set(), set(), 0)

        with db.transaction(conn):
            db.merge_source(conn, "asn", "1", {"a", "b"}, now, 90)
            db.merge_source(conn, "asn", "2", {"b"}, now, 90)
        added, removed, c1 = db.get_changes(conn, {"asn"}, cursor=0)
        assert (added, removed) == ({"a", "b"}, set())

        # Значение b держит второй источник: исчезает только a
        with db.transaction(conn):
            db.purge_source(conn, "asn", "1")
        assert db.get_changes(conn, {"asn"}, cursor=c1)[:2] == (set(), {"a"})

        # Удалено и возвращено в одном интервале - изменений нет; повтор с новым курсором пуст
        c2 = db.get_changes(conn, {"asn"}, cursor=c1)[2]
        with db.transaction(conn):
            db.purge_source(conn, "asn", "2")
            db.merge_source(conn, "asn", "2", {"b"}, now, 90)
        added, removed, c3 = db.get_changes(conn, {"asn"}, cursor=c2)
        assert (added, removed) == (set(), set()) and c3 > c2
        assert db.get_changes(conn, {"fqdn"}, cursor=0)[:2] == (set(), set())  # фильтр по типу

        # Очистка сдвигает горизонт: старый курсор недействителен, текущий - нет; курсор «из будущего» тоже
        with db.transaction(conn):
            assert db.prune_changes(conn, now + datetime.timedelta(days=2), 1) > 0
        assert db.get_changes(conn, {"asn"}, cursor=0) is None
        assert db.get_changes(conn, {"asn"}, cursor=c3 + 1) is None
        assert db.get_changes(conn, {"asn"}, cursor=c3) == (set(), set(), c3)

        # Пакет больше порции проверки наличия (PRESENCE_BATCH): результат полный и точный
        head = db.journal_head(conn)
        bulk = {f"10.{i // 250}.{i % 250}.0/24" for i in range(2 * db.PRESENCE_BATCH + 200)}
        with db.transaction(conn):
            db.merge_source(conn, "asn", "3", bulk, now, 90)
        assert db.get_changes(conn, {"asn"}, cursor=head)[:2] == (bulk, set())
        head = db.journal_head(conn)
        with db.transaction(conn):
            db.purge_source(conn, "asn", "3")
        assert db.get_changes(conn, {"asn"}, cursor=head)[:2] == (set(), bulk)


def test_backup_rotation_and_integrity(files):
    now = datetime.datetime.now()
    with db.session() as conn:
        with db.transaction(conn):
            db.merge_source(conn, "asn", "1", {"a"}, now, 90)
        for minutes in range(3):
            db.backup_database(conn, now + datetime.timedelta(minutes=minutes), keep=2)

    # Ротация оставляет две самые свежие копии; копия самодостаточна и содержит данные
    backups = db.list_backups(cc.BACKUP_DIR)
    assert len(backups) == 2 and not glob.glob(os.path.join(cc.BACKUP_DIR, "*.tmp"))
    assert backups[0].endswith(f"{(now + datetime.timedelta(minutes=2)).astimezone(datetime.timezone.utc):%Y%m%dT%H%M%SZ}.db")
    assert sqlite3.connect(backups[0]).execute("SELECT COUNT(*) FROM addresses").fetchone()[0] == 1


def test_restore_from_backup(files):
    now = datetime.datetime.now()
    with db.session() as conn:
        with db.transaction(conn):
            db.merge_source(conn, "asn", "1", {"a", "b"}, now, 90)
        old_cursor = db.get_changes(conn, {"asn"}, cursor=0)[2]
        db.backup_database(conn, now, keep=3)

    # «Свежая» копия с мусором пропускается; живая база портится
    (files / "backups" / "ripe-99999999T999999Z.db").write_bytes(b"garbage" * 100)
    (files / "ripe.db").write_bytes(b"garbage" * 1000)

    with db.session() as conn:
        assert sorted(db.get_values(conn, "asn")) == ["a", "b"]
        # Курсор, выданный до отката, недействителен (клиент сделает полную выгрузку); новый работает
        assert db.get_changes(conn, {"asn"}, cursor=old_cursor) is None
        head = db.journal_head(conn)
        assert db.get_changes(conn, {"asn"}, cursor=head) == (set(), set(), head)
        with db.transaction(conn):
            db.merge_source(conn, "asn", "1", {"a", "b", "c"}, now, 90)
        assert db.get_changes(conn, {"asn"}, cursor=head)[:2] == ({"c"}, set())
    assert len(glob.glob(str(files / "ripe.db.corrupt-*"))) == 1
    assert json.loads((files / "last_restore.json").read_text())["backup"].endswith(".db")

    # Без исправных копий создаётся новая пустая база (с меткой, см. test_recreated_database_guard)
    for backup in db.list_backups(cc.BACKUP_DIR):
        os.unlink(backup)
    (files / "ripe.db").write_bytes(b"garbage" * 1000)
    with db.session() as conn:
        assert db.get_values(conn) == []
    assert os.path.exists(cc.RECREATED_FILE) and glob.glob(str(files / "ripe.db.corrupt-*"))


def test_recreated_database_guard(files):
    (files / "config.json").write_text('{"asns": [1], "fqdns": []}')
    (files / "ripe.db").write_bytes(b"garbage" * 1000)  # порча, копий нет

    with db.session() as conn:
        assert os.path.exists(cc.RECREATED_FILE)
        # ASN настроен, данных нет - ожидание; FQDN не настроен - пустой список законен
        assert db.recreated_pending(conn, ("asn",)) and not db.recreated_pending(conn, ("fqdn",))
        assert db.get_changes(conn, {"asn"}, cursor=5) is None  # курсор старой базы недействителен

        # Данные собраны заново - сборщик снимает метку
        with db.transaction(conn):
            db.merge_source(conn, "asn", "1", {"a"}, datetime.datetime.now(), 90)
        db.settle_recreated(conn)
    assert not os.path.exists(cc.RECREATED_FILE)


def test_empty_backup_is_not_trusted(files):
    now = datetime.datetime.now()
    later = now + datetime.timedelta(minutes=1)
    (files / "config.json").write_text('{"asns": [], "fqdns": []}')
    with db.session() as conn:
        # Источников нет: пустая база копируется как раньше
        empty_copy = db.backup_database(conn, now, keep=5)
        assert empty_copy and db.list_backups(cc.BACKUP_DIR) == [empty_copy]

        # Источник настроен, данных ещё нет: копия не создаётся (None), существующие остаются
        (files / "config.json").write_text('{"asns": [1], "fqdns": []}')
        assert db.empty_despite_sources(conn) and db.backup_database(conn, later, keep=5) is None
        assert db.list_backups(cc.BACKUP_DIR) == [empty_copy]

    # Порча базы: восстановление из пустой копии не должно дать законный пустой список - ставится метка
    (files / "ripe.db").write_bytes(b"garbage" * 1000)
    with db.session() as conn:
        assert db.get_values(conn) == [] and os.path.exists(cc.RECREATED_FILE)
        assert db.recreated_pending(conn, ("asn",))
        assert json.loads((files / "db_recreated.json").read_text())["reason"].startswith("restored backup")

        # Данные собраны: метка снимается, копия с данными есть, её восстановление метку не ставит
        with db.transaction(conn):
            db.merge_source(conn, "asn", "1", {"a"}, later, 90)
        db.settle_recreated(conn)
        assert not os.path.exists(cc.RECREATED_FILE) and not db.empty_despite_sources(conn)
        good_copy = db.backup_database(conn, later, keep=5)
    assert good_copy
    for stale in glob.glob(str(files / "ripe.db*")):
        os.unlink(stale)
    (files / "ripe.db").write_bytes(b"garbage" * 1000)
    with db.session() as conn:
        assert db.get_values(conn) == ["a"] and not os.path.exists(cc.RECREATED_FILE)