From 2417ebf47aa36f0242a8b3a2a8f193823be4e9e8 Mon Sep 17 00:00:00 2001
From: Salvatore Ingala <6681844+bigspider@users.noreply.github.com>
Date: Tue, 18 Apr 2023 11:49:39 +0200
Subject: [PATCH] Add set of tests on generated address for miniscripts
 covering all fragments

---
 tests/test_get_wallet_address.py | 89 ++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/tests/test_get_wallet_address.py b/tests/test_get_wallet_address.py
index 9206db585..f0c56a235 100644
--- a/tests/test_get_wallet_address.py
+++ b/tests/test_get_wallet_address.py
@@ -1,9 +1,16 @@
+from hashlib import sha256
+import hmac
+import re
+
 from bitcoin_client.ledger_bitcoin import Client, AddressType, MultisigWallet, WalletPolicy
 from bitcoin_client.ledger_bitcoin.exception.errors import IncorrectDataError
 
+from .conftest import testnet_to_regtest_addr as T
 
 import pytest
 
+from test_utils import SpeculosGlobals
+
 # TODO: add tests with UI
 
 
@@ -294,3 +301,85 @@ def test_get_wallet_address_large_addr_index(client: Client):
     # too large address_index, not allowed for an unhardened step
     with pytest.raises(IncorrectDataError):
         client.get_wallet_address(wallet, wallet_hmac, 0, 2**31, False)
+
+
+def test_get_wallet_address_miniscript_all_fragments(client: Client, speculos_globals: SpeculosGlobals, rpc):
+    # Create some miniscripts to exercise all possible fragments at least once,
+    # by comparing with the addresses generated by bitcoin-core.
+    # (currently only covering miniscript in `wsh()`, not on taproot leaves)
+
+    # arbitrary 20-bytes and 32-bytes hex strings
+    H20 = bytes(list(range(20))).hex()
+    H32 = bytes(list(range(32))).hex()
+    fragments = [
+        "or_d(pk(@1/**),0)",                     # 0 and or_d
+        "1",                                     # 1
+        "c:pk_k(@1/**)",                         # pk_k and c:
+        "c:pk_h(@1/**)",                         # pk_h
+        "older(42)",                             # older
+        "after(42)",                             # after
+        f"sha256({H32})",                        # sha256
+        f"ripemd160({H20})",                     # ripemd160
+        f"hash256({H32})",                       # hash256
+        f"hash160({H20})",                       # hash160
+        "andor(pk(@1/**),older(42),pk(@2/**))",  # andor
+        "and_v(v:pk(@1/**),pk(@2/**))",          # and_v and v:
+        "and_b(pk(@1/**),a:pk(@2/**))",          # and_b and a:
+        "or_b(pk(@1/**),a:pk(@2/**))",           # or_b
+        "t:or_c(pk(@1/**),v:pk(@2/**))",         # or_c and t:
+        # or_d is covered
+        "or_i(pk(@1/**),pk(@2/**))",             # or_i
+        "thresh(1,pk(@1/**),a:pk(@2/**))",       # thresh
+        "multi(2,@1/**,@2/**,@3/**)",            # multi
+
+        # WRAPPERS not covered above
+        # a: is covered
+        "and_b(1,s:pk(@1/**))",                  # s:
+        # c: is covered
+        "dv:older(42)",                          # d:
+        # t: is covered
+        # v: is covered
+        "j:pk(@1/**)",                           # j:
+        "n:pk(@1/**)",                           # n:
+        "l:pk(@1/**)",                           # l:
+        "u:pk(@1/**)",                           # u:
+    ]
+
+    def prepend_a(frag):
+        # prepends the a: wrapper (taking into account that `frag` could already start with wrappers)
+        if re.match("^[a-z]+:", frag):
+            return "a" + frag
+        else:
+            return "a:" + frag
+
+    test_keys = [
+        "[f5acc2fd/48'/1'/0'/2']tpubDFAqEGNyad35aBCKUAXbQGDjdVhNueno5ZZVEn3sQbW5ci457gLR7HyTmHBg93oourBssgUxuWz1jX5uhc1qaqFo9VsybY1J5FuedLfm4dK",
+        'tpubDDcmHJ6bsQqSRDzXrF1cgyPfXpFTHmqBUcq5cevfszh83XJtjqXZXDYwP3N82bA51dBVhbe3uaaWwAxW2tEsjgZPXmupQpNwdmULXq1WXDU',
+        'tpubDCXK744twow5CX8HdAvV4Vez413R4xrM3hgD85mA3EpbnwgvtBmhh18eLAGsL5R9E2mwThPTz9fs4x4ZYgCC6GuuKmzSitH9FgWyqaDEKta',
+        'tpubDCLxCbopTq5qisZzRcf5ZJ8dHR3PXEexc1vDUR61eGDnSVcXjvEwC9CFXqRPzCi9vmrMd6xfJtFrZY8yrPo5886K1AjJACAviLuEXMNfvbS'
+    ]
+
+    is_change = False
+    addr_index = 3
+
+    for fr in fragments:
+        # We use "and_b(pk(@0/**),a:<miniscript_to_be_tested>})" as a generic gadget to compute a valid descriptor
+        # that can be registered, as long as the <miniscript_to_be_tested> if valid and safe.
+        # The key placeholders in <miniscript_to_be_tested> must start from @1 (if present).
+
+        desc_tmpl = f"wsh(and_b(pk(@0/**),{prepend_a(fr)}))"
+        n_keys = desc_tmpl.count("@")
+
+        assert n_keys <= len(test_keys), "add more tpubs to the test_keys"
+
+        wallet_policy = WalletPolicy("A policy", desc_tmpl, test_keys[:n_keys])
+
+        wallet_hmac = hmac.new(speculos_globals.wallet_registration_key, wallet_policy.id, sha256).digest()
+        addr_hww = client.get_wallet_address(wallet_policy, wallet_hmac, is_change, addr_index, False)
+
+        desc = wallet_policy.get_descriptor(is_change)
+        # compute descriptor checksum and derive the address
+        desc_chk = rpc.getdescriptorinfo(desc)["descriptor"]
+        addr_core = rpc.deriveaddresses(desc_chk, [3, 3])[0]
+
+        assert T(addr_hww) == addr_core