#!/usr/bin/env python3
# NOTE: публичные диапазоны обезличены для публикации (198.18.0.0/15 — RFC 2544 benchmarking, 198.19.0.0/16 — placeholder).
"""Full matrix: connect-check/resources URLs × all src-nat to-addresses on a MikroTik GW.

Sources = intersection of:
  - /ip firewall nat chain=srcnat to-addresses (single IPs + expanded ranges)
  - /ip address enabled locals (must exist for /tool fetch src-address=)

Runs from jump host with sshpass. Concurrent workers. Incremental TSV.
"""
from __future__ import annotations

import argparse
import concurrent.futures
import csv
import ipaddress
import os
import re
import subprocess
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path


def mt_ssh(host: str, port: int, user: str, password: str, cmd: str, timeout: int = 60) -> str:
    env = os.environ.copy()
    r = subprocess.run(
        [
            "sshpass",
            "-p",
            password,
            "ssh",
            "-n",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "PreferredAuthentications=password",
            "-o",
            "PubkeyAuthentication=no",
            "-o",
            f"ConnectTimeout={min(25, timeout)}",
            "-p",
            str(port),
            f"{user}@{host}",
            cmd,
        ],
        capture_output=True,
        text=True,
        timeout=timeout,
        env=env,
    )
    out = (r.stdout or "") + (r.stderr or "")
    return out.replace("\r", "")


def parse_local_ips(addr_raw: str, wan: str) -> set[str]:
    ips: set[str] = set()
    for m in re.finditer(r"(\d+\.\d+\.\d+\.\d+)/(\d+)", addr_raw):
        ip, pfx = m.group(1), int(m.group(2))
        if ip.startswith("198.18.") or ip == wan or ip.startswith("198.19."):
            if ip.startswith("198.18.209.") and ip.endswith(".1") and ip != wan:
                continue
            ips.add(ip)
    return ips


def expand_to_addresses(token: str) -> list[str]:
    token = token.strip().strip('"')
    if not token or token in ("0.0.0.0", "0.0.0.0/0"):
        return []
    if "-" in token and "/" not in token:
        a, b = token.split("-", 1)
        try:
            start = int(ipaddress.IPv4Address(a.strip()))
            end = int(ipaddress.IPv4Address(b.strip()))
            if end < start or end - start > 4096:
                return []
            return [str(ipaddress.IPv4Address(i)) for i in range(start, end + 1)]
        except Exception:
            return []
    if "/" in token:
        try:
            net = ipaddress.ip_network(token, strict=False)
            if net.num_addresses > 4096:
                return []
            return [str(h) for h in net.hosts()] or [str(net.network_address)]
        except Exception:
            return []
    try:
        ipaddress.IPv4Address(token)
        return [token]
    except Exception:
        return []


def parse_srcnat_tos(nat_raw: str) -> set[str]:
    """Parse RouterOS print for to-addresses=… on srcnat rules."""
    out: set[str] = set()
    # terse or detail: to-addresses=A or to-addresses=A-B
    for m in re.finditer(r"to-addresses=([0-9./\-]+)", nat_raw):
        for ip in expand_to_addresses(m.group(1)):
            if ip.startswith("198.18.") or ip.startswith("198.19."):
                out.add(ip)
    # column layout sometimes shows bare IPs in to-addresses column — also catch 198.18 lines
    for m in re.finditer(r"\b(198\.18\.\d+\.\d+(?:-\d+\.\d+\.\d+\.\d+)?)\b", nat_raw):
        for ip in expand_to_addresses(m.group(1)):
            out.add(ip)
    return out


def classify_fetch(res: str) -> str:
    low = res.lower()
    # as-value path (quiet for MikroTik script,error log)
    if "wrap_fail" in low:
        return "FAIL"
    if "status=finished" in low or "status: finished" in low:
        return "OK_TCP"
    if "status=failed" in low or "status: failed" in low:
        return "FAIL"
    if any(x in low for x in ("timeout", "interrupted", "connecting", "failure:")):
        return "FAIL_TIMEOUT"
    snippet = " ".join(res.split())[:80]
    return f"OTHER:{snippet}"


