Add script equivalent of functions in address.py

1 files changed, 58 insertions, 1 deletions

diff --git a/test/functional/test_framework/script_util.py b/test/functional/test_framework/script_util.py
index 80fbae70bf..318a438705 100755
--- a/test/functional/test_framework/script_util.py
+++ b/test/functional/test_framework/script_util.py
@@ -3,7 +3,8 @@
 # Distributed under the MIT software license, see the accompanying
 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
 """Useful Script constants and utils."""
-from test_framework.script import CScript
+from test_framework.script import CScript, hash160, sha256, OP_0, OP_DUP, OP_HASH160, OP_CHECKSIG, OP_EQUAL, OP_EQUALVERIFY
+from test_framework.util import hex_str_to_bytes
 
 # To prevent a "tx-size-small" policy rule error, a transaction has to have a
 # non-witness size of at least 82 bytes (MIN_STANDARD_TX_NONWITNESS_SIZE in
@@ -24,3 +25,59 @@ from test_framework.script import CScript
 # met.
 DUMMY_P2WPKH_SCRIPT = CScript([b'a' * 21])
 DUMMY_2_P2WPKH_SCRIPT = CScript([b'b' * 21])
+
+def keyhash_to_p2pkh_script(hash, main = False):
+    assert len(hash) == 20
+    return CScript([OP_DUP, OP_HASH160, hash, OP_EQUALVERIFY, OP_CHECKSIG])
+
+def scripthash_to_p2sh_script(hash, main = False):
+    assert len(hash) == 20
+    return CScript([OP_HASH160, hash, OP_EQUAL])
+
+def key_to_p2pkh_script(key, main = False):
+    key = check_key(key)
+    return keyhash_to_p2pkh_script(hash160(key), main)
+
+def script_to_p2sh_script(script, main = False):
+    script = check_script(script)
+    return scripthash_to_p2sh_script(hash160(script), main)
+
+def key_to_p2sh_p2wpkh_script(key, main = False):
+    key = check_key(key)
+    p2shscript = CScript([OP_0, hash160(key)])
+    return script_to_p2sh_script(p2shscript, main)
+
+def program_to_witness_script(version, program, main = False):
+    if isinstance(program, str):
+        program = hex_str_to_bytes(program)
+    assert 0 <= version <= 16
+    assert 2 <= len(program) <= 40
+    assert version > 0 or len(program) in [20, 32]
+    return CScript([version, program])
+
+def script_to_p2wsh_script(script, main = False):
+    script = check_script(script)
+    return program_to_witness_script(0, sha256(script), main)
+
+def key_to_p2wpkh_script(key, main = False):
+    key = check_key(key)
+    return program_to_witness_script(0, hash160(key), main)
+
+def script_to_p2sh_p2wsh_script(script, main = False):
+    script = check_script(script)
+    p2shscript = CScript([OP_0, sha256(script)])
+    return script_to_p2sh_script(p2shscript, main)
+
+def check_key(key):
+    if isinstance(key, str):
+        key = hex_str_to_bytes(key) # Assuming this is hex string
+    if isinstance(key, bytes) and (len(key) == 33 or len(key) == 65):
+        return key
+    assert False
+
+def check_script(script):
+    if isinstance(script, str):
+        script = hex_str_to_bytes(script) # Assuming this is hex string
+    if isinstance(script, bytes) or isinstance(script, CScript):
+        return script
+    assert False