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