feat(homelab): enable Wake-on-LAN on the mini PCs (#509)

scripts/tests/test_wake.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import io
+import unittest
+
+from scripts.wake import (
+    Host,
+    HostSelectionError,
+    Reporter,
+    WakeConfig,
+    WakeSession,
+    magic_packet,
+    resolve_hosts,
+)
+
+
+class FakeClock:
+    def __init__(self) -> None:
+        self.now = 0.0
+
+    def __call__(self) -> float:
+        return self.now
+
+    def sleep(self, seconds: float) -> None:
+        self.now += seconds
+
+
+class RecordingSender:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, int]] = []
+
+    def __call__(self, host: Host, packets_per_broadcast: int) -> None:
+        self.calls.append((host.name, packets_per_broadcast))
+
+
+class SequenceProbe:
+    def __init__(self, online_on_call: int) -> None:
+        self._online_on_call = online_on_call
+        self._calls = 0
+
+    def __call__(self, _host: Host) -> bool:
+        self._calls += 1
+        return self._calls >= self._online_on_call
+
+
+class MagicPacketTests(unittest.TestCase):
+    def test_builds_standard_packet(self) -> None:
+        mac = "8c:32:23:36:81:84"
+        packet = magic_packet(mac)
+
+        self.assertEqual(len(packet), 102)
+        self.assertEqual(packet[:6], b"\xff" * 6)
+        self.assertEqual(packet[6:], bytes.fromhex(mac.replace(":", "")) * 16)
+
+    def test_rejects_invalid_mac(self) -> None:
+        with self.assertRaises(ValueError):
+            magic_packet("not-a-mac")
+
+
+class ResolveHostsTests(unittest.TestCase):
+    def setUp(self) -> None:
+        self.hosts = {
+            "alpha": Host("alpha", "00:00:00:00:00:01", "192.0.2.1"),
+            "beta": Host("beta", "00:00:00:00:00:02", "192.0.2.2"),
+        }
+
+    def test_defaults_to_all_hosts(self) -> None:
+        self.assertEqual(resolve_hosts([], hosts=self.hosts), list(self.hosts.values()))
+
+    def test_resolves_alias_and_removes_duplicates(self) -> None:
+        selected = resolve_hosts(
+            ["a", "alpha", "beta"],
+            hosts=self.hosts,
+            aliases={"a": "alpha"},
+        )
+        self.assertEqual(selected, [self.hosts["alpha"], self.hosts["beta"]])
+
+    def test_rejects_unknown_host(self) -> None:
+        with self.assertRaises(HostSelectionError):
+            resolve_hosts(["missing"], hosts=self.hosts, aliases={})
+
+
+class WakeSessionTests(unittest.TestCase):
+    def test_retries_until_host_is_online(self) -> None:
+        clock = FakeClock()
+        sender = RecordingSender()
+        stdout = io.StringIO()
+        stderr = io.StringIO()
+        reporter = Reporter(stream=stdout, error_stream=stderr, clock=clock)
+        host = Host("test", "00:00:00:00:00:01", "192.0.2.1")
+        config = WakeConfig(
+            timeout=30,
+            retry_interval=5,
+            probe_interval=1,
+            progress_interval=5,
+            packets_per_broadcast=2,
+        )
+        session = WakeSession(
+            [host],
+            config,
+            reporter,
+            sender=sender,
+            probe=SequenceProbe(online_on_call=10),
+            clock=clock,
+            sleep=clock.sleep,
+        )
+
+        self.assertEqual(session.run(), 0)
+        self.assertEqual(sender.calls, [("test", 2), ("test", 2)])
+        output = stdout.getvalue()
+        self.assertIn("attempt 2", output)
+        self.assertIn("still waiting for: test (attempt 2, 192.0.2.1)", output)
+        self.assertIn("test is online", output)
+
+    def test_exits_immediately_when_host_is_already_online(self) -> None:
+        clock = FakeClock()
+        sender = RecordingSender()
+        reporter = Reporter(
+            stream=io.StringIO(), error_stream=io.StringIO(), clock=clock
+        )
+        host = Host("test", "00:00:00:00:00:01", "192.0.2.1")
+        session = WakeSession(
+            [host],
+            WakeConfig(),
+            reporter,
+            sender=sender,
+            probe=lambda _host: True,
+            clock=clock,
+            sleep=clock.sleep,
+        )
+
+        self.assertEqual(session.run(), 0)
+        self.assertEqual(sender.calls, [])
+
+    def test_times_out_when_host_never_comes_online(self) -> None:
+        clock = FakeClock()
+        sender = RecordingSender()
+        stdout = io.StringIO()
+        stderr = io.StringIO()
+        reporter = Reporter(stream=stdout, error_stream=stderr, clock=clock)
+        host = Host("test", "00:00:00:00:00:01", "192.0.2.1")
+        session = WakeSession(
+            [host],
+            WakeConfig(timeout=3, retry_interval=5),
+            reporter,
+            sender=sender,
+            probe=lambda _host: False,
+            clock=clock,
+            sleep=clock.sleep,
+        )
+
+        self.assertEqual(session.run(), 1)
+        self.assertIn("timed out waiting for: test", stderr.getvalue())
+
+
+if __name__ == "__main__":
+    unittest.main()

