From f6154b311e82df114f86cf4833fd3c735bea1490 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Wed, 20 Jan 2021 16:22:07 +0100
Subject: [PATCH] chore(trezorlib): Support payment requests and GetNonce
 command.

---
 python/.changelog.d/1430.added |  1 +
 python/src/trezorlib/btc.py    | 57 +++++++++++++++++++---------------
 python/src/trezorlib/misc.py   |  5 +++
 3 files changed, 38 insertions(+), 25 deletions(-)
 create mode 100644 python/.changelog.d/1430.added

diff --git a/python/.changelog.d/1430.added b/python/.changelog.d/1430.added
new file mode 100644
index 000000000..3d2d76614
--- /dev/null
+++ b/python/.changelog.d/1430.added
@@ -0,0 +1 @@
+Support payment requests and GetNonce command.
diff --git a/python/src/trezorlib/btc.py b/python/src/trezorlib/btc.py
index 4fe7b9e77..29647bc4c 100644
--- a/python/src/trezorlib/btc.py
+++ b/python/src/trezorlib/btc.py
@@ -255,6 +255,7 @@ def sign_tx(
     outputs: Sequence[messages.TxOutputType],
     details: Optional[messages.SignTx] = None,
     prev_txes: Optional["TxCacheType"] = None,
+    payment_reqs: Sequence[messages.TxAckPaymentRequest] = (),
     preauthorized: bool = False,
     **kwargs: Any,
 ) -> Tuple[Sequence[Optional[bytes]], bytes]:
@@ -354,34 +355,40 @@ def sign_tx(
         else:
             current_tx = this_tx
 
-        msg = messages.TransactionType()
-
-        if res.request_type == R.TXMETA:
-            msg = copy_tx_meta(current_tx)
-        elif res.request_type in (R.TXINPUT, R.TXORIGINPUT):
-            assert res.details.request_index is not None
-            msg.inputs = [current_tx.inputs[res.details.request_index]]
-        elif res.request_type == R.TXOUTPUT:
-            assert res.details.request_index is not None
-            if res.details.tx_hash:
-                msg.bin_outputs = [current_tx.bin_outputs[res.details.request_index]]
-            else:
-                msg.outputs = [current_tx.outputs[res.details.request_index]]
-        elif res.request_type == R.TXORIGOUTPUT:
+        if res.request_type == R.TXPAYMENTREQ:
             assert res.details.request_index is not None
-            msg.outputs = [current_tx.outputs[res.details.request_index]]
-        elif res.request_type == R.TXEXTRADATA:
-            assert res.details.extra_data_offset is not None
-            assert res.details.extra_data_len is not None
-            assert current_tx.extra_data is not None
-            o, l = res.details.extra_data_offset, res.details.extra_data_len
-            msg.extra_data = current_tx.extra_data[o : o + l]
+            msg = payment_reqs[res.details.request_index]
+            res = client.call(msg)
         else:
-            raise exceptions.TrezorException(
-                f"Unknown request type - {res.request_type}."
-            )
+            msg = messages.TransactionType()
+            if res.request_type == R.TXMETA:
+                msg = copy_tx_meta(current_tx)
+            elif res.request_type in (R.TXINPUT, R.TXORIGINPUT):
+                assert res.details.request_index is not None
+                msg.inputs = [current_tx.inputs[res.details.request_index]]
+            elif res.request_type == R.TXOUTPUT:
+                assert res.details.request_index is not None
+                if res.details.tx_hash:
+                    msg.bin_outputs = [
+                        current_tx.bin_outputs[res.details.request_index]
+                    ]
+                else:
+                    msg.outputs = [current_tx.outputs[res.details.request_index]]
+            elif res.request_type == R.TXORIGOUTPUT:
+                assert res.details.request_index is not None
+                msg.outputs = [current_tx.outputs[res.details.request_index]]
+            elif res.request_type == R.TXEXTRADATA:
+                assert res.details.extra_data_offset is not None
+                assert res.details.extra_data_len is not None
+                assert current_tx.extra_data is not None
+                o, l = res.details.extra_data_offset, res.details.extra_data_len
+                msg.extra_data = current_tx.extra_data[o : o + l]
+            else:
+                raise exceptions.TrezorException(
+                    f"Unknown request type - {res.request_type}."
+                )
 
-        res = client.call(messages.TxAck(tx=msg))
+            res = client.call(messages.TxAck(tx=msg))
 
     if not isinstance(res, messages.TxRequest):
         raise exceptions.TrezorException("Unexpected message")
diff --git a/python/src/trezorlib/misc.py b/python/src/trezorlib/misc.py
index 6399675a8..cd4d2b4c9 100644
--- a/python/src/trezorlib/misc.py
+++ b/python/src/trezorlib/misc.py
@@ -108,3 +108,8 @@ def decrypt_keyvalue(
             iv=iv,
         )
     )
+
+
+@expect(messages.Nonce, field="nonce", ret_type=bytes)
+def get_nonce(client: "TrezorClient"):
+    return client.call(messages.GetNonce())