Add typing information

bin/validity-led-dance

@@ -21,10 +21,9 @@ if __name__ == "__main__":
         glow_end_scan()
         sleep(0.05)
 
-    led_script = unhexlify(
+    led_script = unhexlify('39ff100000ff03000001ff002000000000ffff0000ffff0000ff03000001ff00'
-        '39ff100000ff03000001ff002000000000ffff0000ffff0000ff03000001ff00'
+                           '200000000000000000ffff0000ff03000001ff002000000000ffff0000000000'
-        '200000000000000000ffff0000ff03000001ff002000000000ffff0000000000'
+                           '0000000000000000000000000000000000000000000000000000000000000000'
-        '0000000000000000000000000000000000000000000000000000000000000000'
+                           '0000000000000000000000000000000000000000000000000000000000')
-        '0000000000000000000000000000000000000000000000000000000000')
 
     assert_status(tls.app(led_script))

validitysensor/db.py

@@ -1,15 +1,17 @@
 from binascii import hexlify
+from struct import pack, unpack
+import typing
 
 from .blobs import db_write_enable
 from .flash import call_cleanups
-from .sid import *
+from .sid import SidIdentity, sid_from_bytes
 from .tls import tls
 from .util import assert_status
 from .winbio_constants import finger_names
 
 
 class UserStorage:
-    def __init__(self, dbid, name):
+    def __init__(self, dbid: int, name: str):
         self.dbid = dbid
         self.name = name
         self.users = []
@@ -20,10 +22,10 @@ class UserStorage:
 
 
 class User:
-    def __init__(self, dbid, identity):
+    def __init__(self, dbid: int, identity: str):
         self.dbid = dbid
         self.identity = identity
-        self.fingers = []
+        self.fingers: typing.List[typing.Mapping[str, int]] = []
 
     def __repr__(self):
         return '<User: dbid=%04x identity=%s fingers=%s>' % (self.dbid, repr(
@@ -35,7 +37,7 @@ def subtype_to_string(s: int):
     return finger_name or 'Unknown'
 
 
-def parse_user_storage(rsp):
+def parse_user_storage(rsp: bytes):
     rc, = unpack('<H', rsp[:2])
 
     if rc == 0x04b3:
@@ -62,7 +64,7 @@ def parse_user_storage(rsp):
     return storage
 
 
-def parse_identity(b):
+def parse_identity(b: bytes):
     t, b = b[:4], b[4:]
     t, = unpack('<L', t)
 
@@ -127,7 +129,7 @@ class DbRecord:
 
 class Db:
     class Info:
-        def __init__(self, total, used, free, records, roots):
+        def __init__(self, total: int, used: int, free: int, records: int, roots):
             self.total = total  # partition size
             self.used = used  # used (not deleted)
             self.free = free  # unallocated space
@@ -154,7 +156,7 @@ class Db:
         rc = self.get_record_children(stg.dbid).children
         return [i['dbid'] for i in rc if i['type'] == 8]  # 8 == "data" type
 
-    def get_user(self, dbid):
+    def get_user(self, dbid: int):
         return parse_user(tls.cmd(pack('<BHHH', 0x4a, dbid, 0, 0)))
 
     def lookup_user(self, identity: str) -> typing.Optional[User]:
@@ -169,7 +171,7 @@ class Db:
         else:
             return parse_user(rsp)
 
-    def get_record_value(self, dbid):
+    def get_record_value(self, dbid: int):
         rsp = tls.cmd(pack('<BH', 0x49, dbid))
         assert_status(rsp)
 
@@ -179,7 +181,7 @@ class Db:
 
         return rec
 
-    def get_record_children(self, dbid):
+    def get_record_children(self, dbid: int):
         rsp = tls.cmd(pack('<BH', 0x46, dbid))
         assert_status(rsp)
 
@@ -193,7 +195,7 @@ class Db:
 
         return rec
 
-    def del_record(self, dbid):
+    def del_record(self, dbid: int):
         assert_status(tls.cmd(pack('<BH', 0x48, dbid)))
 
     def db_info(self):
@@ -208,7 +210,7 @@ class Db:
 
         return Db.Info(total, used, free, records, roots)
 
-    def new_record(self, parent, typ, storage, data):
+    def new_record(self, parent: int, typ: int, storage: int, data: bytes):
         self.db_info()  # TODO check free space, compact the partition when out of storage
         assert_status(tls.cmd(db_write_enable))
         try:
@@ -219,21 +221,21 @@ class Db:
             call_cleanups()
         return recid
 
-    def new_user(self, identity):
+    def new_user(self, identity: str):
         data = identity_to_bytes(identity)
 
         stg = self.get_user_storage(name='StgWindsor')
         rec = self.new_record(stg.dbid, 5, stg.dbid, data)
         return rec
 
-    def new_finger(self, userid, template):
+    def new_finger(self, userid: int, template: bytes):
         stg = self.get_user_storage(name='StgWindsor')
         # We ask to create an object of type 0xb,
         # but because of the magical `db_write_enable` in the new_record() it ends up being 0x6
         rec = self.new_record(userid, 0xb, stg.dbid, template)
         return rec
 
-    def new_data(self, parent, data):
+    def new_data(self, parent: int, data: bytes):
         stg = self.get_user_storage(name='StgWindsor')
         data = pack('<HH', 1, len(data)) + data
         rec = self.new_record(parent, 0x8, stg.dbid, data)

validitysensor/flash.py

@@ -1,28 +1,18 @@
+import typing
 from struct import pack, unpack
 
 from .blobs import db_write_enable
-from .hw_tables import flash_ic_table_lookup
+from .hw_tables import flash_ic_table_lookup, FlashIcInfo
 from .tls import tls
 from .util import assert_status, unhex
 
 
-class FlashInfo:
-    def __init__(self, ic, blocks, unknown0, blocksize, unknown1, partitions):
-        self.ic, self.blocks, self.unknown0, self.blocksize, self.unknown1, self.partitions = \
-            ic, blocks, unknown0, blocksize, unknown1, partitions
-
-    def __repr__(self):
-        return 'FlashInfo(%s, 0x%x, 0x%x, 0x%x, 0x%x, %s)' % (repr(
-            self.ic), self.blocks, self.unknown0, self.blocksize, self.unknown1,
-                                                              repr(self.partitions))
-
-
 # type 01 is for the firmware - written blocks are decrypted on the fly using fw key
 # access lvl:
 # 2 -- write only
 # 7 -- can read/write even without TLS
 class PartitionInfo:
-    def __init__(self, id, type, access_lvl, offset, size):
+    def __init__(self, id: int, type: int, access_lvl: int, offset: int, size: int):
         self.id, self.type, self.access_lvl, self.offset, self.size = id, type, access_lvl, offset, size
 
     def __repr__(self):
@@ -30,6 +20,22 @@ class PartitionInfo:
             self.id, self.type, self.access_lvl, self.offset, self.size)
 
 
+class FlashInfo:
+    def __init__(self, ic: FlashIcInfo, blocks: int, unknown0: int, blocksize: int, unknown1: int,
+                 partitions: typing.Sequence[PartitionInfo]):
+        self.ic = ic
+        self.blocks = blocks
+        self.unknown0 = unknown0
+        self.blocksize = blocksize
+        self.unknown1 = unknown1
+        self.partitions = partitions
+
+    def __repr__(self):
+        return 'FlashInfo(%s, 0x%x, 0x%x, 0x%x, 0x%x, %s)' % (repr(
+            self.ic), self.blocks, self.unknown0, self.blocksize, self.unknown1,
+                                                              repr(self.partitions))
+
+
 def get_flash_info():
     rsp = tls.cmd(unhex('3e'))
     assert_status(rsp)
@@ -65,7 +71,7 @@ def get_flash_info():
 #      0200 6637     0100    0c00    f0220200
 #      0100 2556     0000    0100    60040000
 class ModuleInfo:
-    def __init__(self, type, subtype, major, minor, size):
+    def __init__(self, type: int, subtype: int, major: int, minor: int, size: int):
         self.type, self.subtype, self.major, self.minor, self.size = type, subtype, major, minor, size
 
     def __repr__(self):
@@ -74,7 +80,7 @@ class ModuleInfo:
 
 
 class FirmwareInfo:
-    def __init__(self, major, minor, buildtime, modules):
+    def __init__(self, major: int, minor: int, buildtime: int, modules: typing.Sequence[ModuleInfo]):
         self.major, self.minor, self.buildtime, self.modules = major, minor, buildtime, modules
 
     def __repr__(self):
@@ -82,7 +88,7 @@ class FirmwareInfo:
                                                  repr(self.modules))
 
 
