1  #!/usr/bin/env python3
  2  # Copyright (c) 2018-2021 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  
  8  from 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  )
 17  from test_framework.key import ECKey
 18  from test_framework.script_util import (
 19      key_to_p2pkh_script,
 20      key_to_p2wpkh_script,
 21      keys_to_multisig_script,
 22      script_to_p2sh_script,
 23      script_to_p2wsh_script,
 24  )
 25  
 26  Key = namedtuple('Key', ['privkey',
 27                           'pubkey',
 28                           'p2pkh_script',
 29                           'p2pkh_addr',
 30                           'p2wpkh_script',
 31                           'p2wpkh_addr',
 32                           'p2sh_p2wpkh_script',
 33                           'p2sh_p2wpkh_redeem_script',
 34                           'p2sh_p2wpkh_addr'])
 35  
 36  Multisig = namedtuple('Multisig', ['privkeys',
 37                                     'pubkeys',
 38                                     'p2sh_script',
 39                                     'p2sh_addr',
 40                                     'redeem_script',
 41                                     'p2wsh_script',
 42                                     'p2wsh_addr',
 43                                     'p2sh_p2wsh_script',
 44                                     'p2sh_p2wsh_addr'])
 45  
 46  def get_key(node):
 47      """Generate a fresh key on node
 48  
 49      Returns a named tuple of privkey, pubkey and all address and scripts."""
 50      addr = node.getnewaddress()
 51      pubkey = node.getaddressinfo(addr)['pubkey']
 52      return Key(privkey=node.dumpprivkey(addr),
 53                 pubkey=pubkey,
 54                 p2pkh_script=key_to_p2pkh_script(pubkey).hex(),
 55                 p2pkh_addr=key_to_p2pkh(pubkey),
 56                 p2wpkh_script=key_to_p2wpkh_script(pubkey).hex(),
 57                 p2wpkh_addr=key_to_p2wpkh(pubkey),
 58                 p2sh_p2wpkh_script=script_to_p2sh_script(key_to_p2wpkh_script(pubkey)).hex(),
 59                 p2sh_p2wpkh_redeem_script=key_to_p2wpkh_script(pubkey).hex(),
 60                 p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey))
 61  
 62  def get_generate_key():
 63      """Generate a fresh key
 64  
 65      Returns a named tuple of privkey, pubkey and all address and scripts."""
 66      privkey, pubkey = generate_keypair(wif=True)
 67      return Key(privkey=privkey,
 68                 pubkey=pubkey.hex(),
 69                 p2pkh_script=key_to_p2pkh_script(pubkey).hex(),
 70                 p2pkh_addr=key_to_p2pkh(pubkey),
 71                 p2wpkh_script=key_to_p2wpkh_script(pubkey).hex(),
 72                 p2wpkh_addr=key_to_p2wpkh(pubkey),
 73                 p2sh_p2wpkh_script=script_to_p2sh_script(key_to_p2wpkh_script(pubkey)).hex(),
 74                 p2sh_p2wpkh_redeem_script=key_to_p2wpkh_script(pubkey).hex(),
 75                 p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey))
 76  
 77  def get_multisig(node):
 78      """Generate a fresh 2-of-3 multisig on node
 79  
 80      Returns a named tuple of privkeys, pubkeys and all address and scripts."""
 81      addrs = []
 82      pubkeys = []
 83      for _ in range(3):
 84          addr = node.getaddressinfo(node.getnewaddress())
 85          addrs.append(addr['address'])
 86          pubkeys.append(addr['pubkey'])
 87      script_code = keys_to_multisig_script(pubkeys, k=2)
 88      witness_script = script_to_p2wsh_script(script_code)
 89      return Multisig(privkeys=[node.dumpprivkey(addr) for addr in addrs],
 90                      pubkeys=pubkeys,
 91                      p2sh_script=script_to_p2sh_script(script_code).hex(),
 92                      p2sh_addr=script_to_p2sh(script_code),
 93                      redeem_script=script_code.hex(),
 94                      p2wsh_script=witness_script.hex(),
 95                      p2wsh_addr=script_to_p2wsh(script_code),
 96                      p2sh_p2wsh_script=script_to_p2sh_script(witness_script).hex(),
 97                      p2sh_p2wsh_addr=script_to_p2sh_p2wsh(script_code))
 98  
 99  def test_address(node, address, **kwargs):
100      """Get address info for `address` and test whether the returned values are as expected."""
101      addr_info = node.getaddressinfo(address)
102      for key, value in kwargs.items():
103          if value is None:
104              if key in addr_info.keys():
105                  raise AssertionError("key {} unexpectedly returned in getaddressinfo.".format(key))
106          elif addr_info[key] != value:
107              raise AssertionError("key {} value {} did not match expected value {}".format(key, addr_info[key], value))
108  
109  def bytes_to_wif(b, compressed=True):
110      if compressed:
111          b += b'\x01'
112      return byte_to_base58(b, 239)
113  
114  def generate_keypair(compressed=True, wif=False):
115      """Generate a new random keypair and return the corresponding ECKey /
116      bytes objects. The private key can also be provided as WIF (wallet
117      import format) string instead, which is often useful for wallet RPC
118      interaction."""
119      privkey = ECKey()
120      privkey.generate(compressed)
121      pubkey = privkey.get_pubkey().get_bytes()
122      if wif:
123          privkey = bytes_to_wif(privkey.get_bytes(), compressed)
124      return privkey, pubkey
125  
126  class WalletUnlock():
127      """
128      A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.
129      """
130  
131      MAXIMUM_TIMEOUT = 999000
132  
133      def __init__(self, wallet, passphrase, timeout=MAXIMUM_TIMEOUT):
134          self.wallet = wallet
135          self.passphrase = passphrase
136          self.timeout = timeout
137  
138      def __enter__(self):
139          self.wallet.walletpassphrase(self.passphrase, self.timeout)
140  
141      def __exit__(self, *args):
142          _ = args
143          self.wallet.walletlock()