test: add script_util helper for creating bare multisig scripts

5 files changed, 23 insertions, 18 deletions

diff --git a/test/functional/feature_segwit.py b/test/functional/feature_segwit.py
index 5abe989e55..acb7469c6a 100755
--- a/test/functional/feature_segwit.py
+++ b/test/functional/feature_segwit.py
@@ -30,8 +30,6 @@ from test_framework.script import (
     CScript,
     OP_0,
     OP_1,
-    OP_2,
-    OP_CHECKMULTISIG,
     OP_DROP,
     OP_TRUE,
 )
@@ -39,6 +37,7 @@ from test_framework.script_util import (
     key_to_p2pk_script,
     key_to_p2pkh_script,
     key_to_p2wpkh_script,
+    keys_to_multisig_script,
     script_to_p2sh_script,
     script_to_p2wsh_script,
 )
@@ -149,7 +148,7 @@ class SegWitTest(BitcoinTestFramework):
             key = get_generate_key()
             self.pubkey.append(key.pubkey)
 
-            multiscript = CScript([OP_1, bytes.fromhex(self.pubkey[-1]), OP_1, OP_CHECKMULTISIG])
+            multiscript = keys_to_multisig_script([self.pubkey[-1]])
             p2sh_ms_addr = self.nodes[i].createmultisig(1, [self.pubkey[-1]], 'p2sh-segwit')['address']
             bip173_ms_addr = self.nodes[i].createmultisig(1, [self.pubkey[-1]], 'bech32')['address']
             assert_equal(p2sh_ms_addr, script_to_p2sh_p2wsh(multiscript))
@@ -389,7 +388,7 @@ class SegWitTest(BitcoinTestFramework):
             # Money sent to P2SH of multisig of this should only be seen after importaddress with the BASE58 P2SH address.
 
             multisig_without_privkey_address = self.nodes[0].addmultisigaddress(2, [pubkeys[3], pubkeys[4]])['address']
-            script = CScript([OP_2, bytes.fromhex(pubkeys[3]), bytes.fromhex(pubkeys[4]), OP_2, OP_CHECKMULTISIG])
+            script = keys_to_multisig_script([pubkeys[3], pubkeys[4]])
             solvable_after_importaddress.append(script_to_p2sh_script(script))
 
             for i in compressed_spendable_address:
diff --git a/test/functional/mempool_accept.py b/test/functional/mempool_accept.py
index 2ee440bcb7..71be2b4a82 100755
--- a/test/functional/mempool_accept.py
+++ b/test/functional/mempool_accept.py
@@ -22,13 +22,11 @@ from test_framework.messages import (
 from test_framework.script import (
     CScript,
     OP_0,
-    OP_2,
-    OP_3,
-    OP_CHECKMULTISIG,
     OP_HASH160,
     OP_RETURN,
 )
 from test_framework.script_util import (
+    keys_to_multisig_script,
     script_to_p2sh_script,
 )
 from test_framework.util import (
@@ -283,7 +281,7 @@ class MempoolAcceptanceTest(BitcoinTestFramework):
         key = ECKey()
         key.generate()
         pubkey = key.get_pubkey().get_bytes()
-        tx.vout[0].scriptPubKey = CScript([OP_2, pubkey, pubkey, pubkey, OP_3, OP_CHECKMULTISIG])  # Some bare multisig script (2-of-3)
+        tx.vout[0].scriptPubKey = keys_to_multisig_script([pubkey] * 3, k=2)  # Some bare multisig script (2-of-3)
         self.check_mempool_result(
             result_expected=[{'txid': tx.rehash(), 'allowed': False, 'reject-reason': 'bare-multisig'}],
             rawtxs=[tx.serialize().hex()],
diff --git a/test/functional/test_framework/blocktools.py b/test/functional/test_framework/blocktools.py
index 85e3c2a383..5d0113465f 100644
--- a/test/functional/test_framework/blocktools.py
+++ b/test/functional/test_framework/blocktools.py
@@ -32,13 +32,13 @@ from .script import (
     CScriptNum,
     CScriptOp,
     OP_1,
-    OP_CHECKMULTISIG,
     OP_RETURN,
     OP_TRUE,
 )
 from .script_util import (
     key_to_p2pk_script,
     key_to_p2wpkh_script,
+    keys_to_multisig_script,
     script_to_p2wsh_script,
 )
 from .util import assert_equal
@@ -209,7 +209,7 @@ def witness_script(use_p2wsh, pubkey):
         pkscript = key_to_p2wpkh_script(pubkey)
     else:
         # 1-of-1 multisig
-        witness_script = CScript([OP_1, bytes.fromhex(pubkey), OP_1, OP_CHECKMULTISIG])
+        witness_script = keys_to_multisig_script([pubkey])
         pkscript = script_to_p2wsh_script(witness_script)
     return pkscript.hex()
 
@@ -218,7 +218,7 @@ def create_witness_tx(node, use_p2wsh, utxo, pubkey, encode_p2sh, amount):
 
     Optionally wrap the segwit output using P2SH."""
     if use_p2wsh:
-        program = CScript([OP_1, bytes.fromhex(pubkey), OP_1, OP_CHECKMULTISIG])
+        program = keys_to_multisig_script([pubkey])
         addr = script_to_p2sh_p2wsh(program) if encode_p2sh else script_to_p2wsh(program)
     else:
         addr = key_to_p2sh_p2wpkh(pubkey) if encode_p2sh else key_to_p2wpkh(pubkey)
diff --git a/test/functional/test_framework/script_util.py b/test/functional/test_framework/script_util.py
index 82a9067dd2..cbc4a560db 100755
--- a/test/functional/test_framework/script_util.py
+++ b/test/functional/test_framework/script_util.py
@@ -5,7 +5,9 @@
 """Useful Script constants and utils."""
 from test_framework.script import (
     CScript,
+    CScriptOp,
     OP_0,
+    OP_CHECKMULTISIG,
     OP_CHECKSIG,
     OP_DUP,
     OP_EQUAL,
@@ -41,6 +43,17 @@ def key_to_p2pk_script(key):
     return CScript([key, OP_CHECKSIG])
 
 
+def keys_to_multisig_script(keys, *, k=None):
+    n = len(keys)
+    if k is None:  # n-of-n multisig by default
+        k = n
+    assert k <= n
+    op_k = CScriptOp.encode_op_n(k)
+    op_n = CScriptOp.encode_op_n(n)
+    checked_keys = [check_key(key) for key in keys]
+    return CScript([op_k] + checked_keys + [op_n, OP_CHECKMULTISIG])
+
+
 def keyhash_to_p2pkh_script(hash):
     assert len(hash) == 20
     return CScript([OP_DUP, OP_HASH160, hash, OP_EQUALVERIFY, OP_CHECKSIG])
diff --git a/test/functional/test_framework/wallet_util.py b/test/functional/test_framework/wallet_util.py
index 1ee55aa3b7..c307ded542 100755
--- a/test/functional/test_framework/wallet_util.py
+++ b/test/functional/test_framework/wallet_util.py
@@ -15,15 +15,10 @@ from test_framework.address import (
     script_to_p2wsh,
 )
 from test_framework.key import ECKey
-from test_framework.script import (
-    CScript,
-    OP_2,
-    OP_3,
-    OP_CHECKMULTISIG,
-)
 from test_framework.script_util import (
     key_to_p2pkh_script,
     key_to_p2wpkh_script,
+    keys_to_multisig_script,
     script_to_p2sh_script,
     script_to_p2wsh_script,
 )
@@ -92,7 +87,7 @@ def get_multisig(node):
         addr = node.getaddressinfo(node.getnewaddress())
         addrs.append(addr['address'])
         pubkeys.append(addr['pubkey'])
-    script_code = CScript([OP_2] + [bytes.fromhex(pubkey) for pubkey in pubkeys] + [OP_3, OP_CHECKMULTISIG])
+    script_code = keys_to_multisig_script(pubkeys, k=2)
     witness_script = script_to_p2wsh_script(script_code)
     return Multisig(privkeys=[node.dumpprivkey(addr) for addr in addrs],
                     pubkeys=pubkeys,