-def get_fw_info(partition):
+def get_fw_info(partition: int):
     rsp = tls.cmd(pack('<BB', 0x43, partition))
 
     # don't want to throw exception here - it is normal not to have FW when we're about to upload it
@@ -113,7 +119,7 @@ def call_cleanups():
     assert_status(rsp)
 
 
-def erase_flash(partition):
+def erase_flash(partition: int):
     assert_status(tls.cmd(db_write_enable))
     try:
         assert_status(tls.cmd(pack('<BB', 0x3f, partition)))
@@ -121,7 +127,7 @@ def erase_flash(partition):
         call_cleanups()
 
 
-def read_flash(partition, addr, size):
+def read_flash(partition: int, addr: int, size: int) -> bytes:
     cmd = pack('<BBBHLL', 0x40, partition, 1, 0, addr, size)
     rsp = tls.cmd(cmd)
     assert_status(rsp)
@@ -130,7 +136,7 @@ def read_flash(partition, addr, size):
     return rsp[8:8 + sz]
 
 
-def write_flash(partition, addr, buf):
+def write_flash(partition: int, addr: int, buf: bytes):
     tls.cmd(db_write_enable)
     cmd = pack('<BBBHLL', 0x41, partition, 1, 0, addr, len(buf)) + buf
     try:
@@ -140,7 +146,7 @@ def write_flash(partition, addr, buf):
         call_cleanups()
 
 
-def write_flash_all(partition, ptr, buf):
+def write_flash_all(partition: int, ptr: int, buf: bytes):
     bs = 0x1000
     while len(buf) > 0:
         chunk, buf = buf[:bs], buf[bs:]
@@ -148,13 +154,13 @@ def write_flash_all(partition, ptr, buf):
         ptr += len(chunk)
 
 
