215 lines
6.0 KiB
Python
215 lines
6.0 KiB
Python
#!/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, "nostr_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())
|