1#!/usr/bin/env python3 2# Copyright (c) 2018-2020 The Bitcoin Core developers 3# Distributed under the MIT software license, see the accompanying 4# file COPYING or http://www.opensource.org/licenses/mit-license.php. 5"""Useful util functions for testing the wallet""" 6from collections import namedtuple 7 8from test_framework.address import ( 9 byte_to_base58, 10 key_to_p2pkh, 11 key_to_p2sh_p2wpkh, 12 key_to_p2wpkh, 13 script_to_p2sh, 14 script_to_p2sh_p2wsh, 15 script_to_p2wsh, 16) 17from test_framework.key import ECKey 18from test_framework.script import ( 19 CScript, 20 OP_0, 21 OP_2, 22 OP_3, 23 OP_CHECKMULTISIG, 24 OP_CHECKSIG, 25 OP_DUP, 26 OP_EQUAL, 27 OP_EQUALVERIFY, 28 OP_HASH160, 29 hash160, 30 sha256, 31) 32from test_framework.util import hex_str_to_bytes 33 34Key = namedtuple('Key', ['privkey', 35 'pubkey', 36 'p2pkh_script', 37 'p2pkh_addr', 38 'p2wpkh_script', 39 'p2wpkh_addr', 40 'p2sh_p2wpkh_script', 41 'p2sh_p2wpkh_redeem_script', 42 'p2sh_p2wpkh_addr']) 43 44Multisig = namedtuple('Multisig', ['privkeys', 45 'pubkeys', 46 'p2sh_script', 47 'p2sh_addr', 48 'redeem_script', 49 'p2wsh_script', 50 'p2wsh_addr', 51 'p2sh_p2wsh_script', 52 'p2sh_p2wsh_addr']) 53 54def get_key(node): 55 """Generate a fresh key on node 56 57 Returns a named tuple of privkey, pubkey and all address and scripts.""" 58 addr = node.getnewaddress() 59 pubkey = node.getaddressinfo(addr)['pubkey'] 60 pkh = hash160(hex_str_to_bytes(pubkey)) 61 return Key(privkey=node.dumpprivkey(addr), 62 pubkey=pubkey, 63 p2pkh_script=CScript([OP_DUP, OP_HASH160, pkh, OP_EQUALVERIFY, OP_CHECKSIG]).hex(), 64 p2pkh_addr=key_to_p2pkh(pubkey), 65 p2wpkh_script=CScript([OP_0, pkh]).hex(), 66 p2wpkh_addr=key_to_p2wpkh(pubkey), 67 p2sh_p2wpkh_script=CScript([OP_HASH160, hash160(CScript([OP_0, pkh])), OP_EQUAL]).hex(), 68 p2sh_p2wpkh_redeem_script=CScript([OP_0, pkh]).hex(), 69 p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey)) 70 71def get_generate_key(): 72 """Generate a fresh key 73 74 Returns a named tuple of privkey, pubkey and all address and scripts.""" 75 eckey = ECKey() 76 eckey.generate() 77 privkey = bytes_to_wif(eckey.get_bytes()) 78 pubkey = eckey.get_pubkey().get_bytes().hex() 79 pkh = hash160(hex_str_to_bytes(pubkey)) 80 return Key(privkey=privkey, 81 pubkey=pubkey, 82 p2pkh_script=CScript([OP_DUP, OP_HASH160, pkh, OP_EQUALVERIFY, OP_CHECKSIG]).hex(), 83 p2pkh_addr=key_to_p2pkh(pubkey), 84 p2wpkh_script=CScript([OP_0, pkh]).hex(), 85 p2wpkh_addr=key_to_p2wpkh(pubkey), 86 p2sh_p2wpkh_script=CScript([OP_HASH160, hash160(CScript([OP_0, pkh])), OP_EQUAL]).hex(), 87 p2sh_p2wpkh_redeem_script=CScript([OP_0, pkh]).hex(), 88 p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey)) 89 90def get_multisig(node): 91 """Generate a fresh 2-of-3 multisig on node 92 93 Returns a named tuple of privkeys, pubkeys and all address and scripts.""" 94 addrs = [] 95 pubkeys = [] 96 for _ in range(3): 97 addr = node.getaddressinfo(node.getnewaddress()) 98 addrs.append(addr['address']) 99 pubkeys.append(addr['pubkey']) 100 script_code = CScript([OP_2] + [hex_str_to_bytes(pubkey) for pubkey in pubkeys] + [OP_3, OP_CHECKMULTISIG]) 101 witness_script = CScript([OP_0, sha256(script_code)]) 102 return Multisig(privkeys=[node.dumpprivkey(addr) for addr in addrs], 103 pubkeys=pubkeys, 104 p2sh_script=CScript([OP_HASH160, hash160(script_code), OP_EQUAL]).hex(), 105 p2sh_addr=script_to_p2sh(script_code), 106 redeem_script=script_code.hex(), 107 p2wsh_script=witness_script.hex(), 108 p2wsh_addr=script_to_p2wsh(script_code), 109 p2sh_p2wsh_script=CScript([OP_HASH160, witness_script, OP_EQUAL]).hex(), 110 p2sh_p2wsh_addr=script_to_p2sh_p2wsh(script_code)) 111 112def test_address(node, address, **kwargs): 113 """Get address info for `address` and test whether the returned values are as expected.""" 114 addr_info = node.getaddressinfo(address) 115 for key, value in kwargs.items(): 116 if value is None: 117 if key in addr_info.keys(): 118 raise AssertionError("key {} unexpectedly returned in getaddressinfo.".format(key)) 119 elif addr_info[key] != value: 120 raise AssertionError("key {} value {} did not match expected value {}".format(key, addr_info[key], value)) 121 122def bytes_to_wif(b, compressed=True): 123 if compressed: 124 b += b'\x01' 125 return byte_to_base58(b, 239) 126 127def generate_wif_key(): 128 # Makes a WIF privkey for imports 129 k = ECKey() 130 k.generate() 131 return bytes_to_wif(k.get_bytes(), k.is_compressed) 132