-def read_flash_all(partition, start, size):
+def read_flash_all(partition: int, start: int, size: int):
     bs = 0x1000
     blocks = [read_flash(partition, addr, bs) for addr in range(start, start + size, bs)]
     return b''.join(blocks)[:size]
 
 
-def write_fw_signature(partition, signature):
+def write_fw_signature(partition: int, signature: bytes):
     rsp = tls.cmd(pack('<BBxH', 0x42, partition, len(signature)) + signature)
     assert_status(rsp)
 

validitysensor/hw_tables.py

@@ -1,5 +1,5 @@
 class DeviceInfo:
-    def __init__(self, major, type, version, version_mask, name):
+    def __init__(self, major: int, type: int, version: int, version_mask: int, name: str):
         self.major, self.type, self.version, self.version_mask, self.name = major, type, version, version_mask, name
 
     def __repr__(self):
@@ -428,7 +428,7 @@ dev_info_table = [
 ]
 
 
-def dev_info_lookup(major, ver):
+def dev_info_lookup(major: int, ver: int):
     fuzzy_match = None
 
     for i in dev_info_table:
@@ -446,8 +446,8 @@ def dev_info_lookup(major, ver):
 
 
 class FlashIcInfo:
-    def __init__(self, name, size, f18, jid0, jid1, f1b, f1c, f1e, secror_size, sector_erase_cmd,
+    def __init__(self, name: str, size: int, f18: int, jid0: int, jid1: int, f1b: int, f1c: int,
-                 f25, f26):
+                 f1e: int, secror_size: int, sector_erase_cmd: int, f25: int, f26: int):
         self.name, self.size, self.f18, self.jid0, self.jid1, self.f1b, self.f1c, self.f1e, self.secror_size, self.sector_erase_cmd, self.f25, self.f26 = name, size, f18, jid0, jid1, f1b, f1c, f1e, secror_size, sector_erase_cmd, f25, f26
 
     def __repr__(self):
@@ -480,7 +480,7 @@ flash_ic_table = [
 ]
 
 
-def flash_ic_table_lookup(jedec_id0, jedec_id1, size):
+def flash_ic_table_lookup(jedec_id0: int, jedec_id1: int, size: int):
     for i in flash_ic_table:
         if i.jid0 == jedec_id0 and i.jid1 == jedec_id1 and i.size == size:
             return i

validitysensor/init.py

@@ -48,6 +48,6 @@ def open():
     open_common()
 
 
-def open_devpath(busnum, address):
+def open_devpath(busnum: int, address: int):
     usb.open_devpath(busnum, address)
     open_common()

validitysensor/init_db.py

@@ -4,7 +4,7 @@ from struct import pack, unpack
 from .db import db
 
 
-def machine_id_rec_value(b):
+def machine_id_rec_value(b: str):
     b = b.encode('utf-16le')
     b = b + b'\0' * (0x94 - len(b))
     return pack('<HH', 0x102, len(b)) + b  # 0x102 = Machine ID/GUID?

validitysensor/init_flash.py

@@ -1,6 +1,7 @@
 import hmac
 import logging
 import os
+import typing
 from binascii import unhexlify
 from hashlib import sha256
 from struct import pack, unpack
@@ -11,7 +12,8 @@ from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
 from .blobs import reset_blob
-from .flash import write_flash, erase_flash, call_cleanups, PartitionInfo, get_flash_info
+from .flash import write_flash, erase_flash, call_cleanups, PartitionInfo, get_flash_info, FlashInfo
+from .hw_tables import FlashIcInfo
 from .sensor import reboot, RomInfo
 from .tls import tls, hs_key, crt_hardcoded
 from .usb import usb
@@ -45,7 +47,7 @@ def get_partition_signature():
     return partition_signature
 
 
-def with_hdr(id, buf):
+def with_hdr(id: int, buf: bytes):
     return pack('<HH', id, len(buf)) + buf
 
 
@@ -78,17 +80,17 @@ def make_cert(client_public):
     return msg
 
 
-def serialize_flash_params(ic):
+def serialize_flash_params(ic: FlashIcInfo):
     return pack('<LLxxBx', ic.size, ic.secror_size, ic.sector_erase_cmd)
 
 
-def serialize_partition(p):
+def serialize_partition(p: PartitionInfo):
     b = pack('<BBHLL', p.id, p.type, p.access_lvl, p.offset, p.size)
     b = b + b'\0' * 4 + sha256(b).digest()
     return b
 
 
-def partition_flash(info, layout, client_public):
+def partition_flash(info: FlashInfo, layout: typing.List[PartitionInfo], client_public):
     logging.info('Detected Flash IC: %s, %d bytes' % (info.ic.name, info.ic.size))
 
     cmd = unhex('4f 0000 0000')

validitysensor/sensor.py

@@ -60,14 +60,14 @@ def get_prg_status2():
     return tls.app(unhexlify('5100200000'))
 
 
-def read_hw_reg32(addr):
+def read_hw_reg32(addr: int):
     rsp = tls.cmd(pack('<BLB', 7, addr, 4))
     assert_status(rsp)
     rsp, = unpack('<L', rsp[2:])
     return rsp
 
 
