v0.0.32 - Fix KB2040 Python WebUSB handshake and robust frame reassembly for get-public-key

examples/kb2040_hidden_signer_client.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+"""
+KB2040 Hidden Signer host client (WebUSB/vendor interface via pyusb).
+
+Protocol:
+  4-byte big-endian length + UTF-8 JSON body.
+
+Default target:
+  VID:PID 239a:8104
+
+Examples:
+  python3 examples/kb2040_hidden_signer_client.py status
+  python3 examples/kb2040_hidden_signer_client.py set-mnemonic --mnemonic "abandon ... about"
+  python3 examples/kb2040_hidden_signer_client.py get-public-key
+  python3 examples/kb2040_hidden_signer_client.py sign-event --event '{"kind":1,"content":"hello"}'
+"""
+
+import argparse
+import json
+import struct
+import sys
+import time
+
+try:
+    import usb.core
+    import usb.util
+except ImportError:
+    print("Missing pyusb. Install with: pip install pyusb", file=sys.stderr)
+    sys.exit(2)
+
+
+def find_vendor_interface(dev):
+    cfg = dev.get_active_configuration()
+    for itf in cfg:
+        if itf.bInterfaceClass == 0xFF:
+            ep_out = None
+            ep_in = None
+            for ep in itf:
+                direction = usb.util.endpoint_direction(ep.bEndpointAddress)
+                if direction == usb.util.ENDPOINT_OUT:
+                    ep_out = ep
+                elif direction == usb.util.ENDPOINT_IN:
+                    ep_in = ep
+            if ep_out is not None and ep_in is not None:
+                return itf, ep_out, ep_in
+    return None, None, None
+
+
+def open_device(vid: int, pid: int):
+    dev = usb.core.find(idVendor=vid, idProduct=pid)
+    if dev is None:
+        raise RuntimeError(f"Device {vid:04x}:{pid:04x} not found")
+
+    try:
+        dev.set_configuration()
+    except usb.core.USBError:
+        pass
+
+    itf, ep_out, ep_in = find_vendor_interface(dev)
+    if itf is None:
+        raise RuntimeError("Vendor/WebUSB interface (class 0xFF) not found")
+
+    itf_num = itf.bInterfaceNumber
+    try:
+        if dev.is_kernel_driver_active(itf_num):
+            dev.detach_kernel_driver(itf_num)
+    except (NotImplementedError, usb.core.USBError):
+        pass
+
+    usb.util.claim_interface(dev, itf_num)
+
+    # WebUSB-style connect handshake (CDC SET_CONTROL_LINE_STATE / request 0x22).
+    # Firmware gates IN transfers on this "connected" state.
+    try:
+        dev.ctrl_transfer(0x21, 0x22, 0x0001, itf_num, None)
+    except usb.core.USBError:
+        # Some stacks may not require/implement this; continue and let RPC I/O decide.
+        pass
+
+    return dev, itf_num, ep_out, ep_in
+
+
+def close_device(dev, itf_num: int):
+    try:
+        usb.util.release_interface(dev, itf_num)
+    except usb.core.USBError:
+        pass
+
+
+_rx_remainder = bytearray()
+
+
+def send_frame(ep_out, payload: bytes):
+    frame = struct.pack(">I", len(payload)) + payload
+    ep_out.write(frame)
+
+
+def recv_exact(ep_in, n: int, timeout_ms: int):
+    global _rx_remainder
+
+    out = bytearray()
+
+    # Consume previously over-read bytes first.
+    if _rx_remainder:
+        take = min(n, len(_rx_remainder))
+        out.extend(_rx_remainder[:take])
+        del _rx_remainder[:take]
+
+    deadline = time.time() + max(1.0, timeout_ms / 1000.0)
+    packet_size = max(64, int(getattr(ep_in, "wMaxPacketSize", 64)))
+
+    while len(out) < n and time.time() < deadline:
+        chunk = ep_in.read(packet_size, timeout=timeout_ms)
+        out.extend(bytes(chunk))
+
+    if len(out) < n:
+        raise TimeoutError(f"Timeout while reading {n} bytes (got {len(out)})")
+
+    if len(out) > n:
+        _rx_remainder.extend(out[n:])
+        del out[n:]
+
+    return bytes(out)
+
+
+def recv_frame(ep_in, timeout_ms: int):
+    hdr = recv_exact(ep_in, 4, timeout_ms)
+    body_len = struct.unpack(">I", hdr)[0]
+    if body_len > 64 * 1024:
+        raise ValueError(f"Refusing oversized frame: {body_len}")
+    body = recv_exact(ep_in, body_len, timeout_ms)
+    return body
+
+
+def rpc(dev_args, method: str, params: dict):
+    global _rx_remainder
+
+    dev, itf_num, ep_out, ep_in = open_device(dev_args.vid, dev_args.pid)
+    try:
+        _rx_remainder.clear()
+        req = {"method": method}
+        req.update(params)
+        payload = json.dumps(req, separators=(",", ":")).encode("utf-8")
+        send_frame(ep_out, payload)
+        resp = recv_frame(ep_in, dev_args.timeout_ms)
+        return json.loads(resp.decode("utf-8", errors="replace"))
+    finally:
+        close_device(dev, itf_num)
+
+
+def cmd_status(args):
+    print(json.dumps(rpc(args, "get_status", {}), indent=2))
+
+
+def cmd_set_mnemonic(args):
+    print(json.dumps(rpc(args, "set_mnemonic", {"mnemonic": args.mnemonic}), indent=2))
+
+
+def cmd_set_auto_approve(args):
+    value = args.value.lower() in ("1", "true", "yes", "on")
+    print(json.dumps(rpc(args, "set_auto_approve", {"value": value}), indent=2))
+
+
+def cmd_get_public_key(args):
+    print(json.dumps(rpc(args, "get_public_key", {}), indent=2))
+
+
+def cmd_sign_event(args):
+    event = json.loads(args.event)
+    print(json.dumps(rpc(args, "sign_event", {"event": event}), indent=2))
+
+
+def build_parser():
+    p = argparse.ArgumentParser()
+    p.add_argument("--vid", type=lambda s: int(s, 16), default=0x239A, help="USB VID in hex (default: 239A)")
+    p.add_argument("--pid", type=lambda s: int(s, 16), default=0x8104, help="USB PID in hex (default: 8104)")
+    p.add_argument("--timeout-ms", type=int, default=3000, help="USB read timeout in ms")
+
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    s = sub.add_parser("status")
+    s.set_defaults(func=cmd_status)
+
+    s = sub.add_parser("set-mnemonic")
+    s.add_argument("--mnemonic", required=True)
+    s.set_defaults(func=cmd_set_mnemonic)
+
+    s = sub.add_parser("set-auto-approve")
+    s.add_argument("--value", required=True, help="true|false")
+    s.set_defaults(func=cmd_set_auto_approve)
+
+    s = sub.add_parser("get-public-key")
+    s.set_defaults(func=cmd_get_public_key)
+
+    s = sub.add_parser("sign-event")
+    s.add_argument("--event", required=True, help="JSON object string")
+    s.set_defaults(func=cmd_sign_event)
+
+    return p
+
+
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+    try:
+        args.func(args)
+        return 0
+    except Exception as e:
+        print(f"ERROR: {e}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())