#!/usr/bin/env python3
"""Identify device models via HTTPS for each IP in cpes.txt.

For Cambium ePMP devices: reads /js/cambium_sku.js to get the SKU number,
maps it to a model name. For Ubiquiti devices: tries API endpoints.

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

import argparse
import json
import os
import re
import ssl
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

# Cambium ePMP SKU → model name
CAMBIUM_SKU_MAP: dict[int, str] = {
    8: "ePMP Force 110",
    28: "ePMP Force 300",
    35: "ePMP Force 180",
    55: "ePMP Force 200",
}

_CTX = ssl.create_default_context()
_CTX.check_hostname = False
_CTX.verify_mode = ssl.CERT_NONE


def http_get(url: str, timeout: float = 8) -> str | None:
    try:
        req = urllib.request.Request(url)
        resp = urllib.request.urlopen(req, context=_CTX, timeout=timeout)
        return resp.read().decode(errors="replace")
    except Exception:
        return None


def identify_model(ip: str) -> str:
    # Try both HTTPS and HTTP
    for proto in ("https", "http"):
        body = http_get(f"{proto}://{ip}/js/cambium_sku.js")
        if body:
            break

    # 1. Try Cambium ePMP SKU file
    if body:
        m = re.search(r"window\.sku\s*=\s*(\d+)", body)
        if m:
            sku = int(m.group(1))
            fw_m = re.search(r"window\.cambiumFWVersion\s*=\s*['\"]([^'\"]+)", body)
            fw = f" (fw={fw_m.group(1)})" if fw_m else ""
            model = CAMBIUM_SKU_MAP.get(sku, f"Cambium SKU {sku}")
            return f"{model}{fw}"

    # 2. Try Ubiquiti airOS 8.x API
    body = http_get(f"https://{ip}/api/v1.0/status")
    if body:
        try:
            data = json.loads(body)
            host = data.get("host", {})
            model = host.get("deviceModel", "")
            fw = host.get("firmwareVersion", "")
            if model:
                return f"{model} (fw={fw})" if fw else model
        except json.JSONDecodeError:
            pass

    # 3. Try Ubiquiti status.cgi
    body = http_get(f"https://{ip}/status.cgi")
    if body:
        # Look for model in the XML/HTML
        for pattern in [
            r"<title>([^<]+)</title>",
            r"model.*?([A-Za-z0-9\-\. ]{5,50})",
        ]:
            m = re.search(pattern, body, re.IGNORECASE)
            if m and m.group(1).strip() not in ("ePMP", "Ubiquiti", "Login"):
                return m.group(1).strip()

    # 4. Fallback: check if it's Cambium at all
    body = http_get(f"https://{ip}/")
    if body:
        if "ePMP" in body:
            return "Cambium ePMP (unknown model)"
        if "Ubiquiti" in body or "airOS" in body or "airMAX" in body:
            return "Ubiquiti (unknown model)"
        # Try to extract title
        m = re.search(r"<title>([^<]+)</title>", body)
        if m:
            return f"Unknown ({m.group(1).strip()})"

    return "no HTTPS response"


def parse_cpes(path: str) -> list[dict[str, str]]:
    entries: list[dict[str, str]] = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            # Extract fields positionally:
            # MAC is first 17 chars (AA:BB:CC:DD:EE:FF)
            mac_m = re.match(r"([0-9A-Fa-f:]{17})", line)
            if not mac_m:
                continue
            mac = mac_m.group(1)
            # IP is last field
            ip_m = re.search(r"(\d+\.\d+\.\d+\.\d+)$", line)
            if not ip_m:
                continue
            ip = ip_m.group(1)
            # Account number: # followed by digits
            acct_m = re.search(r"#(\d+)", line)
            readable = acct_m.group(1) if acct_m else ""
            entries.append({
                "mac": mac,
                "readable": readable,
                "name": "",  # filled in below
                "ip": ip,
            })
    return entries


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input", nargs="?", default="cpes.txt")
    ap.add_argument("-w", "--workers", type=int, 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"Identifying {len(entries)} devices via HTTPS...", file=sys.stderr)

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

    # Output sorted by model
    entries.sort(key=lambda e: models[e["ip"]].lower())

    # Show grouped counts
    groups: dict[str, list] = {}
    for e in entries:
        m = models[e["ip"]]
        groups.setdefault(m, []).append(e)

    for model, devs in sorted(groups.items(), key=lambda x: x[0].lower()):
        print(f"\n## {model} ({len(devs)} devices)")
        for d in devs:
            print(f"  {d['mac']:17}  #{d['readable']:5} {d['name']:24}  {d['ip']}")


if __name__ == "__main__":
    main()