-def write_hw_reg32(addr, val):
+def write_hw_reg32(addr: int, val: int):
     rsp = tls.cmd(pack('<BLLB', 8, addr, val, 4))
     assert_status(rsp)
 
@@ -95,7 +95,7 @@ class RomInfo:
         rsp = rsp[2:]
         return cls(*unpack('<LLBBxBxxxB', rsp[0:0x10]))
 
-    def __init__(self, timestamp, build, major, minor, product, u1):
+    def __init__(self, timestamp: int, build: int, major: int, minor: int, product: int, u1: int):
         self.timestamp, self.build, self.major, self.minor, self.product, self.u1 = timestamp, build, major, minor, product, u1
 
     def __repr__(self):
@@ -124,7 +124,7 @@ def identify_sensor():
 #      d0000000 9400 0e00 0700 0080 07000000 2b23203c2d182e1e30182e1c321d341d341e321c301e1e241e201f201d1c321a301e1c211e21341f1e202024201f
 #      6c010000 1400 0e00 0f00 0080 05550007 7701002805720000080100020811e107
 #      88010000 0c00 0e00 1200 0080 07000000 7002 7800 7002 7800
-def get_factory_bits(tag):
+def get_factory_bits(tag: int):
     rsp = tls.cmd(pack('<B H HL', 0x6f, tag, 0, 0))
     assert_status(rsp)
     rsp = rsp[2:]
@@ -177,15 +177,16 @@ def bitpack(b):
 
 
 class Line:
-    mask = None
+    def __init__(self):
-    flags = None
+        self.mask: typing.Optional[int] = None
-    data = None
+        self.flags: typing.Optional[int] = None
-    v0 = 0
+        self.data: typing.Optional[bytes] = None
-    v1 = 0
+        self.v0 = 0
-    v2 = 0
+        self.v1 = 0
+        self.v2 = 0
 
 
-def clip(x):
+def clip(x: int):
     if x < -128:
         x = -128
 
@@ -195,19 +196,19 @@ def clip(x):
     return x & 0xff
 
 
-def scale(x):
+def scale(x: int):
     x -= 0x80
     x = int(x * 10 / 0x22)  # TODO: scaling factor depends on a device
     return clip(x)
 
 
-def add(l, r):
+def add(l: int, r: int):
     # Make signed
     l, r = unpack('bb', pack('BB', l, r))
     return clip(l + r)
 
 
-def chunks(b, l):
+def chunks(b: bytes, l: int):
     return [b[i:i + l] for i in range(0, len(b), l)]
 
 
@@ -266,7 +267,7 @@ class Sensor:
 
     # This is the exact logic from the DLL.
     # If it looks broken that was probably intended.
-    def patch_timeslot_table(self, b, inc_address, mult):
+    def patch_timeslot_table(self, b: bytes, inc_address: bool, mult: int):
         b = bytearray(b)
         i = 0
         while i + 3 < len(b):
@@ -290,7 +291,7 @@ class Sensor:
 
         return bytes(b)
 
-    def patch_timeslot_again(self, b):
+    def patch_timeslot_again(self, b: bytes):
         b = bytearray(b)
 
         pc = 0
@@ -336,7 +337,7 @@ class Sensor:
 
         return bytes(b)
 
-    def average(self, raw_calib_data):
+    def average(self, raw_calib_data: bytes):
         frame_size = self.lines_per_frame * self.bytes_per_line
         interleave_lines = self.lines_per_frame // self.type_info.lines_per_calibration_data  # 2, TODO: algo is quite different when it is 1
         input_frames = self.calibration_frames
@@ -371,7 +372,7 @@ class Sensor:
 
         return frame
 
-    def process_calibration_results(self, cooked_data):
+    def process_calibration_results(self, cooked_data: bytes):
         frame = chunks(cooked_data, self.bytes_per_line)
 
         # apply scaling factors
@@ -407,7 +408,7 @@ class Sensor:
 
         return key_line
 
-    def line_update_type_1(self, mode, chunks):
+    def line_update_type_1(self, mode: CaptureMode, chunks: typing.List[typing.List[typing.Union[int, bytes]]]):
         for c in chunks:
             # Timeslot Table 2D
             if c[0] == 0x34:
@@ -451,7 +452,7 @@ class Sensor:
         # ---------------- Interleave ---------------
         chunks += [[0x44, pack('<L', 1)]]
 
-        lines = []
+        lines: typing.List[Line] = []
         cnt = 2  # TODO figure out why 2
 
         l = Line()
@@ -510,7 +511,7 @@ class Sensor:
 
         return chunks
 
-    def line_update_type_2(self, mode, chunks):
+    def line_update_type_2(self, mode: CaptureMode, chunks: typing.List[typing.List[typing.Union[int, bytes]]]):
         for c in chunks:
             # patch the 2D params.
             # The following is only needed on some rom versions below 6.5 as reported by cmd_01
@@ -597,7 +598,7 @@ class Sensor:
 
         return chunks
 
-    def build_cmd_02(self, mode):
+    def build_cmd_02(self, mode: CaptureMode):
         chunks = list(prg.split_chunks(self.hardcoded_prog))
 
         if self.rom_info.product != 0x30:
@@ -615,7 +616,7 @@ class Sensor:
 
         return pack('<BHH', 2, self.bytes_per_line, req_lines) + prg.merge_chunks(chunks)
 
