From ee694f36cfda44bd3ad438691be0eb3c26cbfdc7 Mon Sep 17 00:00:00 2001 From: matejcik Date: Tue, 3 May 2022 13:32:39 +0200 Subject: [PATCH] feat(python): implement firmware dumping --- python/src/trezorlib/cli/firmware.py | 16 ++++++++++++++++ python/src/trezorlib/firmware.py | 18 +++++++++++++++++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/src/trezorlib/cli/firmware.py b/python/src/trezorlib/cli/firmware.py index fffbc8dbf..7c8c0e85b 100644 --- a/python/src/trezorlib/cli/firmware.py +++ b/python/src/trezorlib/cli/firmware.py @@ -599,3 +599,19 @@ def get_hash(client: "TrezorClient", hex_challenge: Optional[str]) -> str: """Get a hash of the installed firmware combined with the optional challenge.""" challenge = bytes.fromhex(hex_challenge) if hex_challenge else None return firmware.get_hash(client, challenge).hex() + + +@cli.command() +@click.argument("file", type=click.File("wb")) +@with_client +def extract(client: "TrezorClient", file: BinaryIO) -> None: + """Extract the firmware from the device.""" + if client.features.model == "T": + firmware_size = 13 * 128 * 1024 + elif client.features.model == "1": + firmware_size = 7 * 128 * 1024 + 64 * 1024 + else: + firmware_size = None + + with click.progressbar(length=firmware_size) as bar: + file.write(firmware.get_firmware(client, bar.update)) diff --git a/python/src/trezorlib/firmware.py b/python/src/trezorlib/firmware.py index 54408082c..80a9e7c65 100644 --- a/python/src/trezorlib/firmware.py +++ b/python/src/trezorlib/firmware.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple import construct as c import ecdsa -from . import cosi, messages +from . import cosi, exceptions, messages from .tools import expect, session if TYPE_CHECKING: @@ -523,3 +523,19 @@ def update( @expect(messages.FirmwareHash, field="hash", ret_type=bytes) def get_hash(client: "TrezorClient", challenge: Optional[bytes]): return client.call(messages.GetFirmwareHash(challenge=challenge)) + + +@session +def get_firmware( + client: "TrezorClient", progress_update: Callable[[int], Any] = lambda _: None +) -> bytes: + resp = client.call(messages.GetFirmware()) + result = bytearray() + while isinstance(resp, messages.FirmwareChunk): + result.extend(resp.chunk) + progress_update(len(resp.chunk)) + resp = client.call(messages.FirmwareChunkAck()) + + if not isinstance(resp, messages.Success): + raise exceptions.TrezorException(f"Unexpected message {resp}") + return bytes(result)