From 83bb3a09320500d1e4aeeebef02224dd24b45f44 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Tue, 9 Feb 2021 13:38:42 +0100
Subject: [PATCH] feat(python): show progress bar for firmware upload

---
 python/.changelog.d/noissue.added    |  1 +
 python/src/trezorlib/cli/firmware.py | 12 +++++++-----
 python/src/trezorlib/firmware.py     | 11 +++++++++--
 3 files changed, 17 insertions(+), 7 deletions(-)
 create mode 100644 python/.changelog.d/noissue.added

diff --git a/python/.changelog.d/noissue.added b/python/.changelog.d/noissue.added
new file mode 100644
index 000000000..170c38bb0
--- /dev/null
+++ b/python/.changelog.d/noissue.added
@@ -0,0 +1 @@
+`trezorctl firmware update` shows progress bar (Model T only)
diff --git a/python/src/trezorlib/cli/firmware.py b/python/src/trezorlib/cli/firmware.py
index cf44ebd56..843e918a5 100644
--- a/python/src/trezorlib/cli/firmware.py
+++ b/python/src/trezorlib/cli/firmware.py
@@ -394,7 +394,12 @@ def upload_firmware_into_device(
         if f.major_version == 1 and f.firmware_present is not False:
             # Trezor One does not send ButtonRequest
             click.echo("Please confirm the action on your Trezor device")
-        firmware.update(client, firmware_data)
+
+        click.echo("Uploading...\r", nl=False)
+        with click.progressbar(
+            label="Uploading", length=len(firmware_data), show_eta=False
+        ) as bar:
+            firmware.update(client, firmware_data, bar.update)
     except exceptions.Cancelled:
         click.echo("Update aborted on device.")
     except exceptions.TrezorException as e:
@@ -582,7 +587,4 @@ def update(
     if dry_run:
         click.echo("Dry run. Not uploading firmware to device.")
     else:
-        upload_firmware_into_device(
-            client=client,
-            firmware_data=firmware_data,
-        )
+        upload_firmware_into_device(client=client, firmware_data=firmware_data)
diff --git a/python/src/trezorlib/firmware.py b/python/src/trezorlib/firmware.py
index b3c8fdec0..bda8730ae 100644
--- a/python/src/trezorlib/firmware.py
+++ b/python/src/trezorlib/firmware.py
@@ -485,7 +485,11 @@ def validate(
 
 
 @session
-def update(client: "TrezorClient", data: bytes) -> None:
+def update(
+    client: "TrezorClient",
+    data: bytes,
+    progress_update: Callable[[int], Any] = lambda _: None,
+):
     if client.features.bootloader_mode is False:
         raise RuntimeError("Device must be in bootloader mode")
 
@@ -494,6 +498,7 @@ def update(client: "TrezorClient", data: bytes) -> None:
     # TREZORv1 method
     if isinstance(resp, messages.Success):
         resp = client.call(messages.FirmwareUpload(payload=data))
+        progress_update(len(data))
         if isinstance(resp, messages.Success):
             return
         else:
@@ -503,9 +508,11 @@ def update(client: "TrezorClient", data: bytes) -> None:
     while isinstance(resp, messages.FirmwareRequest):
         assert resp.offset is not None
         assert resp.length is not None
-        payload = data[resp.offset : resp.offset + resp.length]
+        length = resp.length
+        payload = data[resp.offset : resp.offset + length]
         digest = blake2s(payload).digest()
         resp = client.call(messages.FirmwareUpload(payload=payload, hash=digest))
+        progress_update(length)
 
     if isinstance(resp, messages.Success):
         return