-    def persist_clean_slate(self, clean_slate):
+    def persist_clean_slate(self, clean_slate: bytes):
         start = read_flash(6, 0, 0x44)
 
         if start != b'\xff' * 0x44:
@@ -688,7 +689,7 @@ class Sensor:
     def cancel(self):
         usb.cancel = True
 
-    def capture(self, mode):
+    def capture(self, mode: CaptureMode) -> typing.Tuple[int, int, int, int]:
         try:
             assert_status(tls.app(self.build_cmd_02(mode)))
 
@@ -733,7 +734,7 @@ class Sensor:
         finally:
             tls.app(unhexlify('04'))  # capture stop if still running, cleanup
 
-    def enrollment_update_start(self, key):
+    def enrollment_update_start(self, key: int) -> int:
         rsp = tls.app(pack('<BLL', 0x68, key, 0))
         assert_status(rsp)
         new_key, = unpack('<L', rsp[2:])
@@ -749,7 +750,7 @@ class Sensor:
         assert_status(tls.app(pack('<BL', 0x69, 0)))
 
     # Generates interrupt
-    def enrollment_update(self, prev):
+    def enrollment_update(self, prev: bytes):
         write_enable()
         try:
             rsp = tls.app(b'\x6b' + prev)
@@ -759,7 +760,7 @@ class Sensor:
 
         return rsp[2:]
 
-    def append_new_image(self, prev):
+    def append_new_image(self, prev: bytes):
         self.enrollment_update(prev)
 
         usb.wait_int()
@@ -790,7 +791,7 @@ class Sensor:
 
         return header, template, tid
 
-    def make_finger_data(self, subtype, template, tid):
+    def make_finger_data(self, subtype: int, template: bytes, tid: bytes):
         template = pack('<HH', 1, len(template)) + template
         tid = pack('<HH', 2, len(tid)) + tid
 
@@ -804,7 +805,7 @@ class Sensor:
     # TODO: Better typing information needed.
     def enroll(self, identity: SidIdentity, subtype: int,
                update_cb: typing.Callable[[typing.Any, typing.Optional[Exception]], None]):
-        def do_create_finger(final_template, tid):
+        def do_create_finger(final_template: bytes, tid: bytes):
             tinfo = self.make_finger_data(subtype, final_template, tid)
 
             usr = db.lookup_user(identity)
@@ -849,7 +850,7 @@ class Sensor:
         self.enrollment_update_end()  # done twice for some reason
         return do_create_finger(template, tid)
 
-    def parse_dict(self, x):
+    def parse_dict(self, x: bytes):
         rc = {}
 
         while len(x) > 0:
@@ -890,7 +891,7 @@ class Sensor:
             # cleanup, ignore any errors
             tls.app(unhexlify('6200000000'))
 
-    def identify(self, update_cb):
+    def identify(self, update_cb: typing.Callable[[Exception], None]):
         while True:
             try:
                 glow_start_scan()
@@ -908,7 +909,7 @@ class Sensor:
 
         return self.match_finger()
 
-    def get_finger_blobs(self, usrid, subtype):
+    def get_finger_blobs(self, usrid: int, subtype: int):
         usr = db.get_user(usrid)
         fingerids = [f['dbid'] for f in usr.fingers if f['subtype'] == subtype]
 

validitysensor/table_types.py

@@ -1,18 +1,20 @@
+import typing
 from binascii import hexlify, unhexlify
 
 
 class SensorTypeInfo:
-    table = []
+    table: typing.List["SensorTypeInfo"] = []
 
     @classmethod
-    def get_by_type(cls, sensor_type):
+    def get_by_type(cls, sensor_type: int) -> typing.Optional["SensorTypeInfo"]:
+        # noinspection PyUnresolvedReferences
         from . import generated_tables
         for i in cls.table:
             if i.sensor_type == sensor_type:
                 return i
 
-    def __init__(self, sensor_type, bytes_per_line, repeat_multiplier, lines_per_calibration_data,
+    def __init__(self, sensor_type: int, bytes_per_line: int, repeat_multiplier: int, lines_per_calibration_data: int,
-                 line_width, calibration_blob):
+                 line_width: int, calibration_blob: str):
         self.sensor_type = sensor_type
         self.repeat_multiplier = repeat_multiplier
         self.lines_per_calibration_data = lines_per_calibration_data
@@ -50,10 +52,11 @@ def metric(i, rominfo):
 
 
 class SensorCaptureProg:
-    table = []
+    table: typing.List["SensorCaptureProg"] = []
 
     @classmethod
-    def get(cls, rominfo, sensor_type, a0, a1):
+    def get(cls, rominfo, sensor_type: int, a0: int, a1: int):
+        # noinspection PyUnresolvedReferences
         from . import generated_tables
 
         maximum = 0
@@ -80,7 +83,8 @@ class SensorCaptureProg:
         if found is not None:
             return b''.join(found.blobs)
 
-    def __init__(self, major, minor, build, u1, dev_type, a0, a1, blobs):
+    def __init__(self, major: int, minor: int, build: int, u1: int, dev_type: int, a0: int, a1: int,
+                 blobs: typing.Sequence[str]):
         self.major = major
         self.minor = minor
         self.build = build
