#!/usr/bin/env python3
"""Host side of the device Wi-Fi provisioning protocol.

Reuses the device's existing USB Serial/JTAG channel; no second transport is
introduced. The device speaks a versioned, framed, newline-delimited protocol:

    WFP/1 <verb> [key=<percent-encoded value> ...]
    WFP/1 OK|ERROR <verb> [key=value ...]
    WFP/1 DATA <key=value ...>          (zero or more)
    WFP/1 END

Values are percent-encoded (RFC 3986 unreserved set), so spaces, `=`, `&`, `%`
and UTF-8 all round-trip. `usb_encode` / `usb_decode` are the exact inverse of
the firmware's `usb_encode_value` / `usb_decode_value`.

This module is imported by tools/mac/device.sh and by the host tests, so the
framing lives in exactly one place.

Security:
  * A password is never accepted on the command line, so it never appears in
    process argv (`--password` does not exist).
  * `add` reads the password itself through `getpass`, which does not echo.
  * Every rendered line passes through `redact`, and the device response never
    contains the password.
"""
from __future__ import annotations

import argparse
import getpass
import sys
import time

PREFIX = "WFP/1 "
DEFAULT_TIMEOUT = 5.0
# Firmware candidate validation budget. `wifi.add` must wait longer than this so
# the device has headroom for the disconnect, the connect initiation, event
# dispatch and the serial response on top of the connection test itself. The
# firmware connection timeout is deliberately left unchanged.
FIRMWARE_CANDIDATE_TIMEOUT_S = 15.0
ADD_TIMEOUT_S = 25.0
assert ADD_TIMEOUT_S > FIRMWARE_CANDIDATE_TIMEOUT_S

_UNRESERVED = set(
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~"
)


def usb_encode(value: str) -> str:
    """Percent-encode a value exactly like the firmware."""
    out = []
    for byte in value.encode("utf-8"):
        char = chr(byte)
        if char in _UNRESERVED:
            out.append(char)
        else:
            out.append("%%%02X" % byte)
    return "".join(out)


def usb_decode(value: str) -> str:
    """Decode a percent-encoded value. Invalid escapes are preserved literally."""
    out = bytearray()
    i = 0
    while i < len(value):
        char = value[i]
        if char == "%" and i + 2 < len(value):
            try:
                out.append(int(value[i + 1:i + 3], 16))
                i += 3
                continue
            except ValueError:
                pass
        out.extend(char.encode("utf-8"))
        i += 1
    return out.decode("utf-8", errors="replace")


# Wire keys whose values are percent-encoded in responses.
_ENCODED_RESPONSE_KEYS = {"ssid"}


def build_command(verb: str, *, ssid: str | None = None, password: str | None = None) -> str:
    """Build one request line. The password is included but never logged."""
    parts = [f"{PREFIX}{verb}"]
    if ssid is not None:
        parts.append(f"ssid={usb_encode(ssid)}")
    if password is not None:
        parts.append(f"password={usb_encode(password)}")
    return " ".join(parts) + "\n"


def redact(text: str) -> str:
    """Remove any password= value from arbitrary text before display."""
    if "password=" not in text:
        return text
    out = []
    for token in text.split(" "):
        if token.startswith("password="):
            out.append("password=***")
        else:
            out.append(token)
    return " ".join(out)


def parse_response(lines: list[str]) -> dict:
    """Parse a complete framed response into a structured result."""
    result: dict = {"ok": False, "verb": "", "status": "", "data": [], "error": "",
                    "end": False, "summary": "", "fields": {}, "result": ""}
    for raw in lines:
        line = raw.rstrip("\r\n")
        if not line.startswith(PREFIX):
            continue
        body = line[len(PREFIX):]
        if body.startswith("OK ") or body.startswith("ERROR "):
            status, _, rest = body.partition(" ")
            result["status"] = status
            result["ok"] = status == "OK"
            result["verb"] = rest.split(" ")[0] if rest else ""
            result["summary"] = rest[len(result["verb"]):].strip()
            result["fields"] = _parse_fields(result["summary"])
            result["result"] = result["fields"].get("result", "")
            if status == "ERROR":
                result["error"] = rest
        elif body.startswith("DATA "):
            result["data"].append(_parse_fields(body[len("DATA "):]))
        elif body.strip() == "END":
            result["end"] = True
    return result


def _parse_fields(text: str) -> dict:
    fields: dict[str, str] = {}
    for token in text.split(" "):
        key, sep, value = token.partition("=")
        if not sep:
            continue
        fields[key] = usb_decode(value) if key in _ENCODED_RESPONSE_KEYS else value
    return fields


