"""
Tests for the shell helpers next to pxe-nic-check: pxe-dhcp-hook.sh (the
dnsmasq dhcp-script) and pxe-nic-carrier (the networkd-dispatcher hook).
No root, no network: every command they run (ss, conntrack, logger,
systemctl, pxe-nic-check) is a stub on a fake PATH that records its argv.

    python3 -m unittest discover -s tests/pxe-nic-check -v
"""

import os
import shutil
import subprocess
import tempfile
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
HELPERS = os.path.normpath(os.path.join(HERE, "..", "..", "playbook",
                                        "pxe-server-helpers"))
HOOK = os.path.join(HELPERS, "pxe-dhcp-hook.sh")
CARRIER = os.path.join(HELPERS, "pxe-nic-carrier")

# Each stub appends "name arg1 arg2 ..." (args joined by \x1f) to $CALLS.
STUB = r"""#!/bin/bash
{ printf '%s' "$(basename "$0")"; for a in "$@"; do printf '\x1f%s' "$a"; done
  printf '\n'; } >> "$CALLS"
"""

SYSTEMCTL = STUB + r"""
case "$1" in
    is-active) s=$(cat "$STUBDIR/dnsmasq.state" 2>/dev/null || echo active)
               echo "$s"; [ "$s" = active ] ;;
    restart)   [ "$2" = dnsmasq ] && exit "$(cat "$STUBDIR/restart.rc" 2>/dev/null || echo 0)"
               exit 0 ;;
    *)         exit 0 ;;
esac
"""

CHECK = STUB + r"""
[ "$1" = --print-imaging ] && cat "$STUBDIR/imaging" 2>/dev/null
exit 0
"""


class Stubs:
    def __init__(self):
        self.dir = tempfile.mkdtemp(prefix="pxe-helper-test-")
        self.bin = os.path.join(self.dir, "bin")
        os.mkdir(self.bin)
        self.calls = os.path.join(self.dir, "calls")
        for name, body in (("ss", STUB), ("conntrack", STUB),
                           ("logger", STUB), ("systemctl", SYSTEMCTL),
                           ("pxe-nic-check", CHECK)):
            self.write(os.path.join(self.bin, name), body, 0o755)
        self.sys = os.path.join(self.dir, "sys")
        self.run_dir = os.path.join(self.dir, "run")

    def write(self, path, text, mode=0o644):
        with open(path, "w") as f:
            f.write(text)
        os.chmod(path, mode)

    def env(self, **extra):
        env = {"PATH": self.bin + ":/usr/bin:/bin", "CALLS": self.calls,
               "STUBDIR": self.dir,
               "PXE_NIC_CHECK": os.path.join(self.bin, "pxe-nic-check"),
               "PXE_SYSTEMCTL": os.path.join(self.bin, "systemctl"),
               "PXE_SYS_NET": self.sys,
               "PXE_NIC_CARRIER_RUN": self.run_dir,
               "PXE_NIC_CARRIER_SETTLE": "0"}
        env.update(extra)
        return env

    def get_calls(self, name=None):
        try:
            with open(self.calls) as f:
                lines = f.read().splitlines()
        except OSError:
            return []
        out = [ln.split("\x1f") for ln in lines]
        return [c for c in out if name is None or c[0] == name]

    def reset(self):
        if os.path.exists(self.calls):
            os.unlink(self.calls)

    def link(self, iface, carrier):
        d = os.path.join(self.sys, iface)
        os.makedirs(d, exist_ok=True)
        path = os.path.join(d, "carrier")
        if carrier is None:          # admin down: reading carrier fails
            if os.path.exists(path):
                os.unlink(path)
        else:
            self.write(path, "%d\n" % carrier)

    def cleanup(self):
        shutil.rmtree(self.dir, ignore_errors=True)


