diff --git a/src/unisoc_unlock/cli.py b/src/unisoc_unlock/cli.py index da6265e..5ef071a 100755 --- a/src/unisoc_unlock/cli.py +++ b/src/unisoc_unlock/cli.py @@ -8,6 +8,8 @@ from Crypto.PublicKey import RSA import base64 import io +import argparse +import importlib.metadata class OemIdToken: @@ -25,48 +27,96 @@ def __call__(self, fb_msg): self.n += 1 -def sign_token(tok, key_file): - priv_key = RSA.importKey(open(key_file).read()) - h = SHA256.new(tok) - signature = PKCS1_v1_5.new(priv_key).sign(h) - return signature +class BootloaderCmd: + def sign_token(self, tok, key_file): + priv_key = RSA.importKey(open(key_file).read()) + h = SHA256.new(tok) + signature = PKCS1_v1_5.new(priv_key).sign(h) + return signature + def prepare(self): + try: + self.dev = fastboot.FastbootCommands() + self.dev.ConnectDevice() + except usb_exceptions.DeviceNotFoundError as e: + print('No device found: {}'.format(e), file=sys.stderr) + sys.exit(1) + except usb_exceptions.CommonUsbError as e: + print('Could not connect to device: {}'.format(e), file=sys.stderr) + sys.exit(1) -def main(): - try: - dev = fastboot.FastbootCommands() - dev.ConnectDevice() - except usb_exceptions.DeviceNotFoundError as e: - print('No device found: {}'.format(e), file=sys.stderr) - sys.exit(1) - except usb_exceptions.CommonUsbError as e: - print('Could not connect to device: {}'.format(e), file=sys.stderr) - sys.exit(1) + oem_id = OemIdToken() + try: + self.dev.Oem('get_identifier_token', info_cb=oem_id) + except Exception as e: + print(f'Fastboot error: {str(e)}') + sys.exit(1) - oem_id = OemIdToken() - try: - dev.Oem('get_identifier_token', info_cb=oem_id) - except Exception as e: - print(f'Fastboot error: {str(e)}') - sys.exit(1) + print(f'OEM ID: {oem_id.id}') + id = oem_id.id.ljust(2*64, '0') + id_raw = base64.b16decode(id, casefold=True) + pemfile = os.path.join( + os.path.dirname(__file__), + 'rsa4096_vbmeta.pem' + ) + sgn = self.sign_token(id_raw, pemfile) - print(f'OEM ID: {oem_id.id}') - id = oem_id.id.ljust(2*64, '0') - id_raw = base64.b16decode(id, casefold=True) - pemfile = os.path.join( - os.path.dirname(__file__), - 'rsa4096_vbmeta.pem' - ) - sgn = sign_token(id_raw, pemfile) + print('Download signature') + self.dev.Download(io.BytesIO(sgn), source_len=len(sgn)) + + +class BootloaderUnlock(BootloaderCmd): + def __call__(self): + print('Preparing to unlock the bootloader') + self.prepare() + + print('Unlock bootloader, pls follow instructions on device screen') + self.dev._SimpleCommand( + b'flashing unlock_bootloader', timeout_ms=60*1000) - print('Download signature') - dev.Download(io.BytesIO(sgn), source_len=len(sgn)) + print('Bootloader unlocked.') + self.dev.Close() - print('Unlock bootloader, pls follow instructions on device screen') - dev._SimpleCommand(b'flashing unlock_bootloader', timeout_ms=60*1000) - print('Bootloader unlocked.') - dev.Close() +class BootloaderLock(BootloaderCmd): + def __call__(self): + print('Preparing to lock the bootloader') + self.prepare() + + print('Lock bootloader, pls follow instructions on device screen') + self.dev._SimpleCommand( + b'flashing lock_bootloader', timeout_ms=60*1000) + + print('Bootloader locked.') + self.dev.Close() + + +def main(): + parser = argparse.ArgumentParser( + description='Lock/Unlock tool for Spreadtrum/Unisoc bootloader' + ) + parser.add_argument('command', + type=str, + nargs='?', + help='Command (lock|unlock), default=unlock' + ) + parser.add_argument('--version', + action='version', + version='%(prog)s ' + + importlib.metadata.version('unisoc-unlock') + ) + + args = parser.parse_args() + + if args.command == 'lock': + cmd = BootloaderLock() + elif args.command in ['unlock', None]: + cmd = BootloaderUnlock() + else: + print(f'Unknown command {args.command}', file=sys.stderr) + sys.exit(1) + + cmd() if __name__ == '__main__':