#!/usr/bin/env python3
"""把 Cloudflare 官方 IP 段自动同步到 轻量应用服务器(Lighthouse) 防火墙规则。

子命令：
  --print                                    打印当前 CF IPv4 段
  --check --instance-id <id>                 只对比，有差异退出码 2
  --apply --instance-id <id>                 做差量收敛（默认 dry-run）
  --apply --instance-id <id> --no-dry-run    真正执行（cron 用）

依赖（一次性）：ECS 装 aliyun CLI + swas-open 插件 + aliyun configure
  （RAM 子账号 AliyunSWASFullAccess，地域 ap-northeast-2）。
"""

import argparse
import json
import subprocess
import sys
import time
import urllib.request

CF_API = "https://api.cloudflare.com/client/v4/ips"
CF_FALLBACK = "https://www.cloudflare.com/ips-v4"
UA = {"User-Agent": "servicic-cf-ip-sync/1.0"}
DEFAULT_REGION = "ap-northeast-2"
REMARK = "cf-ip-sync"
CHUNK = 10  # 每次 create/delete 调用携带的规则数上限（保守）


def log(msg):
    print(msg)


def fetch_cf_ips():
    """拉取 CF IPv4 段；主 API 失败时回退到静态页。"""
    for url in (CF_API, CF_FALLBACK):
        try:
            req = urllib.request.Request(url, headers=UA)
            with urllib.request.urlopen(req, timeout=30) as r:
                data = json.loads(r.read().decode())
                return set(data["result"]["ipv4_cidrs"])
        except Exception as e:
            log(f"fetch failed ({url}): {e}")
    sys.exit(1)


def run_aliyun(args):
    """执行 aliyun CLI；失败即退出。打印实际命令方便排查/复核。"""
    cmd = ["aliyun"] + args
    log("$ " + " ".join(cmd))
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        sys.stderr.write(p.stderr)
        sys.exit(1)
    if not p.stdout.strip():
        return None
    return json.loads(p.stdout)


def list_rules(instance_id, region):
    """列出实例全部防火墙规则，返回 [{port, source, remark, rule_id}]。

    字段以实际 API 响应为准：RuleId（注意不是 FirewallRuleId）。
    若一条规则都解析不到，dump 原始响应前 500 字符到 stderr 供自诊断。
    """
    data = run_aliyun([
        "swas-open", "list-firewall-rules",
        "--instance-id", instance_id,
        "--biz-region-id", region,
        "--page-size", "100",
    ]) or {}

    rules = []

    def walk(x):
        if isinstance(x, dict):
            rid = x.get("RuleId") or x.get("FirewallRuleId")
            if rid:
                rules.append({
                    "rule_id": str(rid),
                    "port": str(x.get("Port") or ""),
                    "source": str(x.get("SourceCidrIp") or ""),
                    "remark": str(x.get("Remark") or ""),
                })
            for v in x.values():
                walk(v)
        elif isinstance(x, list):
            for item in x:
                walk(item)

    walk(data)

    if not rules and data:
        sys.stderr.write("DEBUG: no firewall rules parsed; raw response head:\n"
                         + json.dumps(data)[:500] + "\n")
    return rules


def main():
    ap = argparse.ArgumentParser(description="Sync Cloudflare IP ranges into Lighthouse firewall rules")
    ap.add_argument("--print", action="store_true", help="print current CF IPv4 CIDRs")
    ap.add_argument("--check", action="store_true", help="compare only; exit 2 if drift")
    ap.add_argument("--apply", action="store_true", help="apply drift (dry-run unless --no-dry-run)")
    ap.add_argument("--no-dry-run", action="store_true", help="really apply")
    ap.add_argument("--instance-id", default=None, help="轻量应用服务器实例ID")
    ap.add_argument("--ports", default="80,8080", help="受管端口，逗号分隔")
    ap.add_argument("--region", default=DEFAULT_REGION)
    args = ap.parse_args()

    want = fetch_cf_ips()

    if args.print:
        for c in sorted(want):
            print(c)
        return

    if not args.instance_id:
        sys.stderr.write("--instance-id is required (except --print)\n")
        sys.exit(1)

    ports = [p.strip() for p in args.ports.split(",") if p.strip()]
    rules = list_rules(args.instance_id, args.region)

    # 只关心我们标记过的规则（其余一概不动）
    by_key = {(r["port"], r["source"]): r["rule_id"]
              for r in rules
              if r["remark"] == REMARK and r["port"] in ports}

    desired = {(p, c) for p in ports for c in want}
    adds = sorted(desired - set(by_key))
    removes = [by_key[k] for k in sorted(set(by_key) - desired)]

    if not adds and not removes:
        log("no drift")
        return

    log(f"drift: +{len(adds)} add, -{len(removes)} remove")
    if args.check:
        sys.exit(2)

    def create_args(chunk):
        out = ["swas-open", "create-firewall-rules",
               "--instance-id", args.instance_id, "--biz-region-id", args.region]
        for port, cidr in chunk:
            out += ["--firewall-rules", f"Port={port}", f"Remark={REMARK}",
                    "RuleProtocol=TCP", f"SourceCidrIp={cidr}"]
        return out

    def delete_args(chunk):
        return ["swas-open", "delete-firewall-rules",
                "--instance-id", args.instance_id, "--biz-region-id", args.region,
                "--rule-ids"] + chunk

    if not args.no_dry_run:
        log("dry-run：加 --no-dry-run 才会真正执行，实际命令如下：")
        for i in range(0, len(adds), CHUNK):
            log("$ " + " ".join(create_args(adds[i:i + CHUNK])))
        for i in range(0, len(removes), CHUNK):
            log("$ " + " ".join(delete_args(removes[i:i + CHUNK])))
        return

    for i in range(0, len(adds), CHUNK):
        run_aliyun(create_args(adds[i:i + CHUNK]))
    for i in range(0, len(removes), CHUNK):
        run_aliyun(delete_args(removes[i:i + CHUNK]))

    # 复查收敛
    time.sleep(2)
    rules2 = list_rules(args.instance_id, args.region)
    managed2 = {(r["port"], r["source"]) for r in rules2
                if r["remark"] == REMARK and r["port"] in ports}
    missing = sorted(desired - managed2)
    log("done" if not missing else f"warning: still missing {len(missing)} entries")
    if missing:
        sys.exit(3)


if __name__ == "__main__":
    main()