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.
This commit is contained in:
n0mad1k
Executable
+516
@@ -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()
|
||||
Reference in New Issue
Block a user