@@ -88,8 +92,7 @@ class SensorCaptureProg:
         self.dev_type = dev_type
         self.a0 = a0
         self.a1 = a1
-        blobs = [unhexlify(b) for b in blobs]
+        self.blobs = [unhexlify(b) for b in blobs]
-        self.blobs = blobs
 
     def __repr__(self):
         blobs = [hexlify(b).decode() for b in self.blobs]

validitysensor/timeslot.py

@@ -1,5 +1,6 @@
 from binascii import hexlify
 from struct import unpack, pack
+import typing
 
 codes = {}
 codes[0x0] = "No Operation"
@@ -98,7 +99,7 @@ insn_to_string = [
 ]
 
 
-def decode_insn(b):
+def decode_insn(b: bytes):
     if b[0] == 0:
         return 0, 1
     elif b[0] == 1:
@@ -135,7 +136,7 @@ def decode_insn(b):
         raise Exception('Unhandled instruction %02x' % b)
 
 
-def disassm_timeslot_table(b, off):
+def disassm_timeslot_table(b: bytes, off: int):
     pc = off
     while len(b) > 0:
         op, sz, *operands = decode_insn(b)
@@ -147,7 +148,7 @@ def disassm_timeslot_table(b, off):
         pc += sz
 
 
-def find_nth_insn(b, opcode, n):
+def find_nth_insn(b: bytes, opcode: int, n: int):
     pc = 0
     while len(b) > 0:
         op, sz, *_ = decode_insn(b)
@@ -164,7 +165,7 @@ def find_nth_insn(b, opcode, n):
         pc += sz
 
 
-def find_nth_regwrite(b, reg_addr, n):
+def find_nth_regwrite(b: bytes, reg_addr: int, n: int):
     pc = 0
     while len(b) > 0:
         op, sz, *operands = decode_insn(b)
@@ -183,20 +184,20 @@ def find_nth_regwrite(b, reg_addr, n):
         pc += sz
 
 
-def split_chunks(b):
+def split_chunks(b: bytes) -> typing.Generator[typing.List[typing.Union[int, bytes]], None, None]:
     while len(b) > 0:
         (typ, sz), b = unpack('<HH', b[:4]), b[4:]
         p, b = b[:sz], b[sz:]
         yield [typ, p]
 
 
-def merge_chunks(cs):
+def merge_chunks(cs: typing.List[typing.List[typing.Union[int, bytes]]]):
     return b''.join([pack('<HH', key, len(val)) + val for key, val in cs])
 
 
 def dump_all(b: bytes):
-    ts = None  # type: typing.Optional[bytes]
+    ts: typing.Optional[bytes] = None
-    ts_off = None  # type: typing.Optional[int]
+    ts_off: typing.Optional[int] = None
     while len(b) > 0:
         (typ, sz), b = unpack('<HH', b[:4]), b[4:]
         p, b = b[:sz], b[sz:]

validitysensor/tls.py

@@ -2,6 +2,7 @@ import hmac
 import logging
 import os
 import pickle
+import typing
 from binascii import hexlify, unhexlify
 from hashlib import sha256
 from struct import pack, unpack
@@ -12,7 +13,7 @@ from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.asymmetric.utils import Prehashed
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
-from .usb import usb
+from .usb import usb, Usb
 from .util import unhex
 
 password_hardcoded = unhexlify('717cd72d0962bc4a2846138dbb2c24192512a76407065f383846139d4bec2033')
@@ -34,7 +35,7 @@ fff000000000000000000000000000000000000000000000000000000000000000000000000
 crypto_backend = default_backend()
 
 
-def prf(secret, seed, length):
+def prf(secret: bytes, seed: bytes, length: int):
     n = (length + 0x20 - 1) // 0x20
 
     res = b''
@@ -55,15 +56,15 @@ def hs_key():
     return int(hs_key[::-1].hex(), 16)
 
 
-def with_2bytes_size(chunk):
+def with_2bytes_size(chunk: bytes):
     return pack('>H', len(chunk)) + chunk
 
 
-def with_3bytes_size(chunk):
+def with_3bytes_size(chunk: bytes):
     return pack('>BH', len(chunk) >> 16, len(chunk)) + chunk
 
 
-def with_1byte_size(chunk):
+def with_1byte_size(chunk: bytes):
     return pack('>B', len(chunk)) + chunk
 
 
@@ -76,18 +77,18 @@ def to_bytes(n):
     return b
 
 
-def pad(b):
+def pad(b: bytes):
     l = 16 - (len(b) % 16)
     return b + bytes([l - 1]) * l
 
 
-def unpad(b):
+def unpad(b: bytes):
     return b[:-1 - b[-1]]
 
 
 # TODO assert the right state transitions
 class Tls:
-    def __init__(self, usb):
+    def __init__(self, usb: Usb):
         self.usb = usb
         self.reset()
         try:
@@ -107,7 +108,7 @@ class Tls:
         self.secure_tx = False
 
     # Info about the host computer
-    def set_hwkey(self, product_name, serial_number):
+    def set_hwkey(self, product_name: str, serial_number: str):
         hw_key = bytes(product_name, 'ascii') + b'\0' + \
                  bytes(serial_number, 'ascii') + b'\0'
 
