#!/usr/bin/env python3
"""Annotate cpes.txt with SNMP-discovered device models.

Queries each IP via SNMPv1 (community kdyyJrT0Mm) for:
  - sysDescr / sysObjectID / product code / firmware version
  - Falls back to MAC OUI for devices that don't respond to SNMP

Output: clean table with MAC, device type, account info.

Usage:
  python3 scripts/annotate_cpes.py [cpes.txt] > annotated.txt
"""
from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

SNMP_COMMUNITY = "kdyyJrT0Mm"

# Ubiquiti product codes from 17713.21.1.1.2.0 → model name
_PRODUCT_MAP: dict[int, str] = {
    8: "NanoStation M5 / NS loco M5",
    35: "NanoBeam M5 / LiteBeam M5",
    55: "PowerBeam M5 / Loco M5",
}

# sysObjectID suffix → model hint
_SYSOID_MAP: dict[str, str] = {
    "17713.21.9.35": "AirOS 5.x (NanoBeam/LiteBeam)",
    "17713.21.9.55": "AirOS 5.x (PowerBeam/Loco)",
    "10002.1": "AirOS (legacy MIB)",
    "41112.1.4": "airOS 8.x",
    "netSnmpAgentOIDs.10": "NET-SNMP agent (likely AirOS)",
}


def parse_cpes(path: str) -> list[dict[str, str]]:
    entries: list[dict[str, str]] = []
    pending: dict[str, str] | None = None
    with open(path) as f:
        for line in f:
            line = line.rstrip()
            # Account detail line (indented, comes after the IP/MAC line)
            m_acct = re.search(r"account=#(\d+)", line)
            if m_acct and pending:
                pending["readable"] = m_acct.group(1)
                acct_m = re.search(r"name='([^']+)'", line)
                pending["name"] = acct_m.group(1) if acct_m else ""
                entries.append(pending)
                pending = None
                continue

            # IP/MAC line
            m = re.search(r"ip=(\S+)", line)
            if not m:
                continue
            ip = m.group(1)
            mac_m = re.search(r"mac=(\S+)", line)
            mac = mac_m.group(1) if mac_m else ""
            pending = {"ip": ip, "mac": mac, "name": "", "readable": ""}
    return entries


def snmp_get(ip: str, oid: str, timeout: float = 8) -> str | None:
    try:
        p = subprocess.run(
            ["snmpget", "-v1", "-c", SNMP_COMMUNITY, "-Oqv", ip, oid],
            capture_output=True, text=True, timeout=timeout,
        )
        if p.returncode == 0:
            return p.stdout.strip().strip('"')
    except Exception:
        pass
    return None


def query_device(entry: dict[str, str]) -> dict[str, str]:
    ip = entry["ip"]
    result = {**entry, "model": "unknown"}

    sys_oid = snmp_get(ip, "1.3.6.1.2.1.1.2.0")
    sys_descr = snmp_get(ip, "1.3.6.1.2.1.1.1.0")
    prod_code = snmp_get(ip, "1.3.6.1.4.1.17713.21.1.1.2.0")
    fw = snmp_get(ip, "1.3.6.1.4.1.17713.21.1.1.1.0")
    snmp_mac = snmp_get(ip, "1.3.6.1.4.1.17713.21.1.1.5.0")

    # Build model string from available info
    parts: list[str] = []

    # 1. Product code (most specific)
    if prod_code:
        try:
            code = int(prod_code)
            if code in _PRODUCT_MAP:
                parts.append(_PRODUCT_MAP[code])
            else:
                parts.append(f"product={code}")
        except ValueError:
            parts.append(f"product={prod_code}")

    # 2. sysObjectID hint
    if sys_oid and not parts:
        for key, hint in _SYSOID_MAP.items():
            if key in sys_oid:
                parts.append(hint)
                break
        else:
            parts.append(sys_oid.split("::")[-1] if "::" in sys_oid else sys_oid)

    # 3. sysDescr (sometimes contains model)
    if sys_descr and sys_descr not in ("vntx", "unknown"):
        if not parts or sys_descr != "vntx":
            if sys_descr not in ("vntx",):
                pass  # sysDescr is usually just "vntx" or the host name, not useful

    # 4. MAC OUI fallback
    if not parts:
        mac = entry["mac"]
        if mac:
            oui = mac[:8].upper()
            parts.append(f"OUI {oui}")
        else:
            parts.append("unknown")

    # 5. Firmware version
    if fw:
        parts.append(f"fw={fw}")

    # 6. SNMP MAC verification
    if snmp_mac:
        snmp_mac_norm = re.sub(r"[^0-9a-f]", "", snmp_mac.lower())
        entry_mac_norm = re.sub(r"[^0-9a-f]", "", entry["mac"].lower())
        if snmp_mac_norm != entry_mac_norm:
            parts.append(f"(SNMP MAC: {snmp_mac})")

    result["model"] = " | ".join(parts) if parts else "no SNMP response"
    if sys_oid:
        result["sys_oid"] = sys_oid.split("::")[-1] if "::" in sys_oid else sys_oid
    if prod_code:
        result["prod_code"] = prod_code
    if fw:
        result["fw"] = fw

    return result


def main():
    ap = argparse.ArgumentParser(description="Annotate cpes.txt with SNMP device models")
    ap.add_argument("input", nargs="?", default="cpes.txt",
                    help="path to cpes.txt (default: cpes.txt)")
    ap.add_argument("-w", "--workers", type=int, default=20,
                    help="parallel SNMP queries (default: 20)")
    args = ap.parse_args()

    if not os.path.exists(args.input):
        sys.exit(f"{args.input} not found")

    entries = parse_cpes(args.input)
    print(f"Querying {len(entries)} devices via SNMP...", file=sys.stderr)

    results: list[dict[str, str]] = []
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(query_device, e): e for e in entries}
        for i, f in enumerate(as_completed(futs), 1):
            results.append(f.result())
            if i % 20 == 0:
                print(f"  {i}/{len(entries)}", file=sys.stderr)
    print(f"  {len(entries)}/{len(entries)} done", file=sys.stderr)

    # Sort by account name
    results.sort(key=lambda r: r["name"].lower())

    # Output
    print(f"{'MAC':17}  {'Model':50}  {'Account':30}  IP")
    print("-" * 120)
    for r in results:
        print(f"{r['mac']:17}  {r['model']:50}  #{r['readable']:5} {r['name']:24}  {r['ip']}")


if __name__ == "__main__":
    main()