#!/usr/bin/env python3
# edr-deploy.py — throttled Bitdefender->Datto EDR rollout via GuruRMM.
#
# For every ONLINE Windows RMM agent belonging to a GPS client with a vaulted
# EDR registration key, and whose hostname is not already in Datto EDR, push the
# official Install-EDR one-liner through the RMM command API. Deliberately slow:
# strictly sequential, poll each command to completion, fixed gap between
# machines, retry once on transient 5xx. Glaz-Tech has no key (stays BD) and
# non-Windows agents are skipped by design.
#
# Env: RMM, TOK. Files: .edr-agents-now.json (fresh `edr.py --json agents` dump).
# Usage: edr-deploy.py --dry | --live [--gap 20]
import json, os, sys, time, ssl, subprocess, urllib.request, re

sys.stdout.reconfigure(encoding="utf-8", errors="replace")
RMM = os.environ["RMM"]; TOK = os.environ["TOK"]
ctx = ssl.create_default_context()
DRY = "--live" not in sys.argv
GAP = int(sys.argv[sys.argv.index("--gap") + 1]) if "--gap" in sys.argv else 20
ROOT = "projects/gps-rmm-audit"
INSTANCE = "https://azcomp4587.infocyte.com"
ONELINER = ('[System.Net.ServicePointManager]::SecurityProtocol = '
            '[Enum]::ToObject([System.Net.SecurityProtocolType], 3072); '
            '(new-object Net.WebClient).DownloadString('
            '"https://raw.githubusercontent.com/Infocyte/PowershellTools/master/AgentDeployment/install_huntagent.ps1")'
            ' | iex; Install-EDR -URL "{u}" -RegKey {k}')

def api(path, body=None, tries=3):
    for i in range(tries):
        try:
            req = urllib.request.Request(
                RMM + path,
                data=json.dumps(body).encode() if body is not None else None,
                method="POST" if body is not None else "GET",
                headers={"Authorization": "Bearer " + TOK, "Content-Type": "application/json"})
            raw = urllib.request.urlopen(req, context=ctx, timeout=30).read().decode("utf-8", "replace")
            return json.loads("".join(c for c in raw if ord(c) >= 32 or c in "\t\n\r"))
        except Exception as e:
            if i == tries - 1: raise
            time.sleep(5)

def norm(s):
    s = re.sub(r"[^a-z0-9 ]", " ", s.lower())
    drop = {"llc", "inc", "corp", "corporate", "ltd", "co", "the", "and", "of", "office", "offices"}
    return frozenset(w for w in s.split() if w and w not in drop)

# --- keys from vault (decrypted in-memory, never written to disk) ---
raw = subprocess.run(["bash", ".claude/scripts/vault.sh", "get", "msp-tools/datto-edr-regkeys.sops.yaml"],
                     capture_output=True, text=True).stdout
keys = {}
in_creds = False
for line in raw.splitlines():
    if line.startswith("credentials:"): in_creds = True; continue
    if in_creds:
        m = re.match(r"\s+(.+?):\s*(\S+)\s*$", line)
        if m: keys[m.group(1)] = m.group(2)
        elif line and not line.startswith(" "): break
OVERRIDES = {  # vault name -> RMM client name (where token match is not enough)
    "Tedards": "Bill Tedards",
    "Altschuler, Janet": "Janet Altschuler",
    "Residential and Renovation Eng": "Residential and Renovation Engineering",
}

# --- current EDR coverage ---
edr = json.load(open(f"{ROOT}/.edr-agents-now.json"))
edr = edr if isinstance(edr, list) else edr.get("agents", edr.get("data", []))
covered = set((a.get("hostname") or a.get("name") or "").lower() for a in edr)

# --- RMM fleet ---
agents = api("/api/agents")
clients = sorted(set(a.get("client_name") or "" for a in agents))
key_by_rmm_client = {}
for vname, k in keys.items():
    target = OVERRIDES.get(vname)
    if target and target in clients:
        key_by_rmm_client[target] = k; continue
    nv = norm(vname)
    best = [c for c in clients if norm(c) and (nv <= norm(c) or norm(c) <= nv)]
    if len(best) == 1:
        key_by_rmm_client[best[0]] = k
    elif len(best) > 1:
        exact = [c for c in best if norm(c) == nv]
        if len(exact) == 1: key_by_rmm_client[exact[0]] = k
        else: print(f"[WARNING] ambiguous key mapping for vault '{vname}': {best} — skipped")
    # zero matches = client has no RMM org yet; fine.

targets = []
for a in agents:
    c = a.get("client_name") or ""
    if c not in key_by_rmm_client: continue
    if (a.get("status") or "").lower() != "online": continue
    osname = (a.get("os") or a.get("os_type") or "").lower()
    if osname and "windows" not in osname: continue
    h = (a.get("hostname") or "").lower()
    if not h or h in covered: continue
    targets.append((c, a["hostname"], a["id"]))
targets.sort()

print(f"EDR deploy plan: {len(targets)} machine(s) across "
      f"{len(set(t[0] for t in targets))} client(s); gap {GAP}s; {'DRY-RUN' if DRY else 'LIVE'}")
for c, h, aid in targets:
    print(f"  {c:45s} {h}")
if DRY or not targets:
    sys.exit(0)

results = []
for n, (c, h, aid) in enumerate(targets, 1):
    cmd = ONELINER.format(u=INSTANCE, k=key_by_rmm_client[c])
    print(f"[{n}/{len(targets)}] {c} / {h} ... ", end="", flush=True)
    try:
        r = api(f"/api/agents/{aid}/command",
                {"command_type": "powershell", "command": cmd, "timeout_seconds": 300, "elevated": True})
        cid = r.get("command_id") or r.get("id")
        status, out = "timeout", ""
        for _ in range(60):  # poll up to 5 min
            time.sleep(5)
            st = api(f"/api/commands/{cid}")
            status = st.get("status")
            if status not in ("running", "pending"):
                out = (st.get("stdout") or "")[-400:]
                break
        ok = status == "completed" and ("installed" in out.lower() or "success" in out.lower()
                                        or "already" in out.lower() or out.strip())
        print(status)
        results.append((c, h, status, out.replace("\n", " ")[:200]))
    except Exception as e:
        print(f"ERROR {e}")
        results.append((c, h, "dispatch-error", str(e)[:200]))
    time.sleep(GAP)

print("\n=== RESULTS ===")
for c, h, s, o in results:
    print(f"{s:16s} {c:40s} {h}\n    {o}")
done = sum(1 for r in results if r[2] == "completed")
print(f"\ncompleted {done}/{len(results)}")
with open(f"{ROOT}/edr-deploy-results.md", "a", encoding="utf-8") as f:
    f.write(f"\n## Run {time.strftime('%Y-%m-%d %H:%M')} — {done}/{len(results)} completed\n")
    for c, h, s, o in results:
        f.write(f"- {s} | {c} | {h} | {o}\n")