class DhcpHook(unittest.TestCase):
    def setUp(self):
        self.s = Stubs()

    def tearDown(self):
        self.s.cleanup()

    def run_hook(self, action, vendor=None):
        env = self.s.env()
        if vendor is not None:
            env["DNSMASQ_VENDOR_CLASS"] = vendor
        r = subprocess.run(["bash", HOOK, action, "02:aa:00:00:00:05",
                            "172.16.9.50", "bay3105"], env=env,
                           capture_output=True, text=True, timeout=10)
        self.assertEqual(r.returncode, 0, r.stderr)
        return self.s.get_calls()

    def assert_smb_only(self, calls):
        ct = [c[1:] for c in calls if c[0] == "conntrack"]
        self.assertEqual(sorted(ct), sorted([
            ["-D", "-p", "tcp", "-s", "172.16.9.50", "--dport", "445"],
            ["-D", "-p", "tcp", "-d", "172.16.9.50", "--sport", "445"],
            ["-D", "-p", "tcp", "-s", "172.16.9.50", "--dport", "139"],
            ["-D", "-p", "tcp", "-d", "172.16.9.50", "--sport", "139"],
        ]))
        ss = [c[1:] for c in calls if c[0] == "ss"]
        self.assertEqual(ss, [["-K", "-t", "( sport = :445 or sport = :139 )"
                                           " and dst = 172.16.9.50"]])
        logs = [c for c in calls if c[0] == "logger"]
        self.assertEqual(len(logs), 1)
        self.assertIn("SMB", logs[0][-1])

    def test_add_and_del_clear_smb_state_whatever_the_vendor(self):
        for action in ("add", "del"):
            for vendor in ("PXEClient:Arch:00007:UNDI:003016",
                           "MSFT 5.0", None):
                self.s.reset()
                self.assert_smb_only(self.run_hook(action, vendor))

    def test_old_only_for_network_boot(self):
        self.s.reset()
        calls = self.run_hook("old", "PXEClient:Arch:00007:UNDI:003016")
        self.assert_smb_only(calls)
        self.assertIn("network boot, PXEClient:Arch:00007",
                      [c for c in calls if c[0] == "logger"][0][-1])
        for vendor in ("MSFT 5.0", None, ""):
            self.s.reset()
            self.assertEqual(self.run_hook("old", vendor), [], vendor)

    def test_never_clears_all_traffic(self):
        calls = self.run_hook("add", "PXEClient")
        for c in calls:
            if c[0] == "conntrack":
                self.assertIn("-p", c)
                self.assertIn("tcp", c)
            if c[0] == "ss":
                self.assertIn("sport = :445", c[-1])

    def test_no_ip_does_nothing(self):
        r = subprocess.run(["bash", HOOK, "add", "02:aa:00:00:00:05", ""],
                           env=self.s.env(), capture_output=True, timeout=10)
        self.assertEqual(r.returncode, 0)
        self.assertEqual(self.s.get_calls(), [])