@@ -116,7 +117,7 @@ class Tls:
         self.psk_validation_key = prf(self.psk_encryption_key, b'GWK_SIGN' + gwk_sign_hardcoded,
                                       0x20)
 
-    def cmd(self, cmd):
+    def cmd(self, cmd: typing.Union[bytes, typing.Callable[[], bytes]]):
         if self.secure_rx and self.secure_tx:
             rsp = self.app(cmd)
         else:
@@ -142,15 +143,15 @@ class Tls:
 
         self.parse_tls_response(rsp)
 
-    def trace(self, s):
+    def trace(self, s: str):
         if self.trace_enabled:
             logging.debug(s)
 
-    def app(self, b):
+    def app(self, b: typing.Union[bytes, typing.Callable[[], bytes]]):
         b = b() if callable(b) else b
         return self.parse_tls_response(self.usb.cmd(self.make_app_data(b)))
 
-    def update_neg(self, b):
+    def update_neg(self, b: bytes):
         self.handshake_hash.update(b)
 
     def make_keys(self):
@@ -187,7 +188,7 @@ class Tls:
             self.secure_rx = d['secure_rx']
             self.secure_tx = d['secure_tx']
 
-    def decrypt(self, c):
+    def decrypt(self, c: bytes):
         iv, c = c[:0x10], c[0x10:]
         cipher = Cipher(algorithms.AES(self.decryption_key), modes.CBC(iv), backend=crypto_backend)
         decryptor = cipher.decryptor()
@@ -195,7 +196,7 @@ class Tls:
         m = unpad(m)
         return m
 
-    def encrypt(self, b):
+    def encrypt(self, b: bytes):
         #iv = unhexlify('454849acdd075174d6b9e713a957c2e7')
         iv = os.urandom(0x10)
         cipher = Cipher(algorithms.AES(self.encryption_key), modes.CBC(iv), backend=crypto_backend)
@@ -204,7 +205,7 @@ class Tls:
         c = encryptor.update(b) + encryptor.finalize()
         return iv + c
 
-    def validate(self, t, b):
+    def validate(self, t: int, b: bytes):
         b, hs = b[:-0x20], b[-0x20:]
 
         hdr = pack('>BBBH', t, 3, 3, len(b))
@@ -216,7 +217,7 @@ class Tls:
         self.trace('<tls< %02x: %s' % (t, hexlify(b).decode()))
         return b
 
-    def sign(self, t, b):
+    def sign(self, t: int, b: bytes):
         self.trace('>tls> %02x: %s' % (t, hexlify(b).decode()))
 
         hdr = pack('>BBBH', t, 3, 3, len(b))
@@ -240,7 +241,7 @@ class Tls:
         cert = pack('>BH', 0, len(self.tls_cert)) + cert  # same
         return self.with_neg_hdr(0x0b, cert)
 
-    def with_neg_hdr(self, t, b):
+    def with_neg_hdr(self, t: int, b: bytes):
         b = pack('>B', t) + with_3bytes_size(b)
         self.update_neg(b)
         return b
@@ -254,7 +255,7 @@ class Tls:
         b = self.priv_key.sign(buf, ec.ECDSA(Prehashed(hashes.SHA256())))
         return self.with_neg_hdr(0x0f, b)
 
-    def handle_server_hello(self, p):
+    def handle_server_hello(self, p: bytes):
         if p[:2] != unhexlify('0303'):
             raise Exception('unexpected TLS version %s' % hexlify(p[:2]).decode())
 
@@ -278,7 +279,7 @@ class Tls:
         if p != b'':
             raise Exception('Not expecting any more data')
 
