#!/usr/bin/env python3
"""SSH from the admin laptop to every machine on every campus.

  python3 tools/fleet.py ard      the command to paste into Apple Remote Desktop (Send UNIX Command, as root)
  python3 tools/fleet.py config   write ~/.ssh/config.d/waters from the machines that reported back

The ARD command, on each Mac: turns on Remote Login, lets the logged-in user's account take this
laptop's key (~/.ssh/id_ed25519_waters), and reports name, user and addresses to North Attleboro's
settings page (/api/fleet/hello). `config` turns that list into shortcuts: ssh na-tracks-mac.
"""
import json
import os
import re
import sys
import urllib.request

HUB = os.environ.get("FLEET_HUB", "http://10.1.1.1:8020")
KEY = os.path.expanduser("~/.ssh/id_ed25519_waters")
OUT = os.path.expanduser("~/.ssh/config.d/waters")


def pubkey():
    return open(KEY + ".pub").read().strip()


def ard():
    key = pubkey()
    return r"""U=$(stat -f%Su /dev/console); if [ "$U" = root ] || [ -z "$U" ]; then echo "SKIPPED $(scutil --get ComputerName): nobody is logged in"; exit 0; fi
H=$(dscl . -read "/Users/$U" NFSHomeDirectory | awk '{print $2}'); KEY='""" + key + r"""'
mkdir -p "$H/.ssh" && touch "$H/.ssh/authorized_keys" && (grep -qF "$KEY" "$H/.ssh/authorized_keys" || echo "$KEY" >> "$H/.ssh/authorized_keys") && chown -R "$U:staff" "$H/.ssh" && chmod 700 "$H/.ssh" && chmod 600 "$H/.ssh/authorized_keys"
nc -z -G 2 127.0.0.1 22 >/dev/null 2>&1 || systemsetup -setremotelogin on >/dev/null 2>&1 || launchctl load -w /System/Library/LaunchDaemons/ssh.plist >/dev/null 2>&1
dseditgroup -o read com.apple.access_ssh >/dev/null 2>&1 && dseditgroup -o edit -a "$U" -t user com.apple.access_ssh >/dev/null 2>&1
IPS=$(ifconfig | awk '/inet /{print $2}' | grep -v '^127\.' | sed 's/.*/"&"/' | paste -sd, -); N=$(scutil --get ComputerName | tr -d '"\\')
curl -s -m 5 -X POST -H 'Content-Type: application/json' -d "{\"user\":\"$U\",\"name\":\"$N\",\"hostname\":\"$(scutil --get LocalHostName).local\",\"os\":\"mac\",\"ips\":[$IPS]}" """ + HUB + r"""/api/fleet/hello >/dev/null
echo "OK $N user=$U ssh=$(nc -z -G 2 127.0.0.1 22 >/dev/null 2>&1 && echo on || echo OFF)"
"""


def slug(s):
    s = re.sub(r"\s*\(\d+\)$", "", s)            # macOS's "Name (2)" after a rename clash
    return re.sub(r"[^a-z0-9]+", "-", s.lower()).strip("-") or "mac"


def config():
    with urllib.request.urlopen(HUB + "/api/fleet", timeout=10) as r:
        machines = json.load(r)["machines"]
    wired = lambda ip: not ip.startswith(("10.6.", "169.254.", "127."))
    lines, used = ["# Written by tools/fleet.py from %s/api/fleet - every machine that took this laptop's key." % HUB], set()
    for m in sorted(machines, key=lambda m: m["name"].lower()):
        ip = next((i for i in m["ips"] if wired(i)), m["ips"][0] if m["ips"] else None)
        if not ip:
            continue
        alias = slug(m["name"])
        while alias in used:
            alias += "-2"
        used.add(alias)
        lines += ["", "Host %s" % alias, "    HostName %s" % ip, "    User %s" % m["user"],
                  "    IdentityFile %s" % KEY, "    IdentitiesOnly yes", "    ConnectTimeout 8"]
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    open(OUT, "w").write("\n".join(lines) + "\n")
    cfg = os.path.expanduser("~/.ssh/config")
    text = open(cfg).read() if os.path.exists(cfg) else ""
    if "Include ~/.ssh/config.d/*" not in text:
        open(cfg, "w").write("Include ~/.ssh/config.d/*\n\n" + text)
    print("%d machines -> %s: %s" % (len(used), OUT, ", ".join(sorted(used))))


if __name__ == "__main__":
    what = (sys.argv[1:] or ["ard"])[0]
    print(ard()) if what == "ard" else config()