def probe_one(
    host: str,
    port: int,
    user: str,
    password: str,
    src: str,
    tag: str,
    url: str,
    duration: str,
) -> tuple[str, str, str, str]:
    # as-value + on-error: failed fetch must NOT emit
    # "script,error executing script from sshd failed" (bare /tool fetch does).
    safe = url.replace("\\", "\\\\").replace('"', '\\"')
    # Trailing :put "DONE" keeps sshd from logging script,error on WRAP_FAIL.
    cmd = (
        f':do {{ :local x [/tool fetch url="{safe}" src-address={src} '
        f'duration={duration} keep-result=no as-value]; '
        f':put ("STATUS=".$x->"status") }} on-error={{ :put "WRAP_FAIL" }}; :put "DONE"'
    )
    try:
        res = mt_ssh(host, port, user, password, cmd, timeout=max(20, int(duration.rstrip("s") or "3") + 15))
    except subprocess.TimeoutExpired:
        return src, tag, url, "FAIL_TIMEOUT"
    except Exception as e:
        return src, tag, url, f"OTHER:{type(e).__name__}"
    return src, tag, url, classify_fetch(res)


def load_urls(path: Path) -> list[tuple[str, str]]:
    rows = []
    for ln in path.read_text().splitlines():
        ln = ln.strip()
        if not ln or ln.startswith("#"):
            continue
        parts = ln.split("\t")
        if len(parts) < 2:
            continue
        rows.append((parts[0], parts[1]))
    return rows


