From 9a7990d087e2be50d91f88a507bc8252376452de Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Fri, 5 Nov 2021 21:24:36 +0100
Subject: [PATCH] feat(tests): Add no_script_type device tests for SignMessage.

---
 python/src/trezorlib/btc.py                | 13 +++++-
 tests/device_tests/test_msg_signmessage.py | 54 +++++++++++++++++++++-
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/python/src/trezorlib/btc.py b/python/src/trezorlib/btc.py
index 8b65d3fc0..de73a3c63 100644
--- a/python/src/trezorlib/btc.py
+++ b/python/src/trezorlib/btc.py
@@ -156,12 +156,21 @@ def get_ownership_proof(
 
 @expect(messages.MessageSignature)
 def sign_message(
-    client, coin_name, n, message, script_type=messages.InputScriptType.SPENDADDRESS
+    client,
+    coin_name,
+    n,
+    message,
+    script_type=messages.InputScriptType.SPENDADDRESS,
+    no_script_type=False,
 ):
     message = normalize_nfc(message)
     return client.call(
         messages.SignMessage(
-            coin_name=coin_name, address_n=n, message=message, script_type=script_type
+            coin_name=coin_name,
+            address_n=n,
+            message=message,
+            script_type=script_type,
+            no_script_type=no_script_type,
         )
     )
 
diff --git a/tests/device_tests/test_msg_signmessage.py b/tests/device_tests/test_msg_signmessage.py
index 87d818950..4ef9d0488 100644
--- a/tests/device_tests/test_msg_signmessage.py
+++ b/tests/device_tests/test_msg_signmessage.py
@@ -43,6 +43,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "44h/0h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "1JAd7XCBzGudGpJQSDSfpmJhiygtLQWaGL",
         "This is an example of a signed message.",
         "20fd8f2f7db5238fcdd077d5204c3e6949c261d700269cefc1d9d2dcef6b95023630ee617f6c8acf9eb40c8edd704c9ca74ea4afc393f43f35b4e8958324cbdd1c",
@@ -52,6 +53,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "49h/0h/0h/0/0",
         S.SPENDP2SHWITNESS,
+        False,
         "3L6TyTisPBmrDAj6RoKmDzNnj4eQi54gD2",
         "This is an example of a signed message.",
         "23744de4516fac5c140808015664516a32fead94de89775cec7e24dbc24fe133075ac09301c4cc8e197bea4b6481661d5b8e9bf19d8b7b8a382ecdb53c2ee0750d",
@@ -61,16 +63,48 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "84h/0h/0h/0/0",
         S.SPENDWITNESS,
+        False,
         "bc1qannfxke2tfd4l7vhepehpvt05y83v3qsf6nfkk",
         "This is an example of a signed message.",
         "28b55d7600d9e9a7e2a49155ddf3cfdb8e796c207faab833010fa41fb7828889bc47cf62348a7aaa0923c0832a589fab541e8f12eb54fb711c90e2307f0f66b194",
     ),
