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