#!/usr/bin/env python3
"""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 threading
import time
from collections import Counter, defaultdict
from pathlib import Path
from urllib.parse import urlparse

_SRC_LOCKS: dict[str, threading.Lock] = {}
_SRC_LOCKS_GUARD = threading.Lock()


def _lock_for_src(src: str) -> threading.Lock:
    with _SRC_LOCKS_GUARD:
        lock = _SRC_LOCKS.get(src)
        if lock is None:
            lock = threading.Lock()
            _SRC_LOCKS[src] = lock
        return lock


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("203.0.113.") or ip == wan or ip.startswith("198.51.100."):
            if ip.startswith("203.0.113.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("203.0.113.") or ip.startswith("198.51.100."):
                out.add(ip)
    # column layout sometimes shows bare IPs in to-addresses column — also catch 203.0.113 lines
    for m in re.finditer(r"\b(203\.0\.113\.\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}"


_TCP_REPLY_STATES = {
    "established",
    "time-wait",
    "close",
    "fin-wait",
    "close-wait",
    "last-ack",
    "closing",
}


def classify_fetch_tcp(res: str) -> str:
    """Split fetch-complete vs TCP-up vs no-reply (H75).

    OK_FETCH = /tool fetch STATUS=finished (TCP implied).
    OK_TCP   = WRAP_FAIL/failed but conntrack shows reply or a post-handshake state.
    FAIL     = WRAP_FAIL and syn-sent / no new socket / repl=0.
    """
    low = res.lower()
    finished = "status=finished" in low or "status: finished" in low
    wrap = "wrap_fail" in low
    failed = "status=failed" in low or "status: failed" in low
    states = [s.lower() for s in re.findall(r"TCP_STATE=(\S+)", res, re.I)]
    repls = [int(x) for x in re.findall(r"TCP_REPL=(\d+)", res, re.I)]
    tcp_ok = bool(repls and max(repls) >= 1) or any(
        any(tok in st for tok in _TCP_REPLY_STATES) for st in states
    )
    if finished:
        return "OK_FETCH"
    if tcp_ok:
        return "OK_TCP"
    if wrap or failed:
        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 _fetch_cmd(src: str, url: str, duration: str, fetch_extra: str) -> str:
    safe = url.replace("\\", "\\\\").replace('"', '\\"')
    extra = f" {fetch_extra.strip()}" if fetch_extra.strip() else ""
    return (
        f':do {{ :local x [/tool fetch url="{safe}" src-address={src} '
        f'duration={duration} keep-result=no as-value{extra}]; '
        f':put ("STATUS=".$x->"status") }} on-error={{ :put "WRAP_FAIL" }}'
    )


def _url_host(url: str) -> str:
    host = urlparse(url).hostname or ""
    return host.strip("[]")


def _conntrack_cmd(src: str, url: str, duration: str, fetch_extra: str) -> str:
    """Fetch, then print GW-originated TCP sockets to this URL's host.

    Dest-filtered (literal or :resolve) so we do not scan the whole
    connection table. No ID snapshot: leftover time-wait to the same
    dest still proves TCP from this src recently. Per-src lock keeps
    concurrent URLs from mixing. src-address~IP is originator only
    (subscriber sessions have src=100.64.x).
    """
    fetch = _fetch_cmd(src, url, duration, fetch_extra)
    host = _url_host(url).replace("\\", "").replace('"', "")
    if not host:
        return fetch + '; :put "TCP_NONE"; :put "DONE"'
    try:
        ipaddress.IPv4Address(host)
        resolve = f':local dip "{host}"'
    except ValueError:
        safe_host = host.replace("'", "")
        resolve = (
            f':local dip ""; '
            f":do {{ :set dip [:resolve \"{safe_host}\"] }} "
            f"on-error={{ :put \"TCP_RESOLVE_FAIL\" }}"
        )
    return (
        f"{resolve}; {fetch}; "
        f":local seen 0; "
        f":if ([:len $dip] > 0) do={{ "
        f":foreach i in=[/ip firewall connection find where "
        f'src-address~"{src}" and dst-address~$dip and protocol=tcp] do={{ '
        f":set seen 1; "
        f':put ("TCP_DST=".[/ip firewall connection get $i dst-address]); '
        f':put ("TCP_STATE=".[/ip firewall connection get $i tcp-state]); '
        f':put ("TCP_REPL=".[/ip firewall connection get $i repl-packets]); '
        f':put ("TCP_ORIG=".[/ip firewall connection get $i orig-packets]) }} }}; '
        f':if ($seen = 0) do={{ :put "TCP_NONE" }}; :put "DONE"'
    )


def probe_one(
    host: str,
    port: int,
    user: str,
    password: str,
    src: str,
    tag: str,
    url: str,
    duration: str,
    fetch_extra: str = "",
    tcp_conntrack: bool = False,
) -> 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).
    # Trailing :put "DONE" keeps sshd from logging script,error on WRAP_FAIL.
    if tcp_conntrack:
        cmd = _conntrack_cmd(src, url, duration, fetch_extra)
        classify = classify_fetch_tcp
        extra_to = 20
    else:
        cmd = _fetch_cmd(src, url, duration, fetch_extra) + '; :put "DONE"'
        classify = classify_fetch
        extra_to = 15
    try:
        if tcp_conntrack:
            with _lock_for_src(src):
                res = mt_ssh(
                    host,
                    port,
                    user,
                    password,
                    cmd,
                    timeout=max(30, int(duration.rstrip("s") or "3") + extra_to),
                )
        else:
            res = mt_ssh(
                host,
                port,
                user,
                password,
                cmd,
                timeout=max(20, int(duration.rstrip("s") or "3") + extra_to),
            )
    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(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)
    by_class = Counter()
    for r in rows:
        cls = r["result"].split(":", 1)[0]
        by_class[cls] += 1
        ok = r["result"].startswith("OK")
        by_src[r["src"]]["ok" if ok else "fail"] += 1
        by_tag[r["tag"]]["ok" if ok else "fail"] += 1
        by_src[r["src"]][cls] += 1

    summary = outdir / "summary.txt"
    lines = [
        f"GW={gw}",
        f"probes={len(rows)}",
        f"srcs={len(by_src)}",
        f"tags={len(by_tag)}",
        f"classes={dict(by_class)}",
        "",
        "=== 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
        extra = ""
        if c.get("OK_FETCH") or c.get("OK_TCP"):
            extra = f"\tOK_FETCH={c.get('OK_FETCH', 0)}\tOK_TCP={c.get('OK_TCP', 0)}"
        lines.append(f"{src}\tfail={c['fail']}/{total}\t{pct:.1f}%{extra}")
    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=39999)
    ap.add_argument("--user", default="nf-ro")
    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(
        "--fetch-extra",
        default="",
        help='Extra /tool fetch args, e.g. check-certificate=no for HTTPS-to-IP',
    )
    ap.add_argument("--limit-srcs", type=int, default=0)
    ap.add_argument("--limit-urls", type=int, default=0)
    ap.add_argument(
        "--src",
        action="append",
        default=[],
        help="Restrict probe to this src-address (repeatable). Must be local on GW.",
    )
    ap.add_argument(
        "--tcp-conntrack",
        action="store_true",
        help=(
            "After fetch, classify via /ip firewall connection (H75): "
            "OK_FETCH / OK_TCP (reply despite WRAP_FAIL) / FAIL. "
            "Serializes probes per src."
        ),
    )
    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.src:
        want = set(args.src)
        missing = sorted(want - local)
        if missing:
            print(f"[!] --src not present/enabled on GW: {missing}", file=sys.stderr)
            return 3
        srcs = sorted(want)
    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")

    # Round-robin by src so --tcp-conntrack per-src locks actually parallelize
    # across workers (src-major order would serialize almost everything).
    jobs: list[tuple[str, str, str]] = []
    pending = {
        s: [(t, u) for t, u in urls if (s, u) not in done] for s in srcs
    }
    while any(pending.values()):
        for s in srcs:
            bucket = pending[s]
            if bucket:
                t, u = bucket.pop(0)
                jobs.append((s, t, u))
    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,
                args.fetch_extra,
                args.tcp_conntrack,
            )
            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())