class Carrier(unittest.TestCase):
    def setUp(self):
        self.s = Stubs()
        self.s.write(os.path.join(self.s.dir, "imaging"), "pxe0\n")
        self.s.link("pxe0", 1)
        self.s.link("corp0", 1)

    def tearDown(self):
        self.s.cleanup()

    def event(self, iface, state="routable"):
        r = subprocess.run(["bash", CARRIER], env=self.s.env(
            IFACE=iface, STATE=state), capture_output=True, text=True,
            timeout=20)
        self.assertEqual(r.returncode, 0, r.stderr)
        return self.s.get_calls()

    def gain(self):
        r = subprocess.run(["bash", CARRIER, "--gain"], env=self.s.env(),
                           capture_output=True, text=True, timeout=20)
        self.assertEqual(r.returncode, 0, r.stderr)
        return self.s.get_calls()

    def systemctl(self, calls):
        return [c[1:] for c in calls if c[0] == "systemctl"]

    def dnsmasq(self, state):
        self.s.write(os.path.join(self.s.dir, "dnsmasq.state"), state + "\n")

    def marker(self):
        return os.path.exists(os.path.join(self.s.run_dir, "stopped-dnsmasq"))

    def test_ge_nic_is_never_touched(self):
        for carrier in (0, 1, None, 1):
            self.s.reset()
            self.s.link("corp0", carrier)
            calls = self.event("corp0", "no-carrier")
            self.assertEqual([c[0] for c in calls], ["pxe-nic-check"])
        self.assertFalse(os.path.exists(os.path.join(self.s.run_dir,
                                                     "link.corp0")))

    def test_unknown_imaging_nic_does_nothing(self):
        self.s.write(os.path.join(self.s.dir, "imaging"), "")
        calls = self.event("pxe0", "no-carrier")
        self.assertEqual([c[0] for c in calls], ["pxe-nic-check"])

    def test_loss_stops_running_dnsmasq_and_marks(self):
        self.event("pxe0", "routable")         # first sight: up
        self.s.reset()
        self.s.link("pxe0", 0)
        sc = self.systemctl(self.event("pxe0", "no-carrier"))
        self.assertIn(["stop", "pxe-nic-carrier.service"], sc)
        self.assertIn(["stop", "--no-block", "dnsmasq"], sc)
        self.assertLess(sc.index(["stop", "pxe-nic-carrier.service"]),
                        sc.index(["stop", "--no-block", "dnsmasq"]))
        self.assertTrue(self.marker())
        # same state again (off after no-carrier): idempotent, no calls
        self.s.reset()
        self.s.link("pxe0", None)
        self.assertEqual(self.systemctl(self.event("pxe0", "off")), [])

    def test_loss_with_dnsmasq_stopped_leaves_no_marker(self):
        self.dnsmasq("inactive")
        self.s.link("pxe0", 0)
        sc = self.systemctl(self.event("pxe0", "no-carrier"))
        self.assertNotIn(["stop", "--no-block", "dnsmasq"], sc)
        self.assertFalse(self.marker())

    def test_gain_hands_off_once(self):
        self.s.link("pxe0", 0)
        self.event("pxe0", "no-carrier")
        self.s.reset()
        self.s.link("pxe0", 1)
        sc = self.systemctl(self.event("pxe0", "carrier"))
        self.assertEqual(sc[-1], ["restart", "--no-block",
                                  "pxe-nic-carrier.service"])
        self.assertNotIn("dnsmasq", [x for c in sc for x in c])
        self.s.reset()                       # routable right after carrier
        self.assertEqual(self.systemctl(self.event("pxe0", "routable")), [])

    def test_bridge_port_events_count(self):
        self.s.write(os.path.join(self.s.dir, "imaging"), "br-pxe\npxe0\n")
        self.s.link("br-pxe", 1)
        self.event("pxe0", "enslaved")
        self.s.reset()
        self.s.link("pxe0", 0)
        sc = self.systemctl(self.event("pxe0", "no-carrier"))
        self.assertIn(["stop", "--no-block", "dnsmasq"], sc)

    def test_gain_restarts_dnsmasq_through_its_guard(self):
        os.makedirs(self.s.run_dir)
        self.s.write(os.path.join(self.s.run_dir, "stopped-dnsmasq"), "")
        self.dnsmasq("inactive")              # the hook stopped it
        sc = self.systemctl(self.gain())
        self.assertIn(["restart", "dnsmasq"], sc)
        self.assertFalse(self.marker())

    def test_gain_restarts_running_dnsmasq(self):
        # loss never seen (fast swap, boot): dnsmasq running, no marker
        sc = self.systemctl(self.gain())
        self.assertIn(["restart", "dnsmasq"], sc)

    def test_gain_logs_crit_when_guard_refuses(self):
        self.s.write(os.path.join(self.s.dir, "restart.rc"), "1\n")
        calls = self.gain()
        logs = [c for c in calls if c[0] == "logger"]
        self.assertTrue(any("daemon.crit" in c and "NOT started" in c[-1]
                            for c in logs), logs)

    def test_gain_leaves_operator_stopped_dnsmasq_alone(self):
        self.dnsmasq("inactive")
        calls = self.gain()
        self.assertNotIn(["restart", "dnsmasq"], self.systemctl(calls))
        self.assertNotIn(["start", "dnsmasq"], self.systemctl(calls))
        self.assertIn(["pxe-nic-check", "--guard"],
                      [c for c in calls if c[0] == "pxe-nic-check"])

    def test_gain_does_nothing_if_link_dropped_again(self):
        self.s.link("pxe0", 0)
        calls = self.gain()
        self.assertEqual(self.systemctl(calls), [])
        self.assertNotIn(["pxe-nic-check", "--guard"], calls)


if __name__ == "__main__":
    unittest.main()