discord-geodata/tools/sync.py

#!/usr/bin/env python3
"""Синхронизирует data/ с апстримными rule-set из tools/sources.json.

Каждый источник кладётся в свой файл — data/<target>.txt для доменов и
data/ip/<target>.txt для сетей, — а сводные категории (data/discord.txt,
data/ip/discord.txt) пересобираются как список include-ов. Рукописные
списки (data/discord-base.txt) остаются нетронутыми: sync перезаписывает
только файлы со своей шапкой GENERATED.

    python3 tools/sync.py                 # синхронизировать всё
    python3 tools/sync.py --only legiz    # только один источник
    python3 tools/sync.py --strict        # падать, если хоть один источник недоступен

Источники в формате .srs скомпилированы: для них нужен бинарь sing-box
(`rule-set decompile`). Если его нет — источник пропускается, а ранее
синхронизированный файл остаётся на месте: устаревшие данные лучше пустых.
"""

from __future__ import annotations

import argparse
import ipaddress
import json
import shutil
import subprocess
import sys
import tempfile
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path

import build
from build import BuildError, Net, Options, Rule

TOOLS = Path(__file__).resolve().parent
USER_AGENT = "discord-geodata-sync (+https://github.com/)"
TIMEOUT = 60

GENERATED_HEADER = f"# {build.GENERATED_MARK} tools/sync.py — не редактируйте вручную."


class SyncError(Exception):
    """Источник не удалось получить или разобрать."""


@dataclass
class Parsed:
    domains: list[Rule] = field(default_factory=list)
    groups: dict[Options, list[Net]] = field(default_factory=dict)
    skipped: list[str] = field(default_factory=list)


# --------------------------------------------------------------------------- #
# Получение
# --------------------------------------------------------------------------- #


def fetch(url: str) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(request, timeout=TIMEOUT) as response:
            return response.read()
    except (urllib.error.URLError, OSError, TimeoutError) as exc:
        raise SyncError(f"не удалось скачать {url}: {exc}") from exc


