108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Single-verb PQ test on a fresh boot.
|
|
|
|
Usage:
|
|
python3 firmware/teensy41/test_pq_single.py --port /dev/ttyACM0 --verb kem-keygen
|
|
python3 firmware/teensy41/test_pq_single.py --port /dev/ttyACM0 --verb dsa-sign
|
|
python3 firmware/teensy41/test_pq_single.py --port /dev/ttyACM0 --read-boot
|
|
"""
|
|
import serial, struct, json, time, sys, argparse
|
|
|
|
DEFAULT_PORT = "/dev/ttyACM0"
|
|
BAUD = 115200
|
|
|
|
|
|
def send_request(ser, req, timeout=180.0):
|
|
payload = json.dumps(req).encode("utf-8")
|
|
ser.write(struct.pack(">I", len(payload)) + payload)
|
|
ser.flush()
|
|
h = b""
|
|
deadline = time.time() + timeout
|
|
while len(h) < 4 and time.time() < deadline:
|
|
c = ser.read(4 - len(h))
|
|
if c:
|
|
h += c
|
|
else:
|
|
time.sleep(0.01)
|
|
if len(h) < 4:
|
|
raise TimeoutError("hdr timeout")
|
|
n = struct.unpack(">I", h)[0]
|
|
if n == 0 or n > 65536:
|
|
raise ValueError("bad len %d" % n)
|
|
p = b""
|
|
while len(p) < n and time.time() < deadline:
|
|
c = ser.read(n - len(p))
|
|
if c:
|
|
p += c
|
|
else:
|
|
time.sleep(0.01)
|
|
if len(p) < n:
|
|
raise TimeoutError("body timeout")
|
|
return json.loads(p.decode())
|
|
|
|
|
|
def drain_boot(ser, wait=6.0):
|
|
time.sleep(wait)
|
|
boot = b""
|
|
end = time.time() + 2.0
|
|
while time.time() < end:
|
|
if ser.in_waiting:
|
|
boot += ser.read(ser.in_waiting)
|
|
else:
|
|
time.sleep(0.05)
|
|
return boot.decode("utf-8", errors="replace")
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--port", default=DEFAULT_PORT)
|
|
ap.add_argument("--verb", choices=["kem-keygen", "dsa-sign"])
|
|
ap.add_argument("--read-boot", action="store_true")
|
|
ap.add_argument("--timeout", type=float, default=180.0)
|
|
args = ap.parse_args()
|
|
|
|
ser = serial.Serial(args.port, BAUD, timeout=2.0)
|
|
boot = drain_boot(ser)
|
|
print("=== BOOT OUTPUT ===")
|
|
print(boot)
|
|
print("=== END BOOT ===", flush=True)
|
|
|
|
if args.read_boot:
|
|
ser.close()
|
|
return 0
|
|
|
|
if args.verb == "kem-keygen":
|
|
req = {"jsonrpc": "2.0", "id": 1, "method": "get_public_key",
|
|
"params": [{"algorithm": "ml-kem-768", "index": 0}]}
|
|
print(f"-> get_public_key ml-kem-768", flush=True)
|
|
try:
|
|
resp = send_request(ser, req, timeout=args.timeout)
|
|
ok = "result" in resp
|
|
print(f"<- {'OK' if ok else 'ERR'} {str(resp)[:300]}", flush=True)
|
|
ser.close()
|
|
return 0 if ok else 1
|
|
except Exception as e:
|
|
print(f"!! CRASH/TIMEOUT: {e}", flush=True)
|
|
ser.close()
|
|
return 2
|
|
|
|
if args.verb == "dsa-sign":
|
|
req = {"jsonrpc": "2.0", "id": 1, "method": "sign",
|
|
"params": ["746573742065643235353139206d657373616765",
|
|
{"algorithm": "ml-dsa-65", "index": 0}]}
|
|
print(f"-> sign ml-dsa-65", flush=True)
|
|
try:
|
|
resp = send_request(ser, req, timeout=args.timeout)
|
|
ok = "result" in resp
|
|
print(f"<- {'OK' if ok else 'ERR'} {str(resp)[:300]}", flush=True)
|
|
ser.close()
|
|
return 0 if ok else 1
|
|
except Exception as e:
|
|
print(f"!! HANG/TIMEOUT: {e}", flush=True)
|
|
ser.close()
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|