n_signer/firmware/teensy41/test_nip04.py

#!/usr/bin/env python3
"""Quick NIP-04 round-trip test for the Teensy 4.1 n_signer.

Flow:
  1. get_info  (sanity)
  2. nostr_get_public_key (role=main, role_path=m/44'1237'0'/0/0)
     -> our x-only secp256k1 pubkey (peer)
  3. nostr_nip04_encrypt [our_pub, "hello via nip04", {role:main, role_path:...}]
       -> ciphertext?iv=...
  4. nostr_nip04_decrypt [our_pub, ciphertext, {role:main, role_path:...}]
       -> should recover "hello via nip04"

Also tests NIP-44 the same way to verify the is_nip44 dispatch fix.

Usage:
    python3 firmware/teensy41/test_nip04.py [--port /dev/ttyACM0]
"""

import serial
import struct
import json
import time
import sys
import argparse

DEFAULT_PORT = "/dev/ttyACM0"
BAUD = 115200


def send_request(ser, req: dict) -> dict:
    payload = json.dumps(req).encode("utf-8")
    header = struct.pack(">I", len(payload))
    ser.write(header + payload)
    ser.flush()

    resp_header = b""
    deadline = time.time() + 30.0
    while len(resp_header) < 4 and time.time() < deadline:
        chunk = ser.read(4 - len(resp_header))
        if chunk:
            resp_header += chunk
        else:
            time.sleep(0.01)
    if len(resp_header) < 4:
        raise TimeoutError("Timeout reading response header")

    resp_len = struct.unpack(">I", resp_header)[0]
    if resp_len == 0 or resp_len > 65536:
        raise ValueError(f"Invalid response length: {resp_len}")

    resp_payload = b""
    while len(resp_payload) < resp_len:
        chunk = ser.read(resp_len - len(resp_payload))
        if chunk:
            resp_payload += chunk
        else:
            time.sleep(0.01)
    return json.loads(resp_payload.decode("utf-8"))


def call(ser, method, params=None, idc=[0]):
    idc[0] += 1
    req = {"jsonrpc": "2.0", "id": idc[0], "method": method}
    if params is not None:
        req["params"] = params
    print(f"\n→ {method} {json.dumps(params) if params else '[]'}")
    resp = send_request(ser, req)
    print(f"← {json.dumps(resp)}")
    return resp


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", default=DEFAULT_PORT)
    args = parser.parse_args()

    print(f"Connecting to {args.port}...")
    ser = serial.Serial(args.port, BAUD, timeout=2.0)
    time.sleep(4)  # wait for boot (auto-generate mnemonic)

    # Drain boot messages
    while ser.in_waiting:
        boot = ser.read(ser.in_waiting).decode("utf-8", errors="replace")
        print(f"Boot: {boot}", end="")
    print()

    ok = True

    # 1. get_info
    r = call(ser, "get_info")
    if "result" not in r:
        print("FAIL: get_info"); ok = False

    # Role + role_path selector (replaces the deprecated nostr_index).
    main_role = {"role": "main", "role_path": "m/44'/1237'/0'/0/0"}

    # 2. our nostr pubkey (x-only, 64 hex)
    r = call(ser, "nostr_get_public_key", [main_role])
    if "result" not in r:
        print("FAIL: nostr_get_public_key"); ok = False; ser.close(); return 1
    our_pub = r["result"]
    if isinstance(our_pub, dict):
        our_pub = our_pub.get("public_key", "")
    print(f"  our_pub = {our_pub}")

    # 3. NIP-04 encrypt to ourselves
    plaintext = "hello via nip04"
    r = call(ser, "nostr_nip04_encrypt", [our_pub, plaintext, main_role])
    if "result" not in r:
        print("FAIL: nostr_nip04_encrypt (this is the crash we are testing)"); ok = False
    else:
        cipher = r["result"]
        print(f"  ciphertext = {cipher}")
        # 4. NIP-04 decrypt
        r = call(ser, "nostr_nip04_decrypt", [our_pub, cipher, main_role])
        if "result" not in r:
            print("FAIL: nostr_nip04_decrypt"); ok = False
        else:
            recovered = r["result"]
            print(f"  recovered  = {recovered}")
            if recovered == plaintext:
                print("  ✅ NIP-04 round-trip OK")
            else:
                print("  ❌ NIP-04 round-trip MISMATCH")
                ok = False

    # 5. NIP-44 encrypt to ourselves (verifies the is_nip44 dispatch fix)
    plaintext44 = "hello via nip44"
    r = call(ser, "nostr_nip44_encrypt", [our_pub, plaintext44, main_role])
    if "result" not in r:
        print("FAIL: nostr_nip44_encrypt"); ok = False
    else:
        cipher44 = r["result"]
        print(f"  nip44 ciphertext = {cipher44[:60]}...")
        r = call(ser, "nostr_nip44_decrypt", [our_pub, cipher44, main_role])
        if "result" not in r:
            print("FAIL: nostr_nip44_decrypt"); ok = False
        else:
            recovered44 = r["result"]
            print(f"  nip44 recovered = {recovered44}")
            if recovered44 == plaintext44:
                print("  ✅ NIP-44 round-trip OK")
            else:
                print("  ❌ NIP-44 round-trip MISMATCH")
                ok = False

    print(f"\n{'='*50}")
    print("RESULT:", "PASS" if ok else "FAIL")
    print(f"{'='*50}")
    ser.close()
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())