diff --git a/validity-sensors-initializer.py b/validity-sensors-initializer.py new file mode 100644 index 0000000..16ae261 --- /dev/null +++ b/validity-sensors-initializer.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# +# 2020 - Marco Trevisan +# +# Initializer for ThinkPad's validity sensors 0090 and 0097 +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import os +import subprocess +import sys +import tempfile +import urllib.request + +from enum import Enum, auto +from time import sleep +from usb import core as usb_core + +from proto9x.calibrate import calibrate +from proto9x.flash import read_flash +from proto9x.init_db import init_db +from proto9x.init_flash import init_flash +from proto9x.sensor import factory_reset +from proto9x.tls import tls as vfs_tls +from proto9x.upload_fwext import upload_fwext +from proto9x.usb import usb as vfs_usb + +class VFS(Enum): + DEV_90 = 0x0090 + DEV_97 = 0x0097 + +DEFAULT_URIS = { + VFS.DEV_90: 'https://download.lenovo.com/pccbbs/mobiles/n1cgn08w.exe', + VFS.DEV_97: 'https://download.lenovo.com/pccbbs/mobiles/n1mgf03w.exe', +} + +DEFAULT_FW_NAMES = { + VFS.DEV_90: '6_07f_Lenovo.xpfwext', + VFS.DEV_97: '6_07f_lenovo_mis.xpfwext', +} + + +class VFSInitializer(): + def __init__(self, usb_dev, dev_type): + self.usb_dev = usb_dev + self.dev_type = dev_type + self.dev_str = repr(usb_dev) + + print('Found device {}'.format(self.dev_str)) + + try: + with open('/sys/class/dmi/id/product_name', 'r') as node: + self.product_name = node.read().strip() + with open('/sys/class/dmi/id/product_serial', 'r') as node: + self.product_serial = node.read().strip() + except: + self.product_name = 'VirtualBox' + self.product_serial = '0' + + vfs_tls.set_hwkey(product_name=self.product_name, + serial_number=self.product_serial) + + def open_device(self, init=False): + print('Opening device',hex(self.dev_type.value)) + vfs_usb.open(product=self.dev_type.value) + + if init: + vfs_usb.send_init() + + # try to init TLS session from the flash + vfs_tls.parseTlsFlash(read_flash(1, 0, 0x1000)) + vfs_tls.open() + + def restart(self): + vfs_tls.reset() + self.open_device(init=True) + + def download_and_extract_fw(self, fwdir, fwuri=None): + fwuri = fwuri if fwuri else DEFAULT_URIS[self.dev_type] + fwarchive = os.path.join(fwdir, 'fwinstaller.exe') + fwname = DEFAULT_FW_NAMES[self.dev_type] + + print('Downloading {} to extract {}'.format(fwuri, fwname)) + + with urllib.request.urlopen(fwuri) as response: + with open(fwarchive, 'wb') as out_file: + out_file.write(response.read()) + + subprocess.check_call(['innoextract', + '--output-dir', fwdir, + '--include', fwname, + '--collisions', 'overwrite', + fwarchive + ]) + + fwpath = subprocess.check_output([ + 'find', fwdir, '-name', fwname]).decode('utf-8').strip() + print('Found firmware at {}'.format(fwpath)) + + if not fwpath: + raise Exception('No {} found in the archive'.format(fwname)) + + return fwpath + + def sleep(self, sec=3): + print('Sleeping...') + sleep(sec) + + def try_factory_reset(self): + self.open_device() + try: + print('Factory reset...') + factory_reset() + except Exception as e: + print('Factory reset failed with {}, this should not happen, but ' \ + 'we can ignore it, if pairing works...'.format(e)) + + def pair(self, fwpath): + print('Pairing the sensor with device {}'.format(self.product_name)) + + max_retries = 5 + for i in range(0, max_retries): + try: + self.open_device() + + print('Initializing flash...') + init_flash() + break + except Exception as e: + err = e + self.sleep() + print('Try {} failed with error: {}'.format(i+1, e)) + finally: + max_retries -= 1 + + if max_retries == 0: + print('Device didn\'t show up after reset, retry...') + raise(err) + + self.sleep() + self.restart() + + print('Uploading firmware...') + upload_fwext(fw_path=fwpath) + + self.sleep() + self.restart() + + print('Calibrating...') + calibrate() + + print('Init database...') + init_db() + + vfs_tls.reset() + + print('That\'s it, pairing with {} finished'.format(self.dev_str)) + + +if __name__ == "__main__": + if os.geteuid() != 0: + raise Exception('This script needs to be executed as root') + + usb_dev = None + for d in VFS: + dev = usb_core.find(idVendor=0x138a, idProduct=d.value) + if dev: + dev_type = d + usb_dev = dev + + if not usb_dev: + raise Exception('No supported validity device found') + + try: + subprocess.check_call(['innoextract', '--version'], + stdout=subprocess.DEVNULL) + except Exception as e: + print('Impossible to run innoextract: {}'.format(e)) + sys.exit(1) + + vfs_initializer = VFSInitializer(usb_dev, dev_type) + + with tempfile.TemporaryDirectory() as fwdir: + fwpath = vfs_initializer.download_and_extract_fw(fwdir) + + input('The device will be now reset to factory and associated to the ' \ + 'current laptop.\nPress Enter to continue (or Ctrl+C to cancel)...') + + vfs_initializer.try_factory_reset() + vfs_initializer.sleep() + vfs_initializer.pair(fwpath)