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

tests/test_kb2040_dual.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""
+Verify KB2040 composite USB exposes both HID and Vendor/WebUSB interfaces.
+
+Usage:
+  python3 tests/test_kb2040_dual.py
+
+Dependencies:
+  pip install pyusb
+
+Notes:
+- May require sudo unless udev permissions are configured.
+- Defaults to VID:PID 239a:8104 per user report.
+"""
+
+import argparse
+import sys
+import time
+
+try:
+    import usb.core
+    import usb.util
+except ImportError:
+    print("ERROR: Missing dependency. Install with: pip install pyusb")
+    sys.exit(2)
+
+
+def find_device(vid: int, pid: int):
+    return usb.core.find(idVendor=vid, idProduct=pid)
+
+
+def claim_vendor_interface(dev, vendor_itf):
+    itf_num = vendor_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)
+
+
+def release_vendor_interface(dev, vendor_itf):
+    itf_num = vendor_itf.bInterfaceNumber
+    try:
+        usb.util.release_interface(dev, itf_num)
+    except usb.core.USBError:
+        pass
+
+
+def find_endpoints(vendor_itf):
+    ep_out = None
+    ep_in = None
+    for ep in vendor_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
+    return ep_out, ep_in
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vid", default="239a", help="USB Vendor ID in hex (default: 239a)")
+    parser.add_argument("--pid", default="8104", help="USB Product ID in hex (default: 8104)")
+    parser.add_argument("--payload", default="hello", help="Payload for vendor echo test")
+    parser.add_argument("--timeout-ms", type=int, default=2000, help="Read timeout in ms")
+    args = parser.parse_args()
+
+    vid = int(args.vid, 16)
+    pid = int(args.pid, 16)
+
+    passed = 0
+    failed = 0
+
+    dev = find_device(vid, pid)
+    if dev is None:
+        print(f"[FAIL] Device {vid:04x}:{pid:04x} not found")
+        return 1
+
+    print(f"[PASS] Device {dev.idVendor:04x}:{dev.idProduct:04x} found")
+    passed += 1
+
+    try:
+        dev.set_configuration()
+    except usb.core.USBError:
+        # Usually already configured; continue.
+        pass
+
+    cfg = dev.get_active_configuration()
+    hid_itf = None
+    vendor_itf = None
+
+    for itf in cfg:
+        if itf.bInterfaceClass == 0x03:  # HID
+            hid_itf = itf
+        elif itf.bInterfaceClass == 0xFF:  # Vendor-specific
+            vendor_itf = itf
+
+    if hid_itf is not None:
+        print(f"[PASS] HID interface found (if={hid_itf.bInterfaceNumber})")
+        passed += 1
+    else:
+        print("[FAIL] HID interface not found")
+        failed += 1
+
+    if vendor_itf is not None:
+        print(f"[PASS] Vendor/WebUSB interface found (if={vendor_itf.bInterfaceNumber})")
+        passed += 1
+    else:
+        print("[FAIL] Vendor/WebUSB interface not found")
+        failed += 1
+
+    if vendor_itf is not None:
+        try:
+            claim_vendor_interface(dev, vendor_itf)
+            ep_out, ep_in = find_endpoints(vendor_itf)
+            if ep_out is None or ep_in is None:
+                print("[FAIL] Could not find bulk IN/OUT endpoints on vendor interface")
+                failed += 1
+            else:
+                payload = args.payload.encode("utf-8")
+                ep_out.write(payload)
+                time.sleep(0.05)
+                data = ep_in.read(ep_in.wMaxPacketSize, timeout=args.timeout_ms)
+                got = bytes(data)
+                expected = b"KB2040:" + payload
+                if got == expected:
+                    print(f"[PASS] Vendor echo OK: sent={payload!r} got={got!r}")
+                    passed += 1
+                else:
+                    print(f"[FAIL] Vendor echo mismatch: expected={expected!r} got={got!r}")
+                    failed += 1
+        except usb.core.USBTimeoutError:
+            print("[FAIL] Vendor read timed out")
+            failed += 1
+        except usb.core.USBError as e:
+            print(f"[FAIL] USB error during vendor test: {e}")
+            failed += 1
+        finally:
+            release_vendor_interface(dev, vendor_itf)
+
+    print(f"\nSummary: {passed} passed, {failed} failed")
+    if failed == 0:
+        print("Dual interface support is working (HID + Vendor/WebUSB).")
+        return 0
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())