strix/tools/stream_entropy.py

44 lines
1.4 KiB
Python

#!/usr/bin/env python3
"""Writes STRIX entropy to stdout continuously, for piping into test suites
(e.g. `python3 tools/stream_entropy.py | qngmeter`). Uses the firmware's `entropy <n> bin`
command (arduino/strix/API.md)."""
import argparse
import sys
import time
import serial
CHUNK = 65536
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--port", default="/dev/ttyACM0")
parser.add_argument("--bytes", type=int, default=0, help="stop after this many bytes (0 = run until killed)")
args = parser.parse_args()
device = serial.Serial(args.port, 115200, timeout=5)
time.sleep(0.5)
device.reset_input_buffer()
out = sys.stdout.buffer
sent = 0
try:
while args.bytes == 0 or sent < args.bytes:
count = CHUNK if args.bytes == 0 else min(CHUNK, args.bytes - sent)
device.write(f"entropy {count} bin\n".encode())
header = device.readline().decode(errors="replace").strip()
if header != f"BIN {count}":
sys.exit(f"device: {header or 'timeout'}")
data = device.read(count)
if len(data) != count:
sys.exit(f"short read: {len(data)} of {count} bytes")
out.write(data)
out.flush()
sent += count
except (BrokenPipeError, KeyboardInterrupt):
pass
if __name__ == "__main__":
main()