scripts/wake-homelab.py

@@ -1,46 +0,0 @@
-#!/usr/bin/env python3
-"""Wake every homelab host whose NixOS config arms Wake-on-LAN."""
-
-from __future__ import annotations
-
-import socket
-import sys
-
-# Keep this list limited to hosts that actually set Wake-on: g in NixOS.
-# Add a host here only after its WoL service and BIOS setting are in place.
-HOSTS = {
-    "aquamarine": "8c:32:23:36:81:84",
-}
-BROADCASTS = ("192.168.5.255", "255.255.255.255")
-PORT = 9
-PACKETS = 3
-
-
-def magic_packet(mac: str) -> bytes:
-    octets = bytes.fromhex(mac.replace(":", ""))
-    if len(octets) != 6:
-        raise ValueError(f"invalid MAC: {mac}")
-    return b"\xff" * 6 + octets * 16
-
-
-def main() -> int:
-    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
-    failed = False
-
-    for host, mac in HOSTS.items():
-        packet = magic_packet(mac)
-        print(f"waking {host} ({mac})")
-        for _ in range(PACKETS):
-            for broadcast in BROADCASTS:
-                try:
-                    sock.sendto(packet, (broadcast, PORT))
-                except OSError as exc:
-                    print(f"  {broadcast}:{PORT}: {exc}", file=sys.stderr)
-                    failed = True
-
-    return 1 if failed else 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

scripts/wake.py