def read_response(port, timeout: float = DEFAULT_TIMEOUT) -> list[str]:
    """Read framed lines until END or timeout. Bounded and never blocks forever."""
    lines: list[str] = []
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        raw = port.readline()
        if not raw:
            continue
        text = raw.decode("utf-8", errors="replace")
        lines.append(text)
        if text.strip() == f"{PREFIX}END":
            break
    return lines


def send_command(port, verb: str, *, ssid: str | None = None,
                 password: str | None = None, timeout: float = DEFAULT_TIMEOUT) -> dict:
    line = build_command(verb, ssid=ssid, password=password)
    # Never print or log `line`: it may contain a password.
    port.reset_input_buffer()
    port.write(line.encode("utf-8"))
    port.flush()
    return parse_response(read_response(port, timeout))


def open_port(path: str):
    import serial  # Imported lazily so --help works without pyserial.

    return serial.Serial(path, 115200, timeout=0.2, write_timeout=2)


# Verdict -> human sentence. The wording never claims more than the device
# reported: UNAVAILABLE explicitly means no connection test happened.
_ADD_VERDICTS = {
    "VERIFIED": "Saved. Connection test: PASS.",
    "REJECTED": "Not saved. Connection test: FAIL (check SSID/password).",
    "UNAVAILABLE": "Not saved. Connection test: unavailable.",
    "INVALID": "Not saved. Invalid SSID or password.",
    "STORAGE_ERROR": "Connected but not saved (storage error).",
}


def format_result(result: dict) -> str:
    if not result.get("end"):
        return "No response from device (is the firmware running?)"
    if result["verb"] == "wifi.add":
        verdict = result.get("result", "")
        return _ADD_VERDICTS.get(verdict, f"Connection test: {verdict.lower() or 'unknown'}.")
    if not result.get("ok"):
        return f"FAILED: {redact(result.get('error') or 'unknown error')}"
    if result["verb"] == "wifi.list":
        if not result["data"]:
            return "(no known networks)"
        lines = []
        for index, row in enumerate(result["data"], start=1):
            proven = " (last connected)" if row.get("proven") == "1" else ""
            lines.append(f"{index}. {row.get('ssid', '?')}{proven}")
        return "\n".join(lines)
    if result["verb"] == "wifi.status":
        return result.get("summary") or "OK"
    return "OK"


def provision(port_path: str, ssid: str, password: str,
              timeout: float = ADD_TIMEOUT_S) -> dict:
    """Add a profile with a password supplied by the caller, not from a TTY.

    This is the non-interactive entry point used by tooling (for example the
    post-flash provisioning script). The password is passed in memory and is
    never written to argv, echoed, or logged. `redact` still guards any text a
    caller chooses to print.
    """
    with open_port(port_path) as port:
        return send_command(port, "wifi.add", ssid=ssid, password=password,
                            timeout=max(timeout, ADD_TIMEOUT_S))


def cli(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(
        description="Device Wi-Fi provisioning over USB Serial "
                    "(the password is prompted for and never placed on the command line)")
    parser.add_argument("port")
    parser.add_argument("action", choices=["status", "list", "add", "remove", "clear"])
    parser.add_argument("--ssid")
    parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
    args = parser.parse_args(argv)

    verb = {
        "status": "wifi.status",
        "list": "wifi.list",
        "add": "wifi.add",
        "remove": "wifi.remove",
        "clear": "wifi.clear",
    }[args.action]

    # Validate before touching the port so a missing SSID never opens a device.
    if args.action in ("add", "remove") and not args.ssid:
        print("SSID is required", file=sys.stderr)
        return 2

    password = None
    if args.action == "add":
        # Hidden input, read by this process. The value never enters argv.
        password = getpass.getpass("Password: ")

    with open_port(args.port) as port:
        if args.action == "add":
            result = send_command(port, verb, ssid=args.ssid, password=password,
                                  timeout=max(args.timeout, ADD_TIMEOUT_S))
        elif args.action == "remove":
            result = send_command(port, verb, ssid=args.ssid, timeout=args.timeout)
        else:
            result = send_command(port, verb, timeout=args.timeout)
    print(redact(format_result(result)))
    if args.action == "add":
        # Success means the device verified and persisted the credentials.
        return 0 if result.get("result") == "VERIFIED" else 1
    return 0 if result.get("ok") and result.get("end") else 1


if __name__ == "__main__":
    raise SystemExit(cli(sys.argv[1:]))