switch-pico/tools/configure_switch2_wake.py

#!/usr/bin/env python3
"""Capture a Switch 2 wake record and generate the AIO firmware header."""

import argparse
import json
import re
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = (
    ROOT / "src" / "firmware" / "platform" / "pico" /
    "switch2_wake_config.h"
)
CAPTURE_PREFIX = "SWITCH2_WAKE_CAPTURE "
ADDRESS_PATTERN = re.compile(r"(?:[0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}")
EXPECTED_PREFIX = bytes.fromhex("0201061BFF53050100037E05")
EXPECTED_SUFFIX = bytes.fromhex("0F00000000000000")


class CaptureError(ValueError):
    pass


def normalize_address(value, field):
    if not isinstance(value, str) or ADDRESS_PATTERN.fullmatch(value) is None:
        raise CaptureError(f"{field} must be a six-byte Bluetooth address.")
    normalized = value.upper()
    if normalized == "00:00:00:00:00:00":
        raise CaptureError(f"{field} must not be all zeroes.")
    return normalized


def parse_capture_line(line):
    line = line.strip()
    if line.startswith(CAPTURE_PREFIX):
        line = line[len(CAPTURE_PREFIX):]
    if not line.startswith("{"):
        return None
    try:
        record = json.loads(line)
    except json.JSONDecodeError:
        return None
    if not isinstance(record, dict) or "raw_hex" not in record:
        return None
    return validate_capture(record)


def validate_capture(record):
    if record.get("version", 1) != 1:
        raise CaptureError("unsupported capture version")
    if record.get("address_type") != 0:
        raise CaptureError("capture must use a public advertiser address")
    if record.get("event_type") != 0:
        raise CaptureError("capture must be an ADV_IND wake advertisement")
    advertiser = normalize_address(record.get("advertiser"), "advertiser")
    console = normalize_address(record.get("console"), "console")
    raw_hex = record.get("raw_hex")
    if not isinstance(raw_hex, str) or re.fullmatch(r"[0-9A-Fa-f]{62}", raw_hex) is None:
        raise CaptureError("raw_hex must contain exactly 31 bytes")
    raw = bytes.fromhex(raw_hex)
    if raw[:12] != EXPECTED_PREFIX:
        raise CaptureError("capture is not a Nintendo Switch 2 advertisement")
    if raw[14:16] != b"\x00\x01":
        raise CaptureError("capture has an unexpected reconnect marker")
    if raw[16] != 0x81:
        raise CaptureError("capture does not contain the Switch 2 wake flag")
    if raw[23:] != EXPECTED_SUFFIX:
        raise CaptureError("capture has an unexpected reserved suffix")
    decoded_console = ":".join(f"{byte:02X}" for byte in reversed(raw[17:23]))
    if decoded_console != console:
        raise CaptureError("console address does not match raw wake data")
    pid = int.from_bytes(raw[12:14], "little")
    if record.get("pid", f"{pid:04X}").upper() != f"{pid:04X}":
        raise CaptureError("PID does not match raw wake data")
    payload = raw[7:].hex().upper()
    supplied_payload = record.get("esphome_payload_hex", payload)
    if supplied_payload.upper() != payload:
        raise CaptureError("manufacturer payload does not match raw wake data")
    return {
        "advertiser": advertiser,
        "console": console,
        "pid": pid,
        "raw": raw,
    }


def byte_initializer(data):
    return ", ".join(f"0x{byte:02X}" for byte in data)


def render_header(capture):
    source = bytes.fromhex(capture["advertiser"].replace(":", ""))
    return (
        "#pragma once\n\n"
        "// Generated by tools/configure_switch2_wake.py. Do not commit this file.\n"
        f"// Captured source controller: {capture['advertiser']}\n"
        f"// Target Switch 2: {capture['console']}; PID: 0x{capture['pid']:04X}\n"
        "#define SWITCH2_WAKE_CONFIGURED 1\n"
        "#define SWITCH2_WAKE_SOURCE_ADDRESS_BYTES \\\n"
        f"    {{{byte_initializer(source)}}}\n"
        "#define SWITCH2_WAKE_ADVERTISEMENT_DATA_BYTES \\\n"
        f"    {{{byte_initializer(capture['raw'])}}}\n"
    )


def capture_from_stream(stream, timeout):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        line = stream.readline()
        if not line:
            continue
        if isinstance(line, bytes):
            line = line.decode("utf-8", errors="replace")
        print(line.rstrip())
        capture = parse_capture_line(line)
        if capture is not None:
            return capture
    raise CaptureError("timed out waiting for a Switch 2 wake capture")


def pico_serial_ports():
    try:
        from serial.tools import list_ports
    except ImportError as exc:
        raise CaptureError(
            "pyserial is required; install project dependencies first"
        ) from exc
    return [port.device for port in list_ports.comports() if port.vid == 0x2E8A]


def capture_from_pico(port, timeout):
    try:
        import serial
    except ImportError as exc:
        raise CaptureError(
            "pyserial is required; install project dependencies first"
        ) from exc
    if port is None:
        ports = pico_serial_ports()
        if len(ports) != 1:
            detail = ", ".join(ports) if ports else "none"
            raise CaptureError(
                f"expected one Pico USB serial port, found {detail}; use --port"
            )
        port = ports[0]
    print(f"Waiting for automatic wake capture on {port}...")
    with serial.Serial(port, 115200, timeout=0.25, write_timeout=1) as device:
        deadline = time.monotonic() + timeout
        next_request = 0.0
        while time.monotonic() < deadline:
            now = time.monotonic()
            if now >= next_request:
                device.write(b"p\n")
                next_request = now + 1.0
            line = device.readline()
            if not line:
                continue
            text = line.decode("utf-8", errors="replace").strip()
            print(text)
            capture = parse_capture_line(text)
            if capture is not None:
                return capture
    raise CaptureError("timed out waiting for a Switch 2 wake capture")


def write_header(path, capture):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(render_header(capture), encoding="utf-8")
    temporary.replace(path)


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description=(
            "Read the one-shot Pico wake capture and configure the AIO build."
        )
    )
    parser.add_argument("--port", help="Pico USB serial port; auto-detected by default")
    parser.add_argument(
        "--input",
        type=Path,
        help="read a saved capture log instead of a live Pico",
    )
    parser.add_argument("--timeout", type=float, default=90.0)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    return parser.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    try:
        if args.timeout <= 0:
            raise CaptureError("timeout must be positive")
        if args.input is None:
            capture = capture_from_pico(args.port, args.timeout)
        else:
            with args.input.open("r", encoding="utf-8") as stream:
                capture = capture_from_stream(stream, args.timeout)
        write_header(args.output, capture)
    except (CaptureError, OSError) as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1
    print(f"Configured Switch 2 wake: {args.output}")
    print("Build and flash it with: python3 build.py --aio")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())