-    def handle_cert_req(self, p):
+    def handle_cert_req(self, p: bytes):
         (sign_and_hash_algo, ), p = unpack('>H', p[:2]), p[2:]
         if sign_and_hash_algo != 0x140:
             raise Exception(
@@ -292,24 +293,24 @@ class Tls:
         if p != b'':
             raise Exception('Not expecting any more data')
 
-    def handle_server_hello_done(self, p):
+    def handle_server_hello_done(self, p: bytes):
         if p != b'':
             raise Exception('Not expecting any body for "server hello done" pkt: %s' %
                             hexlify(p).decode())
 
-    def handle_finish(self, b):
+    def handle_finish(self, b: bytes):
         hs_hash = self.handshake_hash.copy().digest()
         verify_data = prf(self.master_secret, b'server finished' + hs_hash, 0xc)
         if verify_data != b:
             raise Exception('Final handshake check failed')
 
-    def handle_app_data(self, b):
+    def handle_app_data(self, b: bytes):
         if not self.secure_rx:
             raise Exception('App payload before secure connection established')
 
         return self.validate(0x17, self.decrypt(b))
 
-    def handle_handshake(self, handshake):
+    def handle_handshake(self, handshake: bytes) -> None:
         if self.secure_rx:
             handshake = self.validate(0x16, self.decrypt(handshake))
 
@@ -335,7 +336,7 @@ class Tls:
 
             self.update_neg(hdr + p)
 
-    def parse_tls_response(self, rsp):
+    def parse_tls_response(self, rsp: bytes):
         app_data = b''
 
         while len(rsp) > 0:
@@ -366,7 +367,7 @@ class Tls:
 
         return app_data
 
-    def make_app_data(self, b):
+    def make_app_data(self, b: bytes):
         if not self.secure_tx:
             raise Exception('App payload before secure connection established')
 
@@ -374,7 +375,7 @@ class Tls:
 
         return unhexlify('170303') + with_2bytes_size(b)
 
-    def make_handshake(self, b):
+    def make_handshake(self, b: bytes):
         if self.secure_tx:
             b = self.encrypt(self.sign(0x16, b))
 
@@ -404,10 +405,10 @@ class Tls:
 
         return self.with_neg_hdr(0x01, h)
 
-    def make_ext(self, id, b):
+    def make_ext(self, id: int, b: bytes):
         return pack('>H', id) + with_2bytes_size(b)
 
-    def parse_tls_flash(self, reply):
+    def parse_tls_flash(self, reply: bytes):
         while len(reply) > 0:
             hdr, reply = reply[:4], reply[4:]
             hs, reply = reply[:0x20], reply[0x20:]
@@ -440,7 +441,7 @@ class Tls:
             else:
                 self.trace('unhandled block id %04x (%d bytes): %s' % (id, sz, hexlify(body)))
 
-    def make_tls_flash_block(self, id, body):
+    def make_tls_flash_block(self, id: int, body: bytes):
         m = sha256()
         m.update(body)
         hdr = pack('<HH', id, len(body))
@@ -457,16 +458,16 @@ class Tls:
         b += b'\xff' * (0x1000 - len(b))
         return b
 
-    def handle_empty(self, body):
+    def handle_empty(self, body: bytes):
         if body != b'\0' * len(body):
             raise Exception('Expected empty block')
 
-    def handle_cert(self, body):
+    def handle_cert(self, body: bytes):
         # TODO validate cert, check if pub keys match
         self.tls_cert = body
         self.trace('TLS cert blob: %s' % hexlify(self.tls_cert))
 
-    def handle_ecdh(self, body):
+    def handle_ecdh(self, body: bytes):
         self.ecdh_blob = body
         key, signature = body[:0x90], body[0x90:]
         x = key[0x8:0x8 + 0x20]
@@ -499,7 +500,7 @@ class Tls:
         # throws InvalidSignature
         fwpub.verify(signature, key, ec.ECDSA(hashes.SHA256()))
 
-    def handle_priv(self, body):
+    def handle_priv(self, body: bytes):
         self.priv_blob = body
         prefix, body = body[0], body[1:]
         if prefix != 2:

validitysensor/upload_fwext.py

@@ -1,6 +1,7 @@
 import logging
 from os.path import basename
 from time import ctime
+import typing
 
 from .flash import write_flash_all, write_fw_signature, get_fw_info
 from .sensor import reboot, write_hw_reg32, read_hw_reg32, identify_sensor
@@ -21,7 +22,7 @@ def default_fwext_name():
     return '6_07f_lenovo_mis_qm.xpfwext'
 
 
-def upload_fwext(fw_path=None):
+def upload_fwext(fw_path: typing.Optional[str]=None):
     fwi = get_fw_info(2)
     if fwi is not None:
         logging.info('Detected firmware version %d.%d (%s))' %

validitysensor/usb.py

@@ -2,6 +2,7 @@ import errno
 import logging
 from binascii import hexlify, unhexlify
 from struct import unpack
+import typing
 
 import usb.core as ucore
 from usb.core import USBError
@@ -22,9 +23,8 @@ class CancelledException(Exception):
 
 class Usb:
     def __init__(self):
-        self.interrupt_cb = None
         self.trace_enabled = False
-        self.dev = None
+        self.dev: typing.Optional[ucore.Device] = None
         self.cancel = False
 
     def open(self, vendor=None, product=None):
@@ -39,7 +39,7 @@ class Usb:
 
         self.open_dev(dev)
 
-    def open_devpath(self, busnum, address):
+    def open_devpath(self, busnum: int, address: int):
         def match(d):
             return d.bus == busnum and d.address == address
 
@@ -47,7 +47,7 @@ class Usb:
 
         self.open_dev(dev)
 
-    def open_dev(self, dev):
+    def open_dev(self, dev: ucore.Device):
         if dev is None:
             raise Exception('No matching devices found')
 
@@ -84,7 +84,7 @@ class Usb:
             logging.info('Clean slate')
             self.cmd(init_hardcoded_clean_slate)
 
-    def cmd(self, out):
+    def cmd(self, out: typing.Union[bytes, typing.Callable[[], bytes]]):
         if callable(out):
             out = out()
             if not out:
@@ -126,7 +126,7 @@ class Usb:
                 else:
                     raise e
 
-    def trace(self, s):
+    def trace(self, s: str):
         if self.trace_enabled:
             logging.debug(s)
 

validitysensor/util.py

@@ -3,7 +3,7 @@ from binascii import unhexlify
 from struct import unpack
 
 
-def assert_status(b):
+def assert_status(b: bytes):
     s, = unpack('<H', b[:2])
     if s != 0:
         if s == 0x44f:
@@ -12,5 +12,5 @@ def assert_status(b):
         raise Exception('Failed: %04x' % s)
 
 
-def unhex(x):
+def unhex(x: str):
     return unhexlify(re.sub(r'\W', '', x))