Add sd_wifi.py -- operator tool for updating WiFi config on SD cards

Interactive and non-interactive CLI to detect, mount, and edit the
netplan 20-wifi.yaml on an armbi_root-labelled SD card without
needing to boot the implant.

scripts/sd_wifi.py

@@ -0,0 +1,516 @@
+#!/usr/bin/env python3
+"""sd_wifi.py — Update WiFi config on a BigBrother SD card.
+
+Detects the SD card by looking for the armbi_root label, mounts it to a
+temp directory, reads/writes /etc/netplan/20-wifi.yaml, and unmounts on exit.
+
+Usage:
+  ./sd_wifi.py                         # interactive menu
+  ./sd_wifi.py --list                  # show current configured networks
+  ./sd_wifi.py --add SSID PASSWORD     # non-interactive add/update
+"""
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import click
+import yaml
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Prompt
+from rich.table import Table
+from rich.text import Text
+
+console = Console()
+
+NETPLAN_PATH = "etc/netplan/20-wifi.yaml"
+SD_LABEL = "armbi_root"
+
+
+# ---------------------------------------------------------------------------
+# SD card detection and mount
+# ---------------------------------------------------------------------------
+
+def find_sd_device() -> Optional[str]:
+    """Return the block device path for the partition labelled armbi_root."""
+    try:
+        result = subprocess.run(
+            ["lsblk", "-o", "NAME,LABEL", "--json"],
+            capture_output=True, text=True, check=True,
+        )
+        data = json.loads(result.stdout)
+    except (subprocess.CalledProcessError, json.JSONDecodeError, FileNotFoundError) as exc:
+        console.print(f"[red]Error running lsblk:[/red] {exc}")
+        return None
+
+    def _search(devices):
+        for dev in devices:
+            label = dev.get("label") or ""
+            if label.strip() == SD_LABEL:
+                return f"/dev/{dev['name']}"
+            children = dev.get("children") or []
+            found = _search(children)
+            if found:
+                return found
+        return None
+
+    return _search(data.get("blockdevices", []))
+
+
+def mount_sd(device: str, mountpoint: str) -> bool:
+    """Mount device to mountpoint using sudo. Returns True on success."""
+    result = subprocess.run(
+        ["sudo", "mount", device, mountpoint],
+        capture_output=True, text=True,
+    )
+    if result.returncode != 0:
+        console.print(f"[red]Mount failed:[/red] {result.stderr.strip()}")
+        return False
+    return True
+
+
+def unmount_sd(mountpoint: str) -> None:
+    """Unmount mountpoint using sudo (best-effort, ignores errors)."""
+    subprocess.run(
+        ["sudo", "umount", mountpoint],
+        capture_output=True, text=True,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Netplan read / write
+# ---------------------------------------------------------------------------
+
+NETPLAN_TEMPLATE: Dict = {
+    "network": {
+        "version": 2,
+        "renderer": "networkd",
+        "wifis": {
+            "wlan0": {
+                "dhcp4": True,
+                "dhcp6": True,
+                "regulatory-domain": "US",
+                "access-points": {},
+            }
+        },
+    }
+}
+
+
+def read_netplan(mountpoint: str) -> Dict:
+    """Read and parse the netplan WiFi config from the mounted SD card.
+
+    Returns the parsed config dict, or a fresh template if the file does not
+    exist or cannot be parsed.
+    """
+    netplan_file = Path(mountpoint) / NETPLAN_PATH
+    if not netplan_file.exists():
+        console.print(f"[yellow]No netplan file found at {netplan_file}. Starting fresh.[/yellow]")
+        return _deep_copy(NETPLAN_TEMPLATE)
+
+    try:
+        with open(netplan_file, "r") as fh:
+            data = yaml.safe_load(fh)
+        if not data or "network" not in data:
+            console.print("[yellow]Netplan file appears empty or malformed. Starting fresh.[/yellow]")
+            return _deep_copy(NETPLAN_TEMPLATE)
+        return data
+    except yaml.YAMLError as exc:
+        console.print(f"[yellow]YAML parse error ({exc}). Starting fresh.[/yellow]")
+        return _deep_copy(NETPLAN_TEMPLATE)
+
+
+def write_netplan(mountpoint: str, config: Dict) -> None:
+    """Write the netplan config to the SD card and set permissions to 600."""
+    netplan_file = Path(mountpoint) / NETPLAN_PATH
+
+    # Ensure the directory exists on the card
+    netplan_file.parent.mkdir(parents=True, exist_ok=True)
+
+    yaml_text = yaml.dump(config, default_flow_style=False, sort_keys=False,
+                          allow_unicode=True)
+
+    # Write via sudo tee to handle root-owned filesystem
+    result = subprocess.run(
+        ["sudo", "tee", str(netplan_file)],
+        input=yaml_text, capture_output=True, text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"Failed to write netplan file: {result.stderr.strip()}")
+
+    # Set permissions 600
+    chmod_result = subprocess.run(
+        ["sudo", "chmod", "600", str(netplan_file)],
+        capture_output=True, text=True,
+    )
+    if chmod_result.returncode != 0:
+        console.print(f"[yellow]Warning: chmod 600 failed: {chmod_result.stderr.strip()}[/yellow]")
+
+
+def _deep_copy(obj):
+    """Deep-copy a plain dict/list/scalar structure via JSON round-trip."""
+    return json.loads(json.dumps(obj))
+
+
+# ---------------------------------------------------------------------------
+# Config accessors
+# ---------------------------------------------------------------------------
+
+def get_access_points(config: Dict) -> Dict[str, str]:
+    """Return {ssid: password} mapping from the config."""
+    try:
+        return dict(
+            config["network"]["wifis"]["wlan0"].get("access-points", {}) or {}
+        )
+    except (KeyError, TypeError):
+        return {}
+
+
+def get_country(config: Dict) -> str:
+    """Return the regulatory-domain value."""
+    try:
+        return config["network"]["wifis"]["wlan0"].get("regulatory-domain", "US")
+    except (KeyError, TypeError):
+        return "US"
+
+
+def set_access_points(config: Dict, aps: Dict[str, str]) -> None:
+    """Overwrite the access-points section."""
+    config["network"]["wifis"]["wlan0"]["access-points"] = aps
+
+
+def set_country(config: Dict, country: str) -> None:
+    """Set the regulatory-domain value."""
+    config["network"]["wifis"]["wlan0"]["regulatory-domain"] = country.upper()
+
+
+# ---------------------------------------------------------------------------
+# Display helpers
+# ---------------------------------------------------------------------------
+
+def print_header() -> None:
+    console.print()
+    console.print(Panel(
+        "[bold cyan]BigBrother[/bold cyan] — SD Card WiFi Configurator\n"
+        "Target label: [yellow]armbi_root[/yellow]  |  "
+        "Netplan: [dim]" + NETPLAN_PATH + "[/dim]",
+        title="sd_wifi",
+        border_style="cyan",
+    ))
+    console.print()
+
+
+def print_networks(config: Dict) -> None:
+    """Print a Rich table of configured access-points."""
+    aps = get_access_points(config)
+    country = get_country(config)
+
+    console.print()
+    if not aps:
+        console.print("  [dim]No WiFi networks configured.[/dim]")
+    else:
+        table = Table(title=f"Configured Networks  (country: {country})",
+                      show_lines=True, border_style="cyan")
+        table.add_column("#", justify="right", min_width=3, style="dim")
+        table.add_column("SSID", min_width=24, style="bold")
+        table.add_column("Password", min_width=20)
+
+        for idx, (ssid, details) in enumerate(aps.items(), start=1):
+            if isinstance(details, dict):
+                pwd = details.get("password", "")
+            else:
+                pwd = str(details) if details else ""
+            table.add_row(str(idx), ssid, pwd or "[dim](none)[/dim]")
+
+        console.print(table)
+    console.print()
+
+
+# ---------------------------------------------------------------------------
+# Interactive menu actions
+# ---------------------------------------------------------------------------
+
+def action_add_network(config: Dict) -> bool:
+    """Prompt for SSID + password and add (or update) the entry. Returns True if changed."""
+    console.print()
+    ssid = Prompt.ask("  [cyan]SSID[/cyan]").strip()
+    if not ssid:
+        console.print("  [yellow]Cancelled — empty SSID.[/yellow]")
+        return False
+
+    password = Prompt.ask("  [cyan]Password[/cyan] (leave blank for open network)").strip()
+
+    aps = get_access_points(config)
+    existed = ssid in aps
+    aps[ssid] = {"password": password} if password else {}
+    set_access_points(config, aps)
+
+    verb = "Updated" if existed else "Added"
+    console.print(f"  [green]{verb}[/green] [bold]{ssid}[/bold]")
+    return True
+
+
+def action_remove_network(config: Dict) -> bool:
+    """Prompt the user to choose a network to remove. Returns True if changed."""
+    aps = get_access_points(config)
+    if not aps:
+        console.print("  [yellow]No networks configured.[/yellow]")
+        return False
+
+    ssid_list = list(aps.keys())
+    console.print()
+    for idx, ssid in enumerate(ssid_list, start=1):
+        console.print(f"    [bold]{idx}.[/bold] {ssid}")
+    console.print()
+
+    choice = Prompt.ask("  Remove # (or [dim]Enter[/dim] to cancel)").strip()
+    if not choice:
+        return False
+
+    try:
+        n = int(choice)
+        if not (1 <= n <= len(ssid_list)):
+            raise ValueError
+    except ValueError:
+        console.print("  [yellow]Invalid selection.[/yellow]")
+        return False
+
+    target = ssid_list[n - 1]
+    del aps[target]
+    set_access_points(config, aps)
+    console.print(f"  [green]Removed[/green] [bold]{target}[/bold]")
+    return True
+
+
+def action_update_password(config: Dict) -> bool:
+    """Update the password for an existing network. Returns True if changed."""
+    aps = get_access_points(config)
+    if not aps:
+        console.print("  [yellow]No networks configured.[/yellow]")
+        return False
+
+    ssid_list = list(aps.keys())
+    console.print()
+    for idx, ssid in enumerate(ssid_list, start=1):
+        console.print(f"    [bold]{idx}.[/bold] {ssid}")
+    console.print()
+
+    choice = Prompt.ask("  Update password for # (or [dim]Enter[/dim] to cancel)").strip()
+    if not choice:
+        return False
+
+    try:
+        n = int(choice)
+        if not (1 <= n <= len(ssid_list)):
+            raise ValueError
+    except ValueError:
+        console.print("  [yellow]Invalid selection.[/yellow]")
+        return False
+
+    target = ssid_list[n - 1]
+    password = Prompt.ask(f"  New password for [bold]{target}[/bold] "
+                          "(leave blank for open network)").strip()
+
+    existing = aps[target]
+    if isinstance(existing, dict):
+        entry = dict(existing)
+    else:
+        entry = {}
+
+    if password:
+        entry["password"] = password
+    else:
+        entry.pop("password", None)
+
+    aps[target] = entry if entry else {}
+    set_access_points(config, aps)
+    console.print(f"  [green]Password updated[/green] for [bold]{target}[/bold]")
+    return True
+
+
+def action_change_country(config: Dict) -> bool:
+    """Prompt for a new regulatory-domain country code. Returns True if changed."""
+    current = get_country(config)
+    console.print()
+    new_country = Prompt.ask(
+        f"  Country code (current: [yellow]{current}[/yellow])"
+    ).strip().upper()
+
+    if not new_country:
+        console.print("  [yellow]Cancelled.[/yellow]")
+        return False
+
+    if len(new_country) != 2 or not new_country.isalpha():
+        console.print("  [yellow]Invalid country code — must be 2 letters (e.g. US, GB, DE).[/yellow]")
+        return False
+
+    set_country(config, new_country)
+    console.print(f"  [green]Country set to[/green] [bold]{new_country}[/bold]")
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Core flow
+# ---------------------------------------------------------------------------
+
+def run_interactive(device: str, mountpoint: str) -> None:
+    """Mount SD card, run interactive menu, write changes, unmount."""
+    console.print(f"  Mounting [cyan]{device}[/cyan] → [dim]{mountpoint}[/dim] ...", end=" ")
+
+    if not mount_sd(device, mountpoint):
+        return
+
+    console.print("[green]OK[/green]")
+
+    try:
+        config = read_netplan(mountpoint)
+        dirty = False
+
+        while True:
+            print_networks(config)
+            country = get_country(config)
+            console.print(f"  Country: [yellow]{country}[/yellow]")
+            console.print()
+            console.print("  [bold]1.[/bold] Add / update network")
+            console.print("  [bold]2.[/bold] Remove network")
+            console.print("  [bold]3.[/bold] Update password")
+            console.print("  [bold]4.[/bold] Change country code")
+            console.print("  [bold]5.[/bold] Done (write & exit)")
+            console.print("  [bold]0.[/bold] Quit without saving")
+            console.print()
+
+            choice = Prompt.ask("  [cyan]Select[/cyan]").strip()
+
+            if choice == "1":
+                if action_add_network(config):
+                    dirty = True
+            elif choice == "2":
+                if action_remove_network(config):
+                    dirty = True
+            elif choice == "3":
+                if action_update_password(config):
+                    dirty = True
+            elif choice == "4":
+                if action_change_country(config):
+                    dirty = True
+            elif choice == "5":
+                break
+            elif choice == "0":
+                console.print("\n  [yellow]Discarding changes.[/yellow]")
+                dirty = False
+                return
+            else:
+                console.print("  [yellow]Invalid option.[/yellow]")
+
+        if dirty:
+            console.print()
+            console.print("  Writing config ... ", end="")
+            try:
+                write_netplan(mountpoint, config)
+                console.print("[green]OK[/green]")
+                console.print(f"  [green]Saved[/green] → [dim]{NETPLAN_PATH}[/dim] (chmod 600)")
+            except RuntimeError as exc:
+                console.print(f"[red]FAILED[/red]\n  {exc}")
+        else:
+            console.print("  [dim]No changes to write.[/dim]")
+
+    finally:
+        console.print()
+        console.print(f"  Unmounting [dim]{mountpoint}[/dim] ... ", end="")
+        unmount_sd(mountpoint)
+        console.print("[green]OK[/green]")
+
+
+def run_list(device: str, mountpoint: str) -> None:
+    """Mount SD card, list networks, unmount."""
+    if not mount_sd(device, mountpoint):
+        return
+
+    try:
+        config = read_netplan(mountpoint)
+        print_networks(config)
+        country = get_country(config)
+        console.print(f"  Regulatory domain: [yellow]{country}[/yellow]")
+    finally:
+        unmount_sd(mountpoint)
+
+
+def run_add(device: str, mountpoint: str, ssid: str, password: str) -> None:
+    """Mount SD card, add/update a network non-interactively, unmount."""
+    if not mount_sd(device, mountpoint):
+        return
+
+    try:
+        config = read_netplan(mountpoint)
+        aps = get_access_points(config)
+        existed = ssid in aps
+        aps[ssid] = {"password": password} if password else {}
+        set_access_points(config, aps)
+
+        write_netplan(mountpoint, config)
+
+        verb = "Updated" if existed else "Added"
+        console.print(f"  [green]{verb}[/green] [bold]{ssid}[/bold]  (chmod 600 applied)")
+    except RuntimeError as exc:
+        console.print(f"  [red]Error:[/red] {exc}")
+    finally:
+        unmount_sd(mountpoint)
+
+
+# ---------------------------------------------------------------------------
+# Click CLI
+# ---------------------------------------------------------------------------
+
+@click.command(context_settings={"help_option_names": ["-h", "--help"]})
+@click.option("--list", "do_list", is_flag=True, default=False,
+              help="Show current configured networks and exit.")
+@click.option("--add", nargs=2, metavar="SSID PASSWORD",
+              default=(None, None),
+              help="Non-interactively add or update a network.")
+def main(do_list: bool, add: Tuple[Optional[str], Optional[str]]) -> None:
+    """Update WiFi config on a BigBrother SD card (armbi_root label)."""
+    print_header()
+
+    # Locate the SD card
+    console.print(f"  Scanning for SD card with label [yellow]{SD_LABEL}[/yellow] ...", end=" ")
+    device = find_sd_device()
+
+    if not device:
+        console.print("[red]NOT FOUND[/red]")
+        console.print()
+        console.print(f"  [red]No block device with label '{SD_LABEL}' detected.[/red]")
+        console.print("  Insert the BigBrother SD card and try again.")
+        sys.exit(1)
+
+    console.print(f"[green]{device}[/green]")
+
+    # Create temp mountpoint
+    mountpoint = tempfile.mkdtemp(prefix="bb_sd_")
+
+    try:
+        ssid, password = add
+        if ssid is not None:
+            # --add mode
+            run_add(device, mountpoint, ssid, password or "")
+        elif do_list:
+            # --list mode
+            run_list(device, mountpoint)
+        else:
+            # Interactive mode
+            run_interactive(device, mountpoint)
+    finally:
+        # Always clean up the temp directory
+        try:
+            os.rmdir(mountpoint)
+        except OSError:
+            pass
+
+
+if __name__ == "__main__":
+    main()