def summarize(outdir: Path, gw: str) -> Path:
    results = outdir / "matrix.tsv"
    rows = []
    if results.is_file():
        for ln in results.read_text().splitlines()[1:]:
            p = ln.split("\t")
            if len(p) >= 4:
                rows.append({"src": p[0], "tag": p[1], "url": p[2], "result": p[3]})
    by_src = defaultdict(Counter)
    by_tag = defaultdict(Counter)
    for r in rows:
        ok = "OK" in r["result"]
        by_src[r["src"]]["ok" if ok else "fail"] += 1
        by_tag[r["tag"]]["ok" if ok else "fail"] += 1

    summary = outdir / "summary.txt"
    lines = [
        f"GW={gw}",
        f"probes={len(rows)}",
        f"srcs={len(by_src)}",
        f"tags={len(by_tag)}",
        "",
        "=== per-src fail rate (worst first) ===",
    ]
    ranked = sorted(
        by_src.items(),
        key=lambda kv: (kv[1]["fail"] / max(1, kv[1]["ok"] + kv[1]["fail"]), kv[1]["fail"]),
        reverse=True,
    )
    for src, c in ranked:
        total = c["ok"] + c["fail"]
        pct = 100.0 * c["fail"] / total if total else 0
        lines.append(f"{src}\tfail={c['fail']}/{total}\t{pct:.1f}%")
    lines.append("")
    lines.append("=== tags with any FAIL (top 80 by fail count) ===")
    tag_ranked = sorted(by_tag.items(), key=lambda kv: kv[1]["fail"], reverse=True)
    for tag, c in tag_ranked[:80]:
        if c["fail"] == 0:
            continue
        total = c["ok"] + c["fail"]
        lines.append(f"{tag}\tfail={c['fail']}/{total}")
    # matrix csv: tag × src OK/FAIL
    srcs = sorted(by_src)
    tags = sorted(by_tag)
    pivot = {(r["tag"], r["src"]): r["result"] for r in rows}
    with (outdir / "matrix_summary.csv").open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["tag"] + srcs)
        for tag in tags:
            w.writerow(
                [tag]
                + [
                    ("OK" if "OK" in pivot.get((tag, s), "") else "FAIL")
                    for s in srcs
                ]
            )
    summary.write_text("\n".join(lines) + "\n")
    return summary


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", required=True)
    ap.add_argument("--port", type=int, default=22)
    ap.add_argument("--user", default="user")
    ap.add_argument("--password", default=os.environ.get("MT_PASS", ""))
    ap.add_argument("--wan", default="")
    ap.add_argument("--urls", type=Path, required=True)
    ap.add_argument("--outdir", type=Path, required=True)
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--duration", default="2s")
    ap.add_argument("--gw-label", default="")
    ap.add_argument("--limit-srcs", type=int, default=0)
    ap.add_argument("--limit-urls", type=int, default=0)
    args = ap.parse_args()
    if not args.password:
        print("MT_PASS / --password required", file=sys.stderr)
        return 2
    wan = args.wan or args.host
    outdir = args.outdir
    outdir.mkdir(parents=True, exist_ok=True)

    print(f"[*] {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} dump addr+nat on {args.host}")
    addr_raw = mt_ssh(args.host, args.port, args.user, args.password, "/ip address print where disabled=no")
    (outdir / "addr_raw.txt").write_text(addr_raw)
    nat_raw = mt_ssh(
        args.host,
        args.port,
        args.user,
        args.password,
        "/ip firewall nat print detail where chain=srcnat and disabled=no",
        timeout=120,
    )
    (outdir / "nat_srcnat_raw.txt").write_text(nat_raw)

    local = parse_local_ips(addr_raw, wan)
    nat_tos = parse_srcnat_tos(nat_raw)
    # Prefer intersection; if nat parse empty, fall back to locals
    srcs = sorted(local & nat_tos) if nat_tos else sorted(local)
    if not srcs and local:
        srcs = sorted(local)
    # Always include WAN if local
    if wan in local and wan not in srcs:
        srcs.append(wan)
    if args.limit_srcs:
        srcs = srcs[: args.limit_srcs]
    (outdir / "srcs.txt").write_text("\n".join(srcs) + "\n")
    (outdir / "nat_tos.txt").write_text("\n".join(sorted(nat_tos)) + "\n")
    (outdir / "local_ips.txt").write_text("\n".join(sorted(local)) + "\n")
    print(f"[*] local={len(local)} nat_tos={len(nat_tos)} probe_srcs={len(srcs)}")

    urls = load_urls(args.urls)
    if args.limit_urls:
        urls = urls[: args.limit_urls]
    print(f"[*] urls={len(urls)} workers={args.workers} duration={args.duration}")
    print(f"[*] total probes={len(srcs) * len(urls)}")

    matrix = outdir / "matrix.tsv"
    done: set[tuple[str, str]] = set()
    if matrix.is_file():
        for ln in matrix.read_text().splitlines()[1:]:
            p = ln.split("\t")
            if len(p) >= 3:
                done.add((p[0], p[2]))
        print(f"[*] resume: already {len(done)} rows")
    else:
        matrix.write_text("src\ttag\turl\tresult\n")

    jobs = [(s, t, u) for s in srcs for t, u in urls if (s, u) not in done]
    t0 = time.time()
    n_done = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = [
            ex.submit(
                probe_one,
                args.host,
                args.port,
                args.user,
                args.password,
                s,
                t,
                u,
                args.duration,
            )
            for s, t, u in jobs
        ]
        with matrix.open("a") as fh:
            for fut in concurrent.futures.as_completed(futs):
                src, tag, url, result = fut.result()
                fh.write(f"{src}\t{tag}\t{url}\t{result}\n")
                fh.flush()
                n_done += 1
                if n_done % 50 == 0 or n_done == len(futs):
                    elapsed = time.time() - t0
                    rate = n_done / elapsed if elapsed else 0
                    eta = (len(futs) - n_done) / rate if rate else 0
                    print(
                        f"[.] {n_done}/{len(futs)} {rate:.2f}/s ETA {eta/60:.1f}m last={src} {result}",
                        flush=True,
                    )

    label = args.gw_label or args.host
    summarize(outdir, label)
    print(f"[+] done {outdir} elapsed={(time.time()-t0)/60:.1f}m")
    return 0


if __name__ == "__main__":
    sys.exit(main())