@@ -0,0 +1,427 @@
+#!/usr/bin/env python3
+"""Wake homelab hosts over Wake-on-LAN and wait until they are reachable."""
+
+from __future__ import annotations
+
+import argparse
+import socket
+import subprocess
+import sys
+import time
+from collections.abc import Callable, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from typing import Self, TextIO
+
+BROADCAST_ADDRESSES = ("192.168.5.255", "255.255.255.255")
+WOL_PORT = 9
+SSH_PORT = 22
+PING_WAIT_SECONDS = 1
+PING_TIMEOUT_SECONDS = PING_WAIT_SECONDS + 1
+SSH_TIMEOUT_SECONDS = 1.0
+
+DEFAULT_TIMEOUT_SECONDS = 180.0
+DEFAULT_RETRY_INTERVAL_SECONDS = 5.0
+DEFAULT_PROBE_INTERVAL_SECONDS = 1.0
+DEFAULT_PROGRESS_INTERVAL_SECONDS = 5.0
+DEFAULT_PACKETS_PER_BROADCAST = 3
+
+
+@dataclass(frozen=True, slots=True)
+class Host:
+    """A Wake-on-LAN-capable host."""
+
+    name: str
+    mac: str
+    ip: str
+
+
+# Keep this list limited to hosts that actually set Wake-on: g in NixOS.
+# Add a host here only after its firmware setting and NixOS link config exist.
+HOSTS: Mapping[str, Host] = {
+    "aquamarine": Host("aquamarine", "8c:32:23:36:81:84", "192.168.5.101"),
+    "shoryu": Host("shoryu", "70:70:fc:01:1b:56", "192.168.5.181"),
+    "shushou": Host("shushou", "58:47:ca:70:0d:72", "192.168.5.182"),
+    "youko": Host("youko", "70:70:fc:01:00:44", "192.168.5.183"),
+}
+
+ALIASES: Mapping[str, str] = {
+    "aqua": "aquamarine",
+}
+
+
+@dataclass(frozen=True, slots=True)
+class WakeConfig:
+    """Timing and packet settings for one wake run."""
+
+    timeout: float = DEFAULT_TIMEOUT_SECONDS
+    retry_interval: float = DEFAULT_RETRY_INTERVAL_SECONDS
+    probe_interval: float = DEFAULT_PROBE_INTERVAL_SECONDS
+    progress_interval: float = DEFAULT_PROGRESS_INTERVAL_SECONDS
+    packets_per_broadcast: int = DEFAULT_PACKETS_PER_BROADCAST
+
+
+class HostSelectionError(ValueError):
+    """Raised when a requested host name is unknown."""
+
+
+class Reporter:
+    """Print timestamped progress messages."""
+
+    def __init__(
+        self,
+        *,
+        stream: TextIO | None = None,
+        error_stream: TextIO | None = None,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self._stream = stream if stream is not None else sys.stdout
+        self._error_stream = error_stream if error_stream is not None else sys.stderr
+        self._clock = clock
+        self._started_at = clock()
+
+    def elapsed(self) -> float:
+        return self._clock() - self._started_at
+
+    def info(self, message: str) -> None:
+        print(f"[{self.elapsed():5.1f}s] {message}", file=self._stream, flush=True)
+
+    def error(self, message: str) -> None:
+        print(
+            f"[{self.elapsed():5.1f}s] {message}", file=self._error_stream, flush=True
+        )
+
+
+class MagicPacketSender:
+    """Send Wake-on-LAN magic packets over a UDP broadcast socket."""
+
+    def __init__(
+        self,
+        *,
+        broadcasts: Sequence[str] = BROADCAST_ADDRESSES,
+        port: int = WOL_PORT,
+    ) -> None:
+        self._broadcasts = tuple(broadcasts)
+        self._port = port
+        self._socket: socket.socket | None = None
+
+    def __enter__(self) -> Self:
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+        return self
+
+    def __exit__(self, *_exc_info: object) -> None:
+        if self._socket is not None:
+            self._socket.close()
+            self._socket = None
+
+    def send(self, host: Host, packets_per_broadcast: int) -> None:
+        """Send one magic-packet burst."""
+        if self._socket is None:
+            raise RuntimeError("sender must be used as a context manager")
+
+        packet = magic_packet(host.mac)
+        for _ in range(packets_per_broadcast):
+            for broadcast in self._broadcasts:
+                self._socket.sendto(packet, (broadcast, self._port))
+
+
+@dataclass(slots=True)
+class _HostState:
+    host: Host
+    attempts: int = 0
+    last_sent_at: float = 0.0
+
+
+class WakeSession:
+    """Retry magic packets until all hosts are online or the timeout expires."""
+
+    def __init__(
+        self,
+        hosts: Sequence[Host],
+        config: WakeConfig,
+        reporter: Reporter,
+        *,
+        sender: Callable[[Host, int], None],
+        probe: Callable[[Host], bool] | None = None,
+        clock: Callable[[], float] = time.monotonic,
+        sleep: Callable[[float], None] = time.sleep,
+    ) -> None:
+        self._states = {host.name: _HostState(host) for host in hosts}
+        self._config = config
+        self._reporter = reporter
+        self._sender = sender
+        self._probe = probe or host_online
+        self._clock = clock
+        self._sleep = sleep
+        self._send_failed = False
+        self._last_progress_at = 0.0
+
+    def run(self) -> int:
+        self._probe_pending(initial=True)
+        if not self._states:
+            return 0
+
+        for name in sorted(self._states):
+            self._send(self._states[name])
+        self._reporter.info(
+            f"waking {self._describe_pending()}; retrying every "
+            f"{self._config.retry_interval:g}s until online"
+        )
+        self._last_progress_at = self._clock()
+
+        while self._states:
+            if self._timed_out():
+                names = ", ".join(sorted(self._states))
+                self._reporter.error(f"timed out waiting for: {names}")
+                return 1
+
+            self._probe_pending(initial=False)
+            if not self._states:
+                break
+
+            self._retry_pending()
+            self._report_progress()
+            self._sleep(self._config.probe_interval)
+
+        return 1 if self._send_failed else 0
+
+    def _timed_out(self) -> bool:
+        return (
+            self._config.timeout > 0
+            and self._reporter.elapsed() >= self._config.timeout
+        )
+
+    def _probe_pending(self, *, initial: bool) -> None:
+        states = list(self._states.values())
+        if not states:
+            return
+
+        with ThreadPoolExecutor(max_workers=len(states)) as executor:
+            futures = {
+                executor.submit(self._probe, state.host): state for state in states
+            }
+            for future in as_completed(futures):
+                state = futures[future]
+                if not future.result():
+                    continue
+
+                self._states.pop(state.host.name, None)
+                if initial:
+                    self._reporter.info(f"{state.host.name} is already online")
+                else:
+                    self._reporter.info(
+                        f"{state.host.name} is online after {state.attempts} magic-packet bursts"
+                    )
+
+    def _send(self, state: _HostState) -> None:
+        state.attempts += 1
+        try:
+            self._sender(state.host, self._config.packets_per_broadcast)
+        except OSError as exc:
+            self._send_failed = True
+            self._reporter.error(
+                f"failed to send magic packets to {state.host.name}: {exc}"
+            )
+            return
+
+        state.last_sent_at = self._clock()
+
+    def _retry_pending(self) -> None:
+        now = self._clock()
+        for name in sorted(self._states):
+            state = self._states[name]
+            if now - state.last_sent_at >= self._config.retry_interval:
+                self._send(state)
+
+    def _report_progress(self) -> None:
+        now = self._clock()
+        if now - self._last_progress_at < self._config.progress_interval:
+            return
+
+        self._reporter.info(f"still waiting for: {self._describe_pending()}")
+        self._last_progress_at = now
+
+    def _describe_pending(self) -> str:
+        return ", ".join(
+            f"{name} (attempt {state.attempts}, {state.host.ip})"
+            for name, state in sorted(self._states.items())
+        )
+
+
+def magic_packet(mac: str) -> bytes:
+    """Build the standard 102-byte Wake-on-LAN magic packet."""
+    try:
+        mac_bytes = bytes.fromhex(mac.replace(":", ""))
+    except ValueError as exc:
+        raise ValueError(f"invalid MAC address: {mac!r}") from exc
+
+    if len(mac_bytes) != 6:
+        raise ValueError(f"invalid MAC address: {mac!r}")
+
+    return b"\xff" * 6 + mac_bytes * 16
+
+
+def resolve_hosts(
+    names: Sequence[str],
+    *,
+    hosts: Mapping[str, Host] = HOSTS,
+    aliases: Mapping[str, str] = ALIASES,
+) -> list[Host]:
+    """Resolve CLI host names and aliases, preserving order and removing duplicates."""
+    if not names:
+        return list(hosts.values())
+
+    selected: dict[str, Host] = {}
+    for requested_name in names:
+        canonical_name = aliases.get(requested_name, requested_name)
+        try:
+            host = hosts[canonical_name]
+        except KeyError as exc:
+            known = ", ".join(sorted(hosts))
+            alias_help = ", ".join(
+                f"{alias} -> {target}" for alias, target in aliases.items()
+            )
+            raise HostSelectionError(
+                f"unknown host {requested_name!r}; known hosts: {known}; aliases: {alias_help}"
+            ) from exc
+        selected[host.name] = host
+
+    return list(selected.values())
+
+
+def host_online(host: Host) -> bool:
+    """Return whether a host answers ICMP and has reached a usable SSH state.
+
+    Requiring SSH prevents a host that is still shutting down, or partway through boot, from being
+    mistaken for online just because it still answers ping.
+    """
+    return _ping(host.ip) and _ssh_port_open(host.ip)
+
+
+def _ping(ip: str) -> bool:
+    try:
+        result = subprocess.run(
+            ["ping", "-c", "1", "-W", str(PING_WAIT_SECONDS), ip],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            timeout=PING_TIMEOUT_SECONDS + 1,
+            check=False,
+        )
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return False
+    return result.returncode == 0
+
+
+def _ssh_port_open(ip: str) -> bool:
+    try:
+        with socket.create_connection((ip, SSH_PORT), timeout=SSH_TIMEOUT_SECONDS):
+            return True
+    except OSError:
+        return False
+
+
+def _positive_float(value: str) -> float:
+    number = float(value)
+    if number <= 0:
+        raise argparse.ArgumentTypeError("must be greater than zero")
+    return number
+
+
+def _non_negative_float(value: str) -> float:
+    number = float(value)
+    if number < 0:
+        raise argparse.ArgumentTypeError("must be zero or greater")
+    return number
+
+
+def _positive_int(value: str) -> int:
+    number = int(value)
+    if number <= 0:
+        raise argparse.ArgumentTypeError("must be greater than zero")
+    return number
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="just wake",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description=(
+            "Wake selected homelab hosts, keep retrying the magic packet, and exit when every "
+            "selected host answers. With no host arguments, wake every WoL-capable host."
+        ),
+        epilog=(
+            "Examples:\n"
+            "  just wake\n"
+            "  just wake aqua youko\n"
+            "  just wake shoryu --timeout 0 --progress-interval 2"
+        ),
+    )
+    parser.add_argument(
+        "hosts",
+        nargs="*",
+        metavar="HOST",
+        help="hosts to wake (default: all); use 'aqua' as an alias for aquamarine",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=_non_negative_float,
+        default=DEFAULT_TIMEOUT_SECONDS,
+        metavar="SECONDS",
+        help="stop waiting after this many seconds (default: 180; 0 waits forever)",
+    )
+    parser.add_argument(
+        "--retry-interval",
+        type=_positive_float,
+        default=DEFAULT_RETRY_INTERVAL_SECONDS,
+        metavar="SECONDS",
+        help="time between magic-packet bursts for hosts still offline (default: 5)",
+    )
+    parser.add_argument(
+        "--probe-interval",
+        type=_positive_float,
+        default=DEFAULT_PROBE_INTERVAL_SECONDS,
+        metavar="SECONDS",
+        help="time between reachability checks (default: 1)",
+    )
+    parser.add_argument(
+        "--progress-interval",
+        type=_positive_float,
+        default=DEFAULT_PROGRESS_INTERVAL_SECONDS,
+        metavar="SECONDS",
+        help="how often to print a waiting-for progress line (default: 5)",
+    )
+    parser.add_argument(
+        "--packets",
+        type=_positive_int,
+        default=DEFAULT_PACKETS_PER_BROADCAST,
+        metavar="N",
+        help="magic packets per broadcast address and burst (default: 3)",
+    )
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    try:
+        hosts = resolve_hosts(args.hosts)
+    except HostSelectionError as exc:
+        parser.error(str(exc))
+
+    config = WakeConfig(
+        timeout=args.timeout,
+        retry_interval=args.retry_interval,
+        probe_interval=args.probe_interval,
+        progress_interval=args.progress_interval,
+        packets_per_broadcast=args.packets,
+    )
+
+    reporter = Reporter()
+    with MagicPacketSender() as sender:
+        return WakeSession(hosts, config, reporter, sender=sender.send).run()
+
+
+if __name__ == "__main__":
+    sys.exit(main())