wallet_util.py raw
1 #!/usr/bin/env python3
2 # Copyright (c) 2018-2021 The Limenka 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"""
6 from collections import namedtuple
7 import unittest
8
9 from test_framework.address import (
10 byte_to_base58,
11 key_to_p2pkh,
12 key_to_p2sh_p2wpkh,
13 key_to_p2wpkh,
14 script_to_p2sh,
15 script_to_p2sh_p2wsh,
16 script_to_p2wsh,
17 )
18 from test_framework.key import ECKey
19 from test_framework.messages import (
20 CTxIn,
21 CTxInWitness,
22 WITNESS_SCALE_FACTOR,
23 )
24 from test_framework.script_util import (
25 key_to_p2pkh_script,
26 key_to_p2wpkh_script,
27 keys_to_multisig_script,
28 script_to_p2sh_script,
29 script_to_p2wsh_script,
30 )
31
32 Key = namedtuple('Key', ['privkey',
33 'pubkey',
34 'p2pkh_script',
35 'p2pkh_addr',
36 'p2wpkh_script',
37 'p2wpkh_addr',
38 'p2sh_p2wpkh_script',
39 'p2sh_p2wpkh_redeem_script',
40 'p2sh_p2wpkh_addr'])
41
42 Multisig = namedtuple('Multisig', ['privkeys',
43 'pubkeys',
44 'p2sh_script',
45 'p2sh_addr',
46 'redeem_script',
47 'p2wsh_script',
48 'p2wsh_addr',
49 'p2sh_p2wsh_script',
50 'p2sh_p2wsh_addr'])
51
52 def get_key(node):
53 """Generate a fresh key on node
54
55 Returns a named tuple of privkey, pubkey and all address and scripts."""
56 addr = node.getnewaddress()
57 pubkey = node.getaddressinfo(addr)['pubkey']
58 return Key(privkey=node.dumpprivkey(addr),
59 pubkey=pubkey,
60 p2pkh_script=key_to_p2pkh_script(pubkey).hex(),
61 p2pkh_addr=key_to_p2pkh(pubkey),
62 p2wpkh_script=key_to_p2wpkh_script(pubkey).hex(),
63 p2wpkh_addr=key_to_p2wpkh(pubkey),
64 p2sh_p2wpkh_script=script_to_p2sh_script(key_to_p2wpkh_script(pubkey)).hex(),
65 p2sh_p2wpkh_redeem_script=key_to_p2wpkh_script(pubkey).hex(),
66 p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey))
67
68 def get_generate_key():
69 """Generate a fresh key
70
71 Returns a named tuple of privkey, pubkey and all address and scripts."""
72 privkey, pubkey = generate_keypair(wif=True)
73 return Key(privkey=privkey,
74 pubkey=pubkey.hex(),
75 p2pkh_script=key_to_p2pkh_script(pubkey).hex(),
76 p2pkh_addr=key_to_p2pkh(pubkey),
77 p2wpkh_script=key_to_p2wpkh_script(pubkey).hex(),
78 p2wpkh_addr=key_to_p2wpkh(pubkey),
79 p2sh_p2wpkh_script=script_to_p2sh_script(key_to_p2wpkh_script(pubkey)).hex(),
80 p2sh_p2wpkh_redeem_script=key_to_p2wpkh_script(pubkey).hex(),
81 p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey))
82
83 def get_multisig(node):
84 """Generate a fresh 2-of-3 multisig on node
85
86 Returns a named tuple of privkeys, pubkeys and all address and scripts."""
87 addrs = []
88 pubkeys = []
89 for _ in range(3):
90 addr = node.getaddressinfo(node.getnewaddress())
91 addrs.append(addr['address'])
92 pubkeys.append(addr['pubkey'])
93 script_code = keys_to_multisig_script(pubkeys, k=2)
94 witness_script = script_to_p2wsh_script(script_code)
95 return Multisig(privkeys=[node.dumpprivkey(addr) for addr in addrs],
96 pubkeys=pubkeys,
97 p2sh_script=script_to_p2sh_script(script_code).hex(),
98 p2sh_addr=script_to_p2sh(script_code),
99 redeem_script=script_code.hex(),
100 p2wsh_script=witness_script.hex(),
101 p2wsh_addr=script_to_p2wsh(script_code),
102 p2sh_p2wsh_script=script_to_p2sh_script(witness_script).hex(),
103 p2sh_p2wsh_addr=script_to_p2sh_p2wsh(script_code))
104
105 def test_address(node, address, **kwargs):
106 """Get address info for `address` and test whether the returned values are as expected."""
107 addr_info = node.getaddressinfo(address)
108 for key, value in kwargs.items():
109 if value is None:
110 if key in addr_info.keys():
111 raise AssertionError("key {} unexpectedly returned in getaddressinfo.".format(key))
112 elif addr_info[key] != value:
113 raise AssertionError("key {} value {} did not match expected value {}".format(key, addr_info[key], value))
114
115 def bytes_to_wif(b, compressed=True):
116 if compressed:
117 b += b'\x01'
118 return byte_to_base58(b, 239)
119
120 def generate_keypair(compressed=True, wif=False):
121 """Generate a new random keypair and return the corresponding ECKey /
122 bytes objects. The private key can also be provided as WIF (wallet
123 import format) string instead, which is often useful for wallet RPC
124 interaction."""
125 privkey = ECKey()
126 privkey.generate(compressed)
127 pubkey = privkey.get_pubkey().get_bytes()
128 if wif:
129 privkey = bytes_to_wif(privkey.get_bytes(), compressed)
130 return privkey, pubkey
131
132 def calculate_input_weight(scriptsig_hex, witness_stack_hex=None):
133 """Given a scriptSig and a list of witness stack items for an input in hex format,
134 calculate the total input weight. If the input has no witness data,
135 `witness_stack_hex` can be set to None."""
136 tx_in = CTxIn(scriptSig=bytes.fromhex(scriptsig_hex))
137 witness_size = 0
138 if witness_stack_hex is not None:
139 tx_inwit = CTxInWitness()
140 for witness_item_hex in witness_stack_hex:
141 tx_inwit.scriptWitness.stack.append(bytes.fromhex(witness_item_hex))
142 witness_size = len(tx_inwit.serialize())
143 return len(tx_in.serialize()) * WITNESS_SCALE_FACTOR + witness_size
144
145 class WalletUnlock():
146 """
147 A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.
148 """
149
150 MAXIMUM_TIMEOUT = 999000
151
152 def __init__(self, wallet, passphrase, timeout=MAXIMUM_TIMEOUT):
153 self.wallet = wallet
154 self.passphrase = passphrase
155 self.timeout = timeout
156
157 def __enter__(self):
158 self.wallet.walletpassphrase(self.passphrase, self.timeout)
159
160 def __exit__(self, *args):
161 _ = args
162 self.wallet.walletlock()
163
164
165 class TestFrameworkWalletUtil(unittest.TestCase):
166 def test_calculate_input_weight(self):
167 SKELETON_BYTES = 32 + 4 + 4 # prevout-txid, prevout-index, sequence
168 SMALL_LEN_BYTES = 1 # bytes needed for encoding scriptSig / witness item lengths < 253
169 LARGE_LEN_BYTES = 3 # bytes needed for encoding scriptSig / witness item lengths >= 253
170
171 # empty scriptSig, no witness
172 self.assertEqual(calculate_input_weight(""),
173 (SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR)
174 self.assertEqual(calculate_input_weight("", None),
175 (SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR)
176 # small scriptSig, no witness
177 scriptSig_small = "00"*252
178 self.assertEqual(calculate_input_weight(scriptSig_small, None),
179 (SKELETON_BYTES + SMALL_LEN_BYTES + 252) * WITNESS_SCALE_FACTOR)
180 # small scriptSig, empty witness stack
181 self.assertEqual(calculate_input_weight(scriptSig_small, []),
182 (SKELETON_BYTES + SMALL_LEN_BYTES + 252) * WITNESS_SCALE_FACTOR + SMALL_LEN_BYTES)
183 # large scriptSig, no witness
184 scriptSig_large = "00"*253
185 self.assertEqual(calculate_input_weight(scriptSig_large, None),
186 (SKELETON_BYTES + LARGE_LEN_BYTES + 253) * WITNESS_SCALE_FACTOR)
187 # large scriptSig, empty witness stack
188 self.assertEqual(calculate_input_weight(scriptSig_large, []),
189 (SKELETON_BYTES + LARGE_LEN_BYTES + 253) * WITNESS_SCALE_FACTOR + SMALL_LEN_BYTES)
190 # empty scriptSig, 5 small witness stack items
191 self.assertEqual(calculate_input_weight("", ["00", "11", "22", "33", "44"]),
192 ((SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR) + SMALL_LEN_BYTES + 5 * SMALL_LEN_BYTES + 5)
193 # empty scriptSig, 253 small witness stack items
194 self.assertEqual(calculate_input_weight("", ["00"]*253),
195 ((SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR) + LARGE_LEN_BYTES + 253 * SMALL_LEN_BYTES + 253)
196 # small scriptSig, 3 large witness stack items
197 self.assertEqual(calculate_input_weight(scriptSig_small, ["00"*253]*3),
198 ((SKELETON_BYTES + SMALL_LEN_BYTES + 252) * WITNESS_SCALE_FACTOR) + SMALL_LEN_BYTES + 3 * LARGE_LEN_BYTES + 3*253)
199 # large scriptSig, 3 large witness stack items
200 self.assertEqual(calculate_input_weight(scriptSig_large, ["00"*253]*3),
201 ((SKELETON_BYTES + LARGE_LEN_BYTES + 253) * WITNESS_SCALE_FACTOR) + SMALL_LEN_BYTES + 3 * LARGE_LEN_BYTES + 3*253)
202