134 lines
4.3 KiB
Python
134 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""ml-dsa-65 sign test with soft-reboot diagnostic recovery.
|
|
|
|
Runs sign; on hang/timeout, triggers a Teensy soft-reboot via DTR/RTS toggle
|
|
(which preserves DMAMEM) and reads the boot output for the reject-count
|
|
diagnostic. Falls back to a hard reboot hint if DTR toggle doesn't work.
|
|
"""
|
|
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(ser, seconds=2.0):
|
|
buf = b""
|
|
end = time.time() + seconds
|
|
while time.time() < end:
|
|
if ser.in_waiting:
|
|
buf += ser.read(ser.in_waiting)
|
|
else:
|
|
time.sleep(0.05)
|
|
return buf.decode("utf-8", errors="replace")
|
|
|
|
|
|
def soft_reboot_and_read(port, wait_boot=8.0):
|
|
"""Toggle DTR/RTS to trigger Teensy bootloader/reset; preserve DMAMEM.
|
|
|
|
Teensy 4.1: setting DTR low then high does not reboot; but a BREAK condition
|
|
or the 1200-bps touch does. We try the 1200-bps open (Teensy bootloader
|
|
trigger) — note this may or may not preserve DMAMEM. As a safer soft-reboot
|
|
we instead pulse RTS/DTR which on some Teensy USB-CDC builds triggers a
|
|
watchdog reset.
|
|
"""
|
|
# Attempt 1: 1200 bps touch (Teensy bootloader reset). This is a hard reset
|
|
# but is the most reliable way to reboot a hung Teensy.
|
|
try:
|
|
s = serial.Serial(port, 1200, timeout=1.0)
|
|
s.dtr = False
|
|
s.rts = False
|
|
time.sleep(0.1)
|
|
s.dtr = True
|
|
time.sleep(0.05)
|
|
s.dtr = False
|
|
s.close()
|
|
except Exception as e:
|
|
print(f" (1200bps touch failed: {e})", flush=True)
|
|
time.sleep(1.0)
|
|
# Reopen at normal baud and drain boot
|
|
try:
|
|
ser = serial.Serial(port, BAUD, timeout=2.0)
|
|
except Exception as e:
|
|
print(f" (reopen failed: {e}; retrying in 5s)", flush=True)
|
|
time.sleep(5)
|
|
ser = serial.Serial(port, BAUD, timeout=2.0)
|
|
boot = drain(ser, seconds=wait_boot)
|
|
ser.close()
|
|
return boot
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--port", default=DEFAULT_PORT)
|
|
ap.add_argument("--timeout", type=float, default=180.0)
|
|
args = ap.parse_args()
|
|
|
|
ser = serial.Serial(args.port, BAUD, timeout=2.0)
|
|
boot = drain(ser, seconds=8.0)
|
|
print("=== BOOT OUTPUT (pre-sign) ===")
|
|
print(boot)
|
|
print("=== END BOOT ===", flush=True)
|
|
|
|
req = {"jsonrpc": "2.0", "id": 1, "method": "sign",
|
|
"params": ["746573742065643235353139206d657373616765",
|
|
{"algorithm": "ml-dsa-65", "index": 0}]}
|
|
print(f"-> sign ml-dsa-65 (timeout={args.timeout}s)", flush=True)
|
|
result = None
|
|
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)
|
|
result = "pass" if ok else "err"
|
|
except Exception as e:
|
|
print(f"!! HANG/TIMEOUT: {e}", flush=True)
|
|
result = "hang"
|
|
ser.close()
|
|
|
|
if result == "hang":
|
|
print("\n=== Attempting soft reboot to read DMAMEM diagnostic ===", flush=True)
|
|
boot = soft_reboot_and_read(args.port, wait_boot=10.0)
|
|
print("=== BOOT OUTPUT (post-hang reboot) ===")
|
|
print(boot)
|
|
print("=== END BOOT ===", flush=True)
|
|
# Extract reject count
|
|
for line in boot.splitlines():
|
|
if "REJECT COUNT" in line:
|
|
print(f"\n>>> DIAGNOSTIC: {line.strip()}", flush=True)
|
|
if "CRASH" in line:
|
|
print(f"\n>>> CRASH: {line.strip()}", flush=True)
|
|
return 2
|
|
return 0 if result == "pass" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|