From a3ac5d752eba524dd6a314f67156d882f2bdd45f Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Tue, 9 Aug 2022 18:27:21 +0200
Subject: [PATCH] feat(trezorctl): Add device set-busy command.

---
 python/.changelog.d/2445.added     |  1 +
 python/src/trezorlib/cli/device.py | 27 +++++++++++++++++++++++++++
 python/src/trezorlib/device.py     | 13 +++++++++++++
 3 files changed, 41 insertions(+)
 create mode 100644 python/.changelog.d/2445.added

diff --git a/python/.changelog.d/2445.added b/python/.changelog.d/2445.added
new file mode 100644
index 0000000000..26b05067b5
--- /dev/null
+++ b/python/.changelog.d/2445.added
@@ -0,0 +1 @@
+Add device set-busy command to trezorctl.
diff --git a/python/src/trezorlib/cli/device.py b/python/src/trezorlib/cli/device.py
index 1e09975a06..57f769213e 100644
--- a/python/src/trezorlib/cli/device.py
+++ b/python/src/trezorlib/cli/device.py
@@ -290,3 +290,30 @@ def reboot_to_bootloader(obj: "TrezorConnection") -> str:
     # which triggers double prompt on device
     with obj.client_context() as client:
         return device.reboot_to_bootloader(client)
+
+
+@cli.command()
+@click.argument("enable", type=ChoiceType({"on": True, "off": False}), required=False)
+@click.option(
+    "-e",
+    "--expiry",
+    type=int,
+    help="Dialog expiry in seconds.",
+)
+@with_client
+def set_busy(
+    client: "TrezorClient", enable: Optional[bool], expiry: Optional[int]
+) -> str:
+    """Show a "Do not disconnect" dialog."""
+    if enable is False:
+        return device.set_busy(client, None)
+
+    if expiry is None:
+        raise click.ClickException("Missing option '-e' / '--expiry'.")
+
+    if expiry <= 0:
+        raise click.ClickException(
+            f"Invalid value for '-e' / '--expiry': '{expiry}' is not a positive integer."
+        )
+
+    return device.set_busy(client, expiry * 1000)
diff --git a/python/src/trezorlib/device.py b/python/src/trezorlib/device.py
index 6f632d1afd..92eb504f03 100644
--- a/python/src/trezorlib/device.py
+++ b/python/src/trezorlib/device.py
@@ -224,3 +224,16 @@ def cancel_authorization(client: "TrezorClient") -> "MessageType":
 @expect(messages.Success, field="message", ret_type=str)
 def reboot_to_bootloader(client: "TrezorClient") -> "MessageType":
     return client.call(messages.RebootToBootloader())
+
+
+@expect(messages.Success, field="message", ret_type=str)
+@session
+def set_busy(client: "TrezorClient", expiry_ms: Optional[int]) -> "MessageType":
+    """Sets or clears the busy state of the device.
+
+    In the busy state the device shows a "Do not disconnect" message instead of the homescreen.
+    Setting `expiry_ms=None` clears the busy state.
+    """
+    ret = client.call(messages.SetBusy(expiry_ms=expiry_ms))
+    client.refresh_features()
+    return ret