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

import argparse
import json
import os
import sys

from .config import (
    build_output_settings,
    get_source_type,
    print_output_settings,
    to_float_or_none,
    validate_config,
)
from .filter import apply_filter, collect_excluded_domains
from .notify import run_notify
from .scoring import (
    build_top_candidates,
    choose_domain,
    parse_scored_records,
    rank_scored_records,
)
from .sources import extract_records, fetch_api_json, load_cfst_rows, parse_domains
from .storage import (
    read_json_file,
    utc_now_iso,
    write_json_file,
    write_text_file,
)


def run_updater(config_path_abs, print_settings_only=False, result_json_file=""):
    """Execute main rotation workflow with strict error handling and state persistence."""
    if not os.path.exists(config_path_abs):
        err_msg = {"status": "error", "error": f"config file not found: {config_path_abs}"}
        print(json.dumps(err_msg, ensure_ascii=True), file=sys.stderr)
        return 1

    cfg = read_json_file(config_path_abs)

    try:
        validate_config(cfg)
    except Exception as e:
        err_msg = {"status": "error", "error": f"invalid config: {e}"}
        print(json.dumps(err_msg, ensure_ascii=True), file=sys.stderr)
        return 1

    if print_settings_only:
        print_output_settings(config_path_abs, cfg)
        return 0

    output_cfg = cfg.get("output", {})
    output_settings = build_output_settings(output_cfg, config_path_abs)
    notify_cfg = cfg.get("notify", {})
    selected_text_file = output_settings["selected_text_path"]
    selected_json_file = output_settings["selected_json_path"]
    state_file = output_settings["state_path"]
    vars_file = output_settings["vars_path"]
    selected_json_key = output_settings["selected_json_key"]
    state_last_good_key = output_settings["state_last_good_key"]
    vars_value_key = output_settings["vars_value_key"]

    state = read_json_file(state_file, default={})
    last_good = state.get(state_last_good_key, "")
    source_type = get_source_type(cfg)

    try:
        top_n = int(cfg.get("selection", {}).get("top_n", 3))

        if source_type == "cfst_local":
            cfst_rows = load_cfst_rows(cfg, config_path_abs)
            parsed = [row["domain"] for row in cfst_rows]
            filtered = apply_filter(parsed, cfg.get("domain_filter", {}))
            filtered_set = set(filtered)
            cfst_rows = [row for row in cfst_rows if row["domain"] in filtered_set]
            if not cfst_rows:
                raise RuntimeError("No valid IP available from cfst result after filtering")

            # 若未显式禁用下载测速（未配置 -dd），优先过滤掉测速失败的 0.00 下载速度记录
            cfst_cfg = cfg.get("cfst_local", {})
            disable_download = "-dd" in [str(a).strip() for a in cfst_cfg.get("run_args", [])]
            valid_speed_rows = [
                row for row in cfst_rows
                if to_float_or_none(row.get("download_speed")) is not None
                and float(to_float_or_none(row.get("download_speed"))) > 0.0
            ]
            if not disable_download and not valid_speed_rows and last_good:
                selected = None
                candidate_domains = []
                top_candidates = []
            else:
                eligible_rows = valid_speed_rows if (not disable_download and valid_speed_rows) else cfst_rows
                selected = eligible_rows[0]["domain"]
                candidate_domains = [row["domain"] for row in eligible_rows[:top_n]]
                top_candidates = build_top_candidates("cfst_local", candidate_domains, cfst_rows=eligible_rows)
        else:
            payload = fetch_api_json(cfg)
            parsed = parse_domains(payload, cfg.get("parser", {}))
            filtered = apply_filter(parsed, cfg.get("domain_filter", {}))

            record_mapping_cfg = cfg.get("record_mapping", {})
            field_map = record_mapping_cfg.get("field_map", {})
            records = extract_records(payload, record_mapping_cfg)

            record_filter_cfg = cfg.get("record_filter", {})
            blocked_domains = collect_excluded_domains(records, field_map, record_filter_cfg)
            if blocked_domains:
                filtered = [d for d in filtered if d not in blocked_domains]

            scoring_cfg = cfg.get("scoring", {})
            scored_records = parse_scored_records(records, field_map, record_mapping_cfg, scoring_cfg)
            filtered_set = set(filtered)
            scored_records = [r for r in scored_records if r["domain"] in filtered_set]
            ranked_scored = rank_scored_records(scored_records, scoring_cfg)
            switch_margin_pct = to_float_or_none(scoring_cfg.get("switch_margin_pct")) or 0.0
            selected, candidate_domains = choose_domain(
                filtered, top_n, ranked_scored, last_good=last_good, switch_margin_pct=switch_margin_pct
            )
            top_candidates = build_top_candidates(
                "api",
                candidate_domains,
                records=records,
                field_map=field_map,
                ranked_scored=ranked_scored,
            )

        status = "ok"
        if not selected and last_good:
            # 复检 last_good: 确保已被拉黑的故障节点不会借回退复活
            if source_type == "api" and apply_filter([last_good], cfg.get("domain_filter", {})) != [last_good]:
                last_good = ""
            else:
                selected = last_good
                status = "fallback_last_good"
        if not selected:
            if source_type == "cfst_local":
                raise RuntimeError("No valid IP available from cfst and no fallback in state")
            raise RuntimeError("No valid domain available from API and no fallback in state")

        write_text_file(selected_text_file, selected + "\n")

        current_json = {
            selected_json_key: selected,
            "updated_at": utc_now_iso(),
            "status": status,
            "source_type": source_type,
            "source_count": len(parsed),
            "top_candidates": top_candidates,
        }
        write_json_file(selected_json_file, current_json)
        write_json_file(
            vars_file,
            {
                vars_value_key: selected,
                "UPDATED_AT": current_json["updated_at"],
                "STATUS": status,
            },
        )

        new_state = {
            "updated_at": current_json["updated_at"],
            state_last_good_key: selected,
            "status": status,
            "source_count": len(parsed),
            "source_type": source_type,
        }
        write_json_file(state_file, new_state)

        run_notify(notify_cfg.get("command", ""), selected, status)
        if result_json_file:
            write_json_file(result_json_file, current_json)
        print(json.dumps(current_json, ensure_ascii=True))
        return 0

    except Exception as e:
        now = utc_now_iso()
        err_state = {
            "updated_at": now,
            "status": "error",
            "error": str(e),
            state_last_good_key: last_good,
            "source_type": source_type,
        }
        write_json_file(state_file, err_state)

        if last_good and source_type == "api":
            if apply_filter([last_good], cfg.get("domain_filter", {})) != [last_good]:
                last_good = ""

        if last_good:
            write_text_file(selected_text_file, last_good + "\n")
            write_json_file(
                selected_json_file,
                {
                    selected_json_key: last_good,
                    "updated_at": now,
                    "status": "error_use_last_good",
                    "error": str(e),
                    "source_type": source_type,
                },
            )
            write_json_file(
                vars_file,
                {
                    vars_value_key: last_good,
                    "UPDATED_AT": now,
                    "STATUS": "error_use_last_good",
                },
            )
            fallback_res = {
                selected_json_key: last_good,
                "status": "error_use_last_good",
                "error": str(e),
            }
            run_notify(notify_cfg.get("command", ""), last_good, "error_use_last_good")
            if result_json_file:
                write_json_file(result_json_file, fallback_res)
            print(json.dumps(fallback_res, ensure_ascii=True))
            return 0

        err_res = {"status": "error", "error": str(e)}
        run_notify(notify_cfg.get("command", ""), "", "error")
        if result_json_file:
            write_json_file(result_json_file, err_res)
        print(json.dumps(err_res, ensure_ascii=True), file=sys.stderr)
        return 1


def main():
    ap = argparse.ArgumentParser(description="Auto select preferred endpoint value")
    ap.add_argument("--config", default="configs/config.server.json", help="Path to config JSON")
    ap.add_argument(
        "--print-output-settings",
        action="store_true",
        help="Print resolved output settings as JSON and exit",
    )
    ap.add_argument(
        "--result-json-file",
        default="",
        help="Optional path to output machine-readable JSON result (decouples IPC from stdout)",
    )
    args = ap.parse_args()

    config_path_abs = os.path.abspath(args.config)
    if not os.path.exists(config_path_abs):
        candidate_configs = os.path.abspath(os.path.join("configs", args.config))
        if os.path.exists(candidate_configs):
            config_path_abs = candidate_configs
        else:
            repo_configs = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "configs", os.path.basename(args.config)))
            if os.path.exists(repo_configs):
                config_path_abs = repo_configs

    rc = run_updater(
        config_path_abs=config_path_abs,
        print_settings_only=args.print_output_settings,
        result_json_file=args.result_json_file,
    )
    if rc != 0:
        sys.exit(rc)


if __name__ == "__main__":
    main()