Consolidate definitions related to firmware files.

This fixes the TODO comments in validity-sensors-firmware

bin/validity-sensors-firmware

@@ -25,52 +25,25 @@ import subprocess
 import sys
 import tempfile
 import urllib.request
-from enum import Enum
 
 from usb import core as usb_core
 
+from validitysensor.usb import SupportedDevices
+from validitysensor.firmware_tables import FIRMWARE_NAMES, FIRMWARE_URIS
+
+
 python_validity_data = '/usr/share/python-validity/'
 
 
-# FIXME: supported usb ids are duplicated in usb.py
-class VFS(Enum):
-    DEV_90 = (0x138a, 0x0090)
-    DEV_97 = (0x138a, 0x0097)
-    DEV_9a = (0x06cb, 0x009a)
-
-
-DEFAULT_URIS = {
-    VFS.DEV_90: {
-        'driver': 'https://download.lenovo.com/pccbbs/mobiles/n1cgn08w.exe',
-        'referral': 'https://support.lenovo.com/us/en/downloads/DS120491',
-    },
-    VFS.DEV_97: {
-        'driver': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe',
-        'referral': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe'
-    },
-    VFS.DEV_9a: {
-        'driver': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe',
-        'referral': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe'
-    }
-}
-
-# FIXME: filenames are duplicated in upload_fwext.py
-DEFAULT_FW_NAMES = {
-    VFS.DEV_90: '6_07f_Lenovo.xpfwext',
-    VFS.DEV_97: '6_07f_lenovo_mis_qm.xpfwext',
-    VFS.DEV_9a: '6_07f_lenovo_mis_qm.xpfwext'
-}
-
-
 def download_and_extract_fw(dev_type, fwdir, fwuri=None):
-    fwuri = fwuri if fwuri else DEFAULT_URIS[dev_type]['driver']
+    fwuri = fwuri if fwuri else FIRMWARE_URIS[dev_type]['driver']
     fwarchive = os.path.join(fwdir, 'fwinstaller.exe')
-    fwname = DEFAULT_FW_NAMES[dev_type]
+    fwname = FIRMWARE_NAMES[dev_type]
 
     print('Downloading {} to extract {}'.format(fwuri, fwname))
 
     req = urllib.request.Request(fwuri)
-    req.add_header('Referer', DEFAULT_URIS[dev_type].get('referral', ''))
+    req.add_header('Referer', FIRMWARE_URIS[dev_type].get('referral', ''))
     req.add_header('User-Agent', 'Mozilla/5.0 (X11; U; Linux)')
 
     with urllib.request.urlopen(req) as response:
@@ -101,7 +74,7 @@ if __name__ == "__main__":
         raise Exception('This script needs to be executed as root')
 
     dev_type = None
-    for d in VFS:
+    for d in SupportedDevices:
         dev = usb_core.find(idVendor=d.value[0], idProduct=d.value[1])
         if dev:
             dev_type = d

validitysensor/firmware_tables.py

@@ -0,0 +1,24 @@
+"""Defines various constants for firmware files"""
+
+from .usb import SupportedDevices
+
+FIRMWARE_URIS = {
+    SupportedDevices.DEV_90: {
+        'driver': 'https://download.lenovo.com/pccbbs/mobiles/n1cgn08w.exe',
+        'referral': 'https://support.lenovo.com/us/en/downloads/DS120491',
+    },
+    SupportedDevices.DEV_97: {
+        'driver': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe',
+        'referral': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe'
+    },
+    SupportedDevices.DEV_9a: {
+        'driver': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe',
+        'referral': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe'
+    }
+}
+
+FIRMWARE_NAMES = {
+    SupportedDevices.DEV_90: '6_07f_Lenovo.xpfwext',
+    SupportedDevices.DEV_97: '6_07f_lenovo_mis_qm.xpfwext',
+    SupportedDevices.DEV_9a: '6_07f_lenovo_mis_qm.xpfwext'
+}

validitysensor/upload_fwext.py

@@ -3,23 +3,21 @@ import typing
 from os.path import basename
 from time import ctime
 
+from .firmware_tables import FIRMWARE_NAMES
 from .flash import write_flash_all, write_fw_signature, get_fw_info
 from .sensor import reboot, write_hw_reg32, read_hw_reg32, identify_sensor
-from .usb import usb
+from .usb import usb, SupportedDevices
 
 firmware_home = '/usr/share/python-validity'
 
 
 def default_fwext_name():
-    if usb.usb_dev().idVendor == 0x138a:
-        if usb.usb_dev().idProduct == 0x0090:
-            return '6_07f_Lenovo.xpfwext'
-
+    dev = SupportedDevices.from_usbid(usb.usb_dev().idVendor, usb.usb_dev().idProduct)
     # It looks like the firmware file must match DLL, not the hardware.
     # Both DLL and xpfwext are universal.
     # The device dependant code seems to be loaded dynamically (via encrypted blobs).
     # So, it is important that xpfwext file is matching the blobs contents.
-    return '6_07f_lenovo_mis_qm.xpfwext'
+    return FIRMWARE_NAMES[dev]
 
 
 def upload_fwext(fw_path: typing.Optional[str] = None):

validitysensor/usb.py

@@ -2,6 +2,7 @@ import errno
 import logging
 import typing
 from binascii import hexlify, unhexlify
+from enum import Enum
 from struct import unpack
 
 import usb.core as ucore
@@ -10,11 +11,19 @@ from usb.core import USBError
 from .blobs import init_hardcoded, init_hardcoded_clean_slate
 from .util import assert_status
 
-supported_devices = [
-    (0x138a, 0x0090),
-    (0x138a, 0x0097),
-    (0x06cb, 0x009a),
-]
+
+class SupportedDevices(Enum):
+    """USB IDs for supported devices"""
+    DEV_90 = (0x138a, 0x0090)
+    DEV_97 = (0x138a, 0x0097)
+    DEV_9a = (0x06cb, 0x009a)
+
+    @classmethod
+    def from_usbid(cls, vendorid, productid):
+        return supported_devices[(vendorid, productid)]
+
+
+supported_devices = dict((dev.value, dev) for dev in SupportedDevices)
 
 
 class CancelledException(Exception):