1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146
|
#!/usr/bin/env python3
import struct
import serial
import subprocess
from random import randbytes
from tempfile import NamedTemporaryFile
from time import sleep
class TK1:
def __init__(self, port="/dev/ttyACM0"):
self.dev = serial.Serial(port, 62500, timeout=0.1)
# Bootloader command
def getNameVersion(self):
"""Request the name and version information from the bootloader"""
cmd = bytearray([0x50, 0x01])
# print(' '.join(['{:02x}'.format(i) for i in cmd]))
self.dev.write(cmd)
rsp = self.dev.read(1 + 32)
# print(' '.join(['{:02x}'.format(i) for i in rsp]))
assert rsp[0] == 0x52
assert rsp[1] == 0x02
response = {}
response["name0"] = "".join([chr(i) for i in rsp[2:6]])
response["name1"] = "".join([chr(i) for i in rsp[6:10]])
response["version"] = int(rsp[10])
return response
# Signer app
def getPubKey(self):
cmd = bytearray([0x58, 0x01])
# print(' '.join(['{:02x}'.format(i) for i in cmd]))
self.dev.write(cmd)
rsp = self.dev.read(1 + 128)
# print(' '.join(['{:02x}'.format(i) for i in rsp]))
assert rsp[0] == 0x5B
assert rsp[1] == 0x02
return rsp[2:34]
def inSignerApp(self):
for i in range(0, 2):
try:
self.dev.write(bytes(128))
key = self.getPubKey()
# print(','.join(['0x{:02x}'.format(i) for i in key]))
# assert(key == bytearray([
# 0x67,0xb1,0x46,0x4a,0xa2,0x4f,0x65,0x93,
# 0xfe,0x67,0x1e,0xc1,0x00,0xf3,0x0e,0x85,
# 0x8c,0xdf,0x7f,0xbb,0x0b,0x46,0x86,0xbd,
# 0xf9,0xca,0x47,0xb5,0xc6,0x48,0xba,0x0f
# ]))
return True
except Exception:
pass
return False
def inBootloader(self):
for i in range(0, 2):
try:
response = self.getNameVersion()
print(response, len(response['name1']))
assert(response['name0'] == 'tk1 ')
assert(response['name1'] == 'mkdf')
assert(response['version'] == 5)
return True
except Exception:
pass
return False
def probe_state():
"""Probe the TK1 to determine if it is running the bootloader or signer"""
try:
key = TK1()
# First, probe for firmware with a getnameversion, then try to
# read the public key If this is successful we assume the
# signer app is loaded
if key.inBootloader():
return "bootloader"
if key.inSignerApp():
return "signer"
except Exception:
pass
return "unknown"
def load_signer_app():
try:
result = subprocess.run(
["../tkey-runapp", "--port", "/dev/ttyACM0", "../apps/signer/app.bin"],
timeout=10,
)
except subprocess.TimeoutExpired:
print("loader process timeout")
def do_signature():
msgf = NamedTemporaryFile()
msgf.write(randbytes(512))
msgf.flush()
try:
result = subprocess.run(
["../tkey-sign", "--port", "/dev/ttyACM0", msgf.name], timeout=1
)
except subprocess.TimeoutExpired:
print("signature process timeout")
def main():
stats = {"restarts": 0, "signatures": 0, "disconnects": 0}
while True:
state = probe_state()
print("Detected key in state: " + state)
if state == "bootloader":
load_signer_app()
sleep(2) # Give time for the app to start
stats["restarts"] += 1
elif state == "signer":
do_signature()
stats["signatures"] += 1
else:
print("Device in unknown state: reconnecting")
stats["disconnects"] += 1
sleep(1)
print(stats)
if __name__ == "__main__":
main()
|