python-validity/bin/validity-sensors-initializer

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# 2020 - Marco Trevisan
#
# Initializer for ThinkPad's validity sensors 138a:0090 and 138a:0097 and 06cb:009a
#
# 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 <http://www.gnu.org/licenses/>.

import argparse
import os
import subprocess
import sys
import tempfile
import urllib.request
import shutil


from enum import Enum, auto
from time import sleep
from usb import core as usb_core

from validitysensor.flash import read_tls_flash
from validitysensor.init_db import init_db
from validitysensor.init_flash import init_flash
from validitysensor.sensor import sensor as vfs_sensor
from validitysensor.sensor import factory_reset
from validitysensor.tls import tls as vfs_tls
from validitysensor.upload_fwext import upload_fwext
from validitysensor.usb import usb as vfs_usb

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'
    }
}

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'
}


class VFSInitializer():
    def __init__(self, args, usb_dev, dev_type):
        self.args = args
        self.usb_dev = usb_dev
        self.dev_type = dev_type
        self.dev_str = repr(usb_dev)

        print('Found device {}'.format(self.dev_str))

        if self.args.host_product or self.args.host_serial:
            vfs_tls.set_hwkey(product_name=self.args.host_product,
                serial_number=self.args.host_serial)

    def open_device(self, init=False):
        print('Opening device',hex(self.dev_type.value[1]))
        vfs_usb.open(vendor=self.dev_type.value[0], product=self.dev_type.value[1])

        if init:
            vfs_usb.send_init()

            # try to init TLS session from the flash
            vfs_tls.parseTlsFlash(read_tls_flash())
            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]['driver']
        fwarchive = os.path.join(fwdir, 'fwinstaller.exe')
        fwname = DEFAULT_FW_NAMES[self.dev_type]

        print('Downloading {} to extract {}'.format(fwuri, fwname))

        req = urllib.request.Request(fwuri)
        req.add_header('Referer', DEFAULT_URIS[self.dev_type].get('referral', ''))
        req.add_header('User-Agent', 'Mozilla/5.0 (X11; U; Linux)')

        with urllib.request.urlopen(req) 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):
        print('Pairing the sensor with device')

        max_retries = 5
        err = None
        for i in range(0, max_retries):
            try:
                self.open_device()

                print('Initializing flash...')
                init_flash()
                break
            except Exception as e:
                if err is None:
                    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()

        self.sleep()
        self.restart()

        print('Calibrating...')
        vfs_sensor.open()
        vfs_sensor.calibrate()

        print('Init database...')
        init_db()

        vfs_tls.reset()

        print('That\'s it, pairing with {} finished'.format(self.dev_str))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--driver-uri')
    parser.add_argument('--firmware-path', type=argparse.FileType('r'))
    parser.add_argument('--calibration-data', type=argparse.FileType('r'))
    parser.add_argument('--host-product')
    parser.add_argument('--host-serial')
    parser.add_argument('--simulate-virtualbox', action='store_true')
    parser.add_argument('--download-fw-only', action='store_true')

    args = parser.parse_args()

    if args.simulate_virtualbox and (args.host_product or args.host_serial):
        parser.error("--simulate-virtualbox is incompatible with host params.")

    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=d.value[0], idProduct=d.value[1])
        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(args, usb_dev, dev_type)

    with tempfile.TemporaryDirectory() as fwdir:
        if args.firmware_path:
            fwpath = args.firmware_path.name
        else:
            fwpath = vfs_initializer.download_and_extract_fw(fwdir,
                fwuri=args.driver_uri)

        shutil.copy(fwpath, '/usr/share/python-validity/')

        if not args.download_fw_only:
            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()