+    case(
+        "p2pkh",
+        "Bitcoin",
+        "44h/0h/0h/0/0",
+        S.SPENDADDRESS,
+        True,
+        "1JAd7XCBzGudGpJQSDSfpmJhiygtLQWaGL",
+        "This is an example of a signed message.",
+        "20fd8f2f7db5238fcdd077d5204c3e6949c261d700269cefc1d9d2dcef6b95023630ee617f6c8acf9eb40c8edd704c9ca74ea4afc393f43f35b4e8958324cbdd1c",
+    ),
+    case(
+        "segwit-p2sh",
+        "Bitcoin",
+        "49h/0h/0h/0/0",
+        S.SPENDP2SHWITNESS,
+        True,
+        "3L6TyTisPBmrDAj6RoKmDzNnj4eQi54gD2",
+        "This is an example of a signed message.",
+        "1f744de4516fac5c140808015664516a32fead94de89775cec7e24dbc24fe133075ac09301c4cc8e197bea4b6481661d5b8e9bf19d8b7b8a382ecdb53c2ee0750d",
+    ),
+    case(
+        "segwit-native",
+        "Bitcoin",
+        "84h/0h/0h/0/0",
+        S.SPENDWITNESS,
+        True,
+        "bc1qannfxke2tfd4l7vhepehpvt05y83v3qsf6nfkk",
+        "This is an example of a signed message.",
+        "20b55d7600d9e9a7e2a49155ddf3cfdb8e796c207faab833010fa41fb7828889bc47cf62348a7aaa0923c0832a589fab541e8f12eb54fb711c90e2307f0f66b194",
+    ),
     # ==== Bitcoin with long message ====
     case(
         "p2pkh long message",
         "Bitcoin",
         "44h/0h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "1JAd7XCBzGudGpJQSDSfpmJhiygtLQWaGL",
         "VeryLongMessage!" * 64,
         "200a46476ceb84d06ef5784828026f922c8815f57aac837b8c013007ca8a8460db63ef917dbebaebd108b1c814bbeea6db1f2b2241a958e53fe715cc86b199d9c3",
@@ -80,6 +114,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "49h/0h/0h/0/0",
         S.SPENDP2SHWITNESS,
+        False,
         "3L6TyTisPBmrDAj6RoKmDzNnj4eQi54gD2",
         "VeryLongMessage!" * 64,
         "236eadee380684f70749c52141c8aa7c3b6afd84d0e5f38cfa71823f3b1105a5f34e23834a5bb6f239ff28ad87f409f44e4ce6269754adc00388b19507a5d9386f",
@@ -89,6 +124,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "84h/0h/0h/0/0",
         S.SPENDWITNESS,
+        False,
         "bc1qannfxke2tfd4l7vhepehpvt05y83v3qsf6nfkk",
         "VeryLongMessage!" * 64,
         "28c6f86e255eaa768c447d635d91da01631ac54af223c2c182d4fa3676cfecae4a199ad33a74fe04fb46c39432acb8d83de74da90f5f01123b3b7d8bc252bc7f71",
@@ -99,6 +135,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "44h/0h/0h/0/1",
         S.SPENDADDRESS,
+        False,
         "1GWFxtwWmNVqotUPXLcKVL2mUKpshuJYo",
         MESSAGE_NFKD,
         NFKD_NFC_SIGNATURE,
@@ -108,6 +145,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bitcoin",
         "44h/0h/0h/0/1",
         S.SPENDADDRESS,
+        False,
         "1GWFxtwWmNVqotUPXLcKVL2mUKpshuJYo",
         MESSAGE_NFC,
         NFKD_NFC_SIGNATURE,
@@ -118,6 +156,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Testnet",
         "44h/1h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "mvbu1Gdy8SUjTenqerxUaZyYjmveZvt33q",
         "This is an example of a signed message.",
         "2030cd7f116c0481d1936cfef48137fd23ee56aaf00787bfa08a94837466ec9909390c3efacfc56bae5782f1db4cf49ae05f242b5f62a47f871ec46bf1a3253e7f",
@@ -127,6 +166,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Testnet",
         "49h/1h/0h/0/0",
         S.SPENDP2SHWITNESS,
+        False,
         "2N4Q5FhU2497BryFfUgbqkAJE87aKHUhXMp",
         "This is an example of a signed message.",
         "23ef39fd388c3425d6aaa04274dcd5c7dd4c283a411b616443474fbcde5dd966050d91bc7c57e9578f28efdd84c9a9bcba415f93c5727b5d3f2bf3de46d7084896",
@@ -136,6 +176,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Testnet",
         "84h/1h/0h/0/0",
         S.SPENDWITNESS,
+        False,
         "tb1qkvwu9g3k2pdxewfqr7syz89r3gj557l3uuf9r9",
         "This is an example of a signed message.",
         "27758b3393396ad9fe48f6ce81f63410145e7b2b69a5dfc1d48b5e6e623e91e08e3afb60bda1546f9c6f9fb5bd0a41887b784c266036dd4b4015a0abc1137daa1d",
@@ -146,6 +187,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Bcash",
         "44h/145h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "bitcoincash:qr08q88p9etk89wgv05nwlrkm4l0urz4cyl36hh9sv",
         "This is an example of a signed message.",
         "1fda7733e666a4ab8ba86f3cfc3728d318ecf824a3bf99597570297aa131607c10316959136b2c500b2b478a73c563ba314c0b7b2a22065b6d9596118f246d360e",
@@ -156,6 +198,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Groestlcoin",
         "44h/17h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "Fj62rBJi8LvbmWu2jzkaUX1NFXLEqDLoZM",
         "test",
         "20d39869afe38fc631cf7983e64f9b65f5268e48c1f55ce857874d1bbf91b015322b7d312fb23dc8c816595bec2f8e82e7242dc6d658d1c45193babd37a6fe6133",
@@ -166,6 +209,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Groestlcoin",
         "49h/17h/0h/0/0",
         S.SPENDP2SHWITNESS,
+        False,
         "31inaRqambLsd9D7Ke4USZmGEVd3PHkh7P",
         "test",
         "23f340fc9f9ea6469e13dbc743b70313e4d076bcd8ce867eddd71ec41160d02a4a462205d21ec6e49502bf3e2a8463d48e895ca56f6b385b15ec2cc7556292ecae",
@@ -176,6 +220,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Groestlcoin",
         "84h/17h/0h/0/0",
         S.SPENDWITNESS,
+        False,
         "grs1qw4teyraux2s77nhjdwh9ar8rl9dt7zww8r6lne",
         "test",
         "288253db4b4a1d5dac059296385310a353ef80992c4777a44133a335d12d3444da6c287d32aec4071ec49ae327e208f89ba0a115a129f106221c8dd5590fd3df13",
@@ -186,6 +231,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Decred",
         "44h/42h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "DsZtHtXHwvNR3nWf1PqfxrEdnRJisKEyzp1",
         "This is an example of a signed message.",
         "206b1f8ba47ef9eaf87aa900e41ab1e97f67e8c09292faa4acf825228d074c4b774484046dcb1d9bbf0603045dbfb328c3e1b0c09c5ae133e89e604a67a1fc6cca",
@@ -196,6 +242,7 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
         "Decred",
         "44h/42h/0h/0/0",
         S.SPENDADDRESS,
+        False,
         "DsZtHtXHwvNR3nWf1PqfxrEdnRJisKEyzp1",
         "",
         "1fd2d57490b44a0361c7809768cad032d41ba1d4b7a297f935fc65ae05f71de7ea0c6c6fd265cc5154f1fa4acd7006b6a00ddd67fb7333c1594aff9120b3ba8024",
@@ -205,14 +252,17 @@ VECTORS = (  # case name, coin_name, path, script_type, address, message, signat
 
 
 @pytest.mark.parametrize(
-    "coin_name, path, script_type, address, message, signature", VECTORS
+    "coin_name, path, script_type, no_script_type, address, message, signature", VECTORS
 )
-def test_signmessage(client, coin_name, path, script_type, address, message, signature):
+def test_signmessage(
+    client, coin_name, path, script_type, no_script_type, address, message, signature
+):
     sig = btc.sign_message(
         client,
         coin_name=coin_name,
         n=parse_path(path),
         script_type=script_type,
+        no_script_type=no_script_type,
         message=message,
     )
     assert sig.address == address