vmess-domain-rotator/scripts/rotator/sources.py

import csv
import json
import os
import re
import subprocess
import urllib.request

from .config import build_url, resolve_path
from .filter import DOMAIN_RE, IPV4_RE, normalize_domain, flatten_values, get_by_json_path, get_values_by_path


def fetch_api_json(cfg):
    api = cfg["api"]
    url = build_url(api["url"], api.get("params", {}))
    method = api.get("method", "GET").upper()
    headers = api.get("headers", {})
    timeout = int(api.get("timeout_sec", 10))
    body_obj = api.get("body")
    body = None
    if body_obj is not None:
        body = json.dumps(body_obj).encode("utf-8")
        headers = {**headers, "Content-Type": "application/json"}

    req = urllib.request.Request(url=url, data=body, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8", errors="replace")
    return json.loads(raw)


def parse_domains(payload, parser_cfg):
    domains = []

    for p in parser_cfg.get("field_paths", []):
        values = get_values_by_path(payload, p)
        domains.extend(flatten_values(values))

    for p in parser_cfg.get("json_paths", []):
        v = get_by_json_path(payload, p)
        if v is not None:
            domains.extend(flatten_values(v))

    if not domains:
        regex_s = parser_cfg.get("regex", r"[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}")
        text = json.dumps(payload, ensure_ascii=True)
        domains.extend(re.findall(regex_s, text))

    clean = []
    seen = set()
    for d in domains:
        d = str(d).strip().lower().rstrip(".")
        if (DOMAIN_RE.match(d) or IPV4_RE.match(d)) and d not in seen:
            seen.add(d)
            clean.append(d)
    return clean


def extract_records(payload, record_mapping):
    records_path = str(record_mapping.get("records_path", "")).strip()
    raw = get_values_by_path(payload, records_path)
    return [x for x in raw if isinstance(x, dict)]


def load_cfst_rows(cfg, config_path_abs):
    cfst_cfg = cfg.get("cfst_local", {})
    config_dir = os.path.dirname(config_path_abs)

    work_dir = resolve_path(config_dir, cfst_cfg.get("work_dir", "./cfst"))
    binary_path = resolve_path(work_dir, cfst_cfg.get("binary", "./cfst"))
    result_file = resolve_path(work_dir, cfst_cfg.get("result_file", "result.csv"))
    encoding = str(cfst_cfg.get("encoding", "utf-8")).strip() or "utf-8"
    skip_run = bool(cfst_cfg.get("skip_run", False))
    timeout_sec = int(cfst_cfg.get("run_timeout_sec", 600))

    run_args = cfst_cfg.get("run_args", ["-o", os.path.basename(result_file)])
    if not isinstance(run_args, list):
        raise ValueError("cfst_local.run_args must be an array")
    command = [binary_path] + [str(x) for x in run_args]

    if not skip_run:
        completed = subprocess.run(
            command,
            cwd=work_dir,
            check=False,
            capture_output=True,
            text=True,
            encoding=encoding,
            errors="replace",
            timeout=timeout_sec,
        )
        if completed.returncode != 0:
            stderr = (completed.stderr or "").strip()
            stdout = (completed.stdout or "").strip()
            details = stderr or stdout or f"exit code {completed.returncode}"
            raise RuntimeError(f"cfst run failed: {details}")

    if not os.path.exists(result_file):
        raise RuntimeError(f"cfst result file not found: {result_file}")

    with open(result_file, "r", encoding=encoding, errors="replace", newline="") as f:
        reader = csv.reader(f)
        rows = [row for row in reader if any(str(col).strip() for col in row)]

    header_rows = int(cfst_cfg.get("header_rows", 1))
    if len(rows) <= header_rows:
        raise RuntimeError("cfst result has no data rows")

    columns_cfg = cfst_cfg.get("columns", {})
    if not isinstance(columns_cfg, dict):
        raise ValueError("cfst_local.columns must be an object")

    def col_index(name, default_index):
        raw = columns_cfg.get(name, default_index)
        try:
            idx = int(raw)
        except Exception as exc:
            raise ValueError(f"cfst_local.columns.{name} must be an integer") from exc
        if idx < 0:
            raise ValueError(f"cfst_local.columns.{name} must be >= 0")
        return idx

    ip_idx = col_index("ip", 0)
    sent_idx = col_index("sent", 1)
    received_idx = col_index("received", 2)
    loss_idx = col_index("loss_rate", 3)
    latency_idx = col_index("avg_latency", 4)
    speed_idx = col_index("download_speed", 5)
    region_idx = col_index("region", 6)

    out = []
    for row in rows[header_rows:]:
        if ip_idx >= len(row):
            continue
        domain = normalize_domain(row[ip_idx])
        if not domain:
            continue
        out.append(
            {
                "domain": domain,
                "ip": domain,
                "sent": row[sent_idx].strip() if sent_idx < len(row) else "",
                "received": row[received_idx].strip() if received_idx < len(row) else "",
                "loss_rate": row[loss_idx].strip() if loss_idx < len(row) else "",
                "avg_latency": row[latency_idx].strip() if latency_idx < len(row) else "",
                "download_speed": row[speed_idx].strip() if speed_idx < len(row) else "",
                "region": row[region_idx].strip() if region_idx < len(row) else "",
            }
        )

    if not out:
        raise RuntimeError("cfst result parsed to zero valid rows")
    return out