def decompile_srs(blob: bytes, singbox: str | None) -> dict:
    """.srs — это zlib-контейнер sing-box; читать его умеет только сам sing-box."""
    if not singbox:
        raise SyncError("формат .srs требует бинаря sing-box (--sing-box или в PATH)")
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / "rule-set.srs"
        dst = Path(tmp) / "rule-set.json"
        src.write_bytes(blob)
        result = subprocess.run(
            [singbox, "rule-set", "decompile", "--output", str(dst), str(src)],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            raise SyncError(f"sing-box rule-set decompile: {result.stderr.strip() or result.stdout.strip()}")
        return json.loads(dst.read_text(encoding="utf-8"))


def load_source(source: dict, singbox: str | None) -> dict:
    url = source.get("json_url") or source["url"]
    blob = fetch(url)
    if url.endswith(".srs") or blob[:3] == b"SRS":
        return decompile_srs(blob, singbox)
    try:
        return json.loads(blob.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise SyncError(f"{url}: не похоже на JSON rule-set: {exc}") from exc


# --------------------------------------------------------------------------- #
# Разбор sing-box rule-set
# --------------------------------------------------------------------------- #


def iter_default_rules(rules: list, where: str):
    """Разворачивает logical-правила: вложенные условия нас интересуют как список."""
    for rule in rules:
        if not isinstance(rule, dict):
            raise SyncError(f"{where}: правило не является объектом")
        if rule.get("type") == "logical":
            yield from iter_default_rules(rule.get("rules", []), where)
        else:
            yield rule


def listify(value: object, where: str, key: str) -> list:
    """В sing-box почти любое поле правила — «listable»: строка или массив строк."""
    if value is None:
        return []
    if isinstance(value, (str, int)):
        return [value]
    if isinstance(value, list):
        return value
    raise SyncError(f"{where}: поле '{key}' неожиданного типа {type(value).__name__}")


def rule_field(rule: dict, key: str, where: str) -> list:
    return listify(rule.get(key), where, key)


def rule_options(rule: dict, where: str) -> Options:
    networks = sorted({str(n).lower() for n in rule_field(rule, "network", where)})
    for network in networks:
        if network not in build.NETWORKS:
            raise SyncError(f"{where}: неизвестный network '{network}'")

    ports: list[tuple[int, int]] = []
    for port in rule_field(rule, "port", where):
        ports.append((int(port), int(port)))
    for span in rule_field(rule, "port_range", where):
        match = build.PORT_RE.match(str(span))
        if not match:
            raise SyncError(f"{where}: невалидный port_range '{span}'")
        start = int(match.group(1))
        ports.append((start, int(match.group(2) or start)))

    return Options(tuple(networks), build.merge_ports(ports))


def parse_ruleset(payload: dict, where: str) -> Parsed:
    if not isinstance(payload, dict) or "rules" not in payload:
        raise SyncError(f"{where}: нет ключа 'rules'")

    parsed = Parsed()
    domains: dict[Rule, None] = {}

    for index, rule in enumerate(iter_default_rules(payload["rules"], where)):
        spot = f"{where}#{index}"
        if rule.get("invert"):
            parsed.skipped.append(f"{spot}: invert-правило не выразимо в geosite/geoip")
            continue

        for key, template in (
            ("domain", "full:{}"),
            ("domain_suffix", "{}"),
            ("domain_keyword", "keyword:{}"),
            ("domain_regex", "regexp:{}"),
        ):
            for value in rule_field(rule, key, spot):
                # Ведущая точка в sing-box значит «только поддомены», а в geosite
                # суффикс всегда покрывает и сам домен — точку просто снимаем.
                token = template.format(str(value).lstrip(".") if key == "domain_suffix" else value)
                try:
                    domains.setdefault(build.parse_token(token, spot), None)
                except BuildError as exc:
                    parsed.skipped.append(str(exc))

        cidrs: list[Net] = []
        for value in rule_field(rule, "ip_cidr", spot):
            try:
                network = ipaddress.ip_network(str(value), strict=False)
            except ValueError as exc:
                parsed.skipped.append(f"{spot}: невалидная сеть '{value}': {exc}")
                continue
            cidrs.append(Net(network.network_address.packed, network.prefixlen))

        for key in ("source_ip_cidr", "source_port", "source_port_range"):
            if rule.get(key):
                parsed.skipped.append(f"{spot}: '{key}' игнорируется")

        if cidrs:
            parsed.groups.setdefault(rule_options(rule, spot), []).extend(cidrs)

    parsed.domains = sorted(domains, key=lambda r: r.sort_key)
    return parsed


# --------------------------------------------------------------------------- #
# Чистка
# --------------------------------------------------------------------------- #


# Чистка внутри файла — та же, что build применяет к развёрнутой категории:
# покрытые суффиксом домены и вложенные сети до .dat всё равно не доедут.
prune_domains = build.collapse_rules
prune_nets = build.collapse_nets


# --------------------------------------------------------------------------- #
# Запись
# --------------------------------------------------------------------------- #


def is_generated(path: Path) -> bool:
    if not path.is_file():
        return True
    head = path.read_text(encoding="utf-8").splitlines()[:3]
    return any(build.GENERATED_MARK in line for line in head)


def write_generated(path: Path, notes: list[str], body: list[str]) -> None:
    if not is_generated(path):
        raise SyncError(f"{path} создан вручную — sync его не трогает")
    content = "\n".join([GENERATED_HEADER, *(f"# {n}" for n in notes), "", *body, ""])
    build.write(path, content)


def short_suffix(options: Options) -> str:
    return "-".join(options.networks) or "port"


def long_suffix(options: Options) -> str:
    """Различает наборы, у которых совпал протокол, но разошлись порты."""
    ports = "-".join(f"{a}" if a == b else f"{a}-{b}" for a, b in options.ports)
    return "-".join(part for part in (short_suffix(options), ports) if part)


def group_names(groups: list[Options], target: str) -> dict[Options, str]:
    """Имя категории на группу: короткое, а при совпадении — с портами."""
    taken = [short_suffix(o) for o in groups]
    names = {}
    for options in groups:
        suffix = short_suffix(options)
        if taken.count(suffix) > 1:
            suffix = long_suffix(options)
        names[options] = f"{target}-{suffix}"
    if len(set(names.values())) != len(names):
        raise SyncError(f"{target}: не удалось развести имена категорий по ограничениям")
    return names


def sync_source(source: dict, singbox: str | None) -> tuple[list[Path], list[str]]:
    """Возвращает записанные файлы и предупреждения."""
    target = source["target"]
    note = source.get("note", source["name"])
    url = source.get("json_url") or source["url"]

    parsed = parse_ruleset(load_source(source, singbox), source["name"])

    written: list[Path] = []
    domains = prune_domains(parsed.domains)
    if domains:
        write_generated(
            build.DATA_DIR / f"{target}.txt",
            [f"Источник: {note}", f"URL: {url}", f"Правил: {len(domains)}"],
            [str(rule) for rule in domains],
        )
        written.append(build.DATA_DIR / f"{target}.txt")

    ordered = sorted(parsed.groups, key=lambda o: (o.networks, o.ports))
    names = group_names([o for o in ordered if o.restricted], target)
    for options in ordered:
        nets = prune_nets(parsed.groups[options])
        if not nets:
            continue
        name = names[options] if options.restricted else target

        directives = [f"!network: {n}" for n in options.networks]
        if options.ports:
            directives.append(f"!port: {options.xray_ports}")
        path = build.DATA_DIR / build.IP_SUBDIR / f"{name}.txt"
        write_generated(
            path,
            [f"Источник: {note}", f"URL: {url}", f"Сетей: {len(nets)}", f"Ограничения: {options}"],
            [*directives, *([""] if directives else []), *(str(net) for net in nets)],
        )
        written.append(path)

    return written, parsed.skipped


def drop_stale(target: str, written: set[Path]) -> list[Path]:
    """Убирает файлы источника, которые в этот раз не появились (апстрим сузился)."""
    removed = []
    for directory in (build.DATA_DIR, build.DATA_DIR / build.IP_SUBDIR):
        for path in sorted(directory.glob(f"{target}*.txt")):
            stem = path.stem
            if stem != target and not stem.startswith(f"{target}-"):
                continue
            if path in written or not is_generated(path):
                continue
            path.unlink()
            removed.append(path)
    return removed


def write_aggregates(aggregate: str) -> None:
    """Сводные категории — просто список include-ов по тому, что есть на диске."""
    for directory, restrict in ((build.DATA_DIR, False), (build.DATA_DIR / build.IP_SUBDIR, True)):
        names = []
        for path in sorted(directory.glob("*.txt")):
            name = path.stem
            if name == aggregate:
                continue
            if restrict:
                # В сводную попадают только безусловные списки: подмешать к ним
                # «udp + порты 50000-50099» значило бы снять это ограничение с
                # чужих сетей (там половина Cloudflare) — такие категории живут
                # отдельными тегами.
                _, options = build.resolve_ip(name)
                if options.restricted:
                    continue
            names.append(name)
        if not names:
            continue
        write_generated(
            directory / f"{aggregate}.txt",
            ["Сводная категория: рукописный список плюс всё, что притянул sync."],
            [f"include:{name}" for name in names],
        )


# --------------------------------------------------------------------------- #
# Точка входа
# --------------------------------------------------------------------------- #


def main(argv: list[str] | None = None) -> int:
    for stream in (sys.stdout, sys.stderr):
        stream.reconfigure(encoding="utf-8", errors="replace")  # type: ignore[union-attr]

    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--sources", type=Path, default=TOOLS / "sources.json")
    parser.add_argument("--only", action="append", default=[], help="синхронизировать только эти источники")
    parser.add_argument("--sing-box", dest="singbox", help="путь к бинарю sing-box для .srs")
    parser.add_argument("--strict", action="store_true", help="ненулевой код возврата при любой ошибке")
    args = parser.parse_args(argv)

    config = json.loads(args.sources.read_text(encoding="utf-8"))
    singbox = args.singbox or shutil.which("sing-box")
    sources = [s for s in config["sources"] if not args.only or s["name"] in args.only]
    if not sources:
        print("error: подходящих источников нет", file=sys.stderr)
        return 1

    failed = []
    for source in sources:
        try:
            written, skipped = sync_source(source, singbox)
            removed = drop_stale(source["target"], set(written))
            files = ", ".join(p.relative_to(build.ROOT).as_posix() for p in written) or "нет данных"
            print(f"{source['name']}: {files}")
            for path in removed:
                print(f"  удалён устаревший {path.relative_to(build.ROOT).as_posix()}")
            for warning in skipped[:10]:
                print(f"  пропущено — {warning}")
            if len(skipped) > 10:
                print(f"  … и ещё {len(skipped) - 10} пропущенных записей")
        except (SyncError, BuildError) as exc:
            failed.append(source["name"])
            print(f"{source['name']}: ОШИБКА — {exc}", file=sys.stderr)
            print("  оставлены прежние данные", file=sys.stderr)

    try:
        write_aggregates(config["aggregate"])
    except (SyncError, BuildError) as exc:
        print(f"error: не удалось пересобрать сводные категории: {exc}", file=sys.stderr)
        return 1

    if failed:
        print(f"источников с ошибкой: {len(failed)} из {len(sources)} ({', '.join(failed)})", file=sys.stderr)
        if args.strict or len(failed) == len(sources):
            return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())