wallet.py raw

   1  #!/usr/bin/env python3
   2  # Copyright (c) 2020-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  """A limited-functionality wallet, which may replace a real wallet in tests"""
   6  
   7  from copy import deepcopy
   8  from decimal import Decimal
   9  from enum import Enum
  10  from typing import (
  11      Any,
  12      Optional,
  13  )
  14  from test_framework.address import (
  15      address_to_scriptpubkey,
  16      create_deterministic_address_bcrt1_p2tr_op_true,
  17      key_to_p2pkh,
  18      key_to_p2sh_p2wpkh,
  19      key_to_p2wpkh,
  20      output_key_to_p2tr,
  21  )
  22  from test_framework.blocktools import COINBASE_MATURITY
  23  from test_framework.descriptors import descsum_create
  24  from test_framework.key import (
  25      ECKey,
  26      compute_xonly_pubkey,
  27  )
  28  from test_framework.messages import (
  29      COIN,
  30      COutPoint,
  31      CTransaction,
  32      CTxIn,
  33      CTxInWitness,
  34      CTxOut,
  35      hash256,
  36  )
  37  from test_framework.script import (
  38      CScript,
  39      OP_NOP,
  40      OP_RETURN,
  41      OP_TRUE,
  42      sign_input_legacy,
  43      taproot_construct,
  44  )
  45  from test_framework.script_util import (
  46      bulk_vout,
  47      key_to_p2pk_script,
  48      key_to_p2pkh_script,
  49      key_to_p2sh_p2wpkh_script,
  50      key_to_p2wpkh_script,
  51  )
  52  from test_framework.util import (
  53      assert_equal,
  54      assert_greater_than_or_equal,
  55      get_fee,
  56  )
  57  from test_framework.wallet_util import (
  58      bytes_to_wif,
  59      generate_keypair,
  60  )
  61  
  62  DEFAULT_FEE = Decimal("0.0001")
  63  
  64  class MiniWalletMode(Enum):
  65      """Determines the transaction type the MiniWallet is creating and spending.
  66  
  67      For most purposes, the default mode ADDRESS_OP_TRUE should be sufficient;
  68      it simply uses a fixed bech32m P2TR address whose coins are spent with a
  69      witness stack of OP_TRUE, i.e. following an anyone-can-spend policy.
  70      However, if the transactions need to be modified by the user (e.g. prepending
  71      scriptSig for testing opcodes that are activated by a soft-fork), or the txs
  72      should contain an actual signature, the raw modes RAW_OP_TRUE and RAW_P2PK
  73      can be useful. In order to avoid mixing of UTXOs between different MiniWallet
  74      instances, a tag name can be passed to the default mode, to create different
  75      output scripts. Note that the UTXOs from the pre-generated test chain can
  76      only be spent if no tag is passed. Summary of modes:
  77  
  78                      |      output       |           |  tx is   | can modify |  needs
  79           mode       |    description    |  address  | standard | scriptSig  | signing
  80      ----------------+-------------------+-----------+----------+------------+----------
  81      ADDRESS_OP_TRUE | anyone-can-spend  |  bech32m  |   yes    |    no      |   no
  82      RAW_OP_TRUE     | anyone-can-spend  |  - (raw)  |   no     |    yes     |   no
  83      RAW_P2PK        | pay-to-public-key |  - (raw)  |   yes    |    yes     |   yes
  84      """
  85      ADDRESS_OP_TRUE = 1
  86      RAW_OP_TRUE = 2
  87      RAW_P2PK = 3
  88  
  89  
  90  class MiniWallet:
  91      def __init__(self, test_node, *, mode=MiniWalletMode.ADDRESS_OP_TRUE, tag_name=None):
  92          self._test_node = test_node
  93          self._utxos = []
  94          self._mode = mode
  95  
  96          assert isinstance(mode, MiniWalletMode)
  97          if mode == MiniWalletMode.RAW_OP_TRUE:
  98              assert tag_name is None
  99              self._scriptPubKey = bytes(CScript([OP_TRUE]))
 100          elif mode == MiniWalletMode.RAW_P2PK:
 101              # use simple deterministic private key (k=1)
 102              assert tag_name is None
 103              self._priv_key = ECKey()
 104              self._priv_key.set((1).to_bytes(32, 'big'), True)
 105              pub_key = self._priv_key.get_pubkey()
 106              self._scriptPubKey = key_to_p2pk_script(pub_key.get_bytes())
 107          elif mode == MiniWalletMode.ADDRESS_OP_TRUE:
 108              internal_key = None if tag_name is None else compute_xonly_pubkey(hash256(tag_name.encode()))[0]
 109              self._address, self._taproot_info = create_deterministic_address_bcrt1_p2tr_op_true(internal_key)
 110              self._scriptPubKey = address_to_scriptpubkey(self._address)
 111  
 112          # When the pre-mined test framework chain is used, it contains coinbase
 113          # outputs to the MiniWallet's default address in blocks 76-100
 114          # (see method BitcoinTestFramework._initialize_chain())
 115          # The MiniWallet needs to rescan_utxos() in order to account
 116          # for those mature UTXOs, so that all txs spend confirmed coins
 117          self.rescan_utxos()
 118  
 119      def _create_utxo(self, *, txid, vout, value, height, coinbase, confirmations):
 120          return {"txid": txid, "vout": vout, "value": value, "height": height, "coinbase": coinbase, "confirmations": confirmations}
 121  
 122      def _bulk_tx(self, tx, target_vsize):
 123          """Pad a transaction with extra outputs until it reaches a target vsize.
 124          returns the tx
 125          """
 126          tx.vout.append(CTxOut(nValue=0, scriptPubKey=CScript([OP_RETURN])))
 127          bulk_vout(tx, target_vsize)
 128  
 129  
 130      def get_balance(self):
 131          return sum(u['value'] for u in self._utxos)
 132  
 133      def rescan_utxos(self, *, include_mempool=True):
 134          """Drop all utxos and rescan the utxo set"""
 135          self._utxos = []
 136          res = self._test_node.scantxoutset(action="start", scanobjects=[self.get_descriptor()])
 137          assert_equal(True, res['success'])
 138          for utxo in res['unspents']:
 139              self._utxos.append(
 140                  self._create_utxo(txid=utxo["txid"],
 141                                    vout=utxo["vout"],
 142                                    value=utxo["amount"],
 143                                    height=utxo["height"],
 144                                    coinbase=utxo["coinbase"],
 145                                    confirmations=res["height"] - utxo["height"] + 1))
 146          if include_mempool:
 147              mempool = self._test_node.getrawmempool(verbose=True)
 148              # Sort tx by ancestor count. See BlockAssembler::SortForBlock in src/node/miner.cpp
 149              sorted_mempool = sorted(mempool.items(), key=lambda item: (item[1]["ancestorcount"], int(item[0], 16)))
 150              for txid, _ in sorted_mempool:
 151                  self.scan_tx(self._test_node.getrawtransaction(txid=txid, verbose=True))
 152  
 153      def scan_tx(self, tx):
 154          """Scan the tx and adjust the internal list of owned utxos"""
 155          for spent in tx["vin"]:
 156              # Mark spent. This may happen when the caller has ownership of a
 157              # utxo that remained in this wallet. For example, by passing
 158              # mark_as_spent=False to get_utxo or by using an utxo returned by a
 159              # create_self_transfer* call.
 160              try:
 161                  self.get_utxo(txid=spent["txid"], vout=spent["vout"])
 162              except StopIteration:
 163                  pass
 164          for out in tx['vout']:
 165              if out['scriptPubKey']['hex'] == self._scriptPubKey.hex():
 166                  self._utxos.append(self._create_utxo(txid=tx["txid"], vout=out["n"], value=out["value"], height=0, coinbase=False, confirmations=0))
 167  
 168      def scan_txs(self, txs):
 169          for tx in txs:
 170              self.scan_tx(tx)
 171  
 172      def sign_tx(self, tx, fixed_length=True):
 173          if self._mode == MiniWalletMode.RAW_P2PK:
 174              # for exact fee calculation, create only signatures with fixed size by default (>49.89% probability):
 175              # 65 bytes: high-R val (33 bytes) + low-S val (32 bytes)
 176              # with the DER header/skeleton data of 6 bytes added, plus 2 bytes scriptSig overhead
 177              # (OP_PUSHn and SIGHASH_ALL), this leads to a scriptSig target size of 73 bytes
 178              tx.vin[0].scriptSig = b''
 179              while not len(tx.vin[0].scriptSig) == 73:
 180                  tx.vin[0].scriptSig = b''
 181                  sign_input_legacy(tx, 0, self._scriptPubKey, self._priv_key)
 182                  if not fixed_length:
 183                      break
 184          elif self._mode == MiniWalletMode.RAW_OP_TRUE:
 185              for i in tx.vin:
 186                  i.scriptSig = CScript([OP_NOP] * 43)  # pad to identical size
 187          elif self._mode == MiniWalletMode.ADDRESS_OP_TRUE:
 188              tx.wit.vtxinwit = [CTxInWitness()] * len(tx.vin)
 189              for i in tx.wit.vtxinwit:
 190                  assert_equal(len(self._taproot_info.leaves), 1)
 191                  leaf_info = list(self._taproot_info.leaves.values())[0]
 192                  i.scriptWitness.stack = [
 193                      leaf_info.script,
 194                      bytes([leaf_info.version | self._taproot_info.negflag]) + self._taproot_info.internal_pubkey,
 195                  ]
 196          else:
 197              assert False
 198  
 199      def generate(self, num_blocks, **kwargs):
 200          """Generate blocks with coinbase outputs to the internal address, and call rescan_utxos"""
 201          blocks = self._test_node.generatetodescriptor(num_blocks, self.get_descriptor(), **kwargs)
 202          # Calling rescan_utxos here makes sure that after a generate the utxo
 203          # set is in a clean state. For example, the wallet will update
 204          # - if the caller consumed utxos, but never used them
 205          # - if the caller sent a transaction that is not mined or got rbf'd
 206          # - after block re-orgs
 207          # - the utxo height for mined mempool txs
 208          # - However, the wallet will not consider remaining mempool txs
 209          self.rescan_utxos()
 210          return blocks
 211  
 212      def get_output_script(self):
 213          return self._scriptPubKey
 214  
 215      def get_descriptor(self):
 216          return descsum_create(f'raw({self._scriptPubKey.hex()})')
 217  
 218      def get_address(self):
 219          assert_equal(self._mode, MiniWalletMode.ADDRESS_OP_TRUE)
 220          return self._address
 221  
 222      def get_utxo(self, *, txid: str = '', vout: Optional[int] = None, mark_as_spent=True, confirmed_only=False) -> dict:
 223          """
 224          Returns a utxo and marks it as spent (pops it from the internal list)
 225  
 226          Args:
 227          txid: get the first utxo we find from a specific transaction
 228          """
 229          self._utxos = sorted(self._utxos, key=lambda k: (k['value'], -k['height']))  # Put the largest utxo last
 230          blocks_height = self._test_node.getblockchaininfo()['blocks']
 231          mature_coins = list(filter(lambda utxo: not utxo['coinbase'] or COINBASE_MATURITY - 1 <= blocks_height - utxo['height'], self._utxos))
 232          if txid:
 233              utxo_filter: Any = filter(lambda utxo: txid == utxo['txid'], self._utxos)
 234          else:
 235              utxo_filter = reversed(mature_coins)  # By default the largest utxo
 236          if vout is not None:
 237              utxo_filter = filter(lambda utxo: vout == utxo['vout'], utxo_filter)
 238          if confirmed_only:
 239              utxo_filter = filter(lambda utxo: utxo['confirmations'] > 0, utxo_filter)
 240          index = self._utxos.index(next(utxo_filter))
 241          if mark_as_spent:
 242              return self._utxos.pop(index)
 243          else:
 244              return self._utxos[index]
 245  
 246      def get_utxos(self, *, include_immature_coinbase=False, mark_as_spent=True, confirmed_only=False):
 247          """Returns the list of all utxos and optionally mark them as spent"""
 248          if not include_immature_coinbase:
 249              blocks_height = self._test_node.getblockchaininfo()['blocks']
 250              utxo_filter = filter(lambda utxo: not utxo['coinbase'] or COINBASE_MATURITY - 1 <= blocks_height - utxo['height'], self._utxos)
 251          else:
 252              utxo_filter = self._utxos
 253          if confirmed_only:
 254              utxo_filter = filter(lambda utxo: utxo['confirmations'] > 0, utxo_filter)
 255          utxos = deepcopy(list(utxo_filter))
 256          if mark_as_spent:
 257              self._utxos = []
 258          return utxos
 259  
 260      def send_self_transfer(self, *, from_node, **kwargs):
 261          """Call create_self_transfer and send the transaction."""
 262          tx = self.create_self_transfer(**kwargs)
 263          self.sendrawtransaction(from_node=from_node, tx_hex=tx['hex'])
 264          return tx
 265  
 266      def send_to(self, *, from_node, scriptPubKey, amount, fee=1000):
 267          """
 268          Create and send a tx with an output to a given scriptPubKey/amount,
 269          plus a change output to our internal address. To keep things simple, a
 270          fixed fee given in Satoshi is used.
 271  
 272          Note that this method fails if there is no single internal utxo
 273          available that can cover the cost for the amount and the fixed fee
 274          (the utxo with the largest value is taken).
 275          """
 276          tx = self.create_self_transfer(fee_rate=0)["tx"]
 277          assert_greater_than_or_equal(tx.vout[0].nValue, amount + fee)
 278          tx.vout[0].nValue -= (amount + fee)           # change output -> MiniWallet
 279          tx.vout.append(CTxOut(amount, scriptPubKey))  # arbitrary output -> to be returned
 280          txid = self.sendrawtransaction(from_node=from_node, tx_hex=tx.serialize().hex())
 281          return {
 282              "sent_vout": 1,
 283              "txid": txid,
 284              "wtxid": tx.wtxid_hex,
 285              "hex": tx.serialize().hex(),
 286              "tx": tx,
 287          }
 288  
 289      def send_self_transfer_multi(self, *, from_node, **kwargs):
 290          """Call create_self_transfer_multi and send the transaction."""
 291          tx = self.create_self_transfer_multi(**kwargs)
 292          self.sendrawtransaction(from_node=from_node, tx_hex=tx["hex"])
 293          return tx
 294  
 295      def create_self_transfer_multi(
 296          self,
 297          *,
 298          utxos_to_spend: Optional[list[dict]] = None,
 299          num_outputs=1,
 300          amount_per_output=0,
 301          version=2,
 302          locktime=0,
 303          sequence=0,
 304          fee_per_output=1000,
 305          target_vsize=0,
 306          confirmed_only=False,
 307      ):
 308          """
 309          Create and return a transaction that spends the given UTXOs and creates a
 310          certain number of outputs with equal amounts. The output amounts can be
 311          set by amount_per_output or automatically calculated with a fee_per_output.
 312          """
 313          utxos_to_spend = utxos_to_spend or [self.get_utxo(confirmed_only=confirmed_only)]
 314          sequence = [sequence] * len(utxos_to_spend) if type(sequence) is int else sequence
 315          assert_equal(len(utxos_to_spend), len(sequence))
 316  
 317          # calculate output amount
 318          inputs_value_total = sum([int(COIN * utxo['value']) for utxo in utxos_to_spend])
 319          outputs_value_total = inputs_value_total - fee_per_output * num_outputs
 320          amount_per_output = amount_per_output or (outputs_value_total // num_outputs)
 321          assert amount_per_output > 0
 322          outputs_value_total = amount_per_output * num_outputs
 323          fee = Decimal(inputs_value_total - outputs_value_total) / COIN
 324  
 325          # create tx
 326          tx = CTransaction()
 327          tx.vin = [CTxIn(COutPoint(int(utxo_to_spend['txid'], 16), utxo_to_spend['vout']), nSequence=seq) for utxo_to_spend, seq in zip(utxos_to_spend, sequence)]
 328          tx.vout = [CTxOut(amount_per_output, bytearray(self._scriptPubKey)) for _ in range(num_outputs)]
 329          tx.version = version
 330          tx.nLockTime = locktime
 331  
 332          self.sign_tx(tx)
 333  
 334          if target_vsize:
 335              self._bulk_tx(tx, target_vsize)
 336  
 337          txid = tx.txid_hex
 338          return {
 339              "new_utxos": [self._create_utxo(
 340                  txid=txid,
 341                  vout=i,
 342                  value=Decimal(tx.vout[i].nValue) / COIN,
 343                  height=0,
 344                  coinbase=False,
 345                  confirmations=0,
 346              ) for i in range(len(tx.vout))],
 347              "fee": fee,
 348              "txid": txid,
 349              "wtxid": tx.wtxid_hex,
 350              "hex": tx.serialize().hex(),
 351              "tx": tx,
 352          }
 353  
 354      def create_self_transfer(
 355              self,
 356              *,
 357              fee_rate=Decimal("0.003"),
 358              fee=Decimal("0"),
 359              utxo_to_spend=None,
 360              target_vsize=0,
 361              confirmed_only=False,
 362              **kwargs,
 363      ):
 364          """Create and return a tx with the specified fee. If fee is 0, use fee_rate, where the resulting fee may be exact or at most one satoshi higher than needed."""
 365          utxo_to_spend = utxo_to_spend or self.get_utxo(confirmed_only=confirmed_only)
 366          assert fee_rate >= 0
 367          assert fee >= 0
 368          # calculate fee
 369          if self._mode in (MiniWalletMode.RAW_OP_TRUE, MiniWalletMode.ADDRESS_OP_TRUE):
 370              vsize = Decimal(104)  # anyone-can-spend
 371          elif self._mode == MiniWalletMode.RAW_P2PK:
 372              vsize = Decimal(168)  # P2PK (73 bytes scriptSig + 35 bytes scriptPubKey + 60 bytes other)
 373          else:
 374              assert False
 375          if target_vsize and not fee:  # respect fee_rate if target vsize is passed
 376              fee = get_fee(target_vsize, fee_rate)
 377          send_value = utxo_to_spend["value"] - (fee or (fee_rate * vsize / 1000))
 378          if send_value <= 0:
 379              raise RuntimeError(f"UTXO value {utxo_to_spend['value']} is too small to cover fees {(fee or (fee_rate * vsize / 1000))}")
 380          # create tx
 381          tx = self.create_self_transfer_multi(
 382              utxos_to_spend=[utxo_to_spend],
 383              amount_per_output=int(COIN * send_value),
 384              target_vsize=target_vsize,
 385              **kwargs,
 386          )
 387          if not target_vsize:
 388              assert_equal(tx["tx"].get_vsize(), vsize)
 389          tx["new_utxo"] = tx.pop("new_utxos")[0]
 390  
 391          return tx
 392  
 393      def sendrawtransaction(self, *, from_node, tx_hex, maxfeerate=0, **kwargs):
 394          txid = from_node.sendrawtransaction(hexstring=tx_hex, maxfeerate=maxfeerate, **kwargs)
 395          self.scan_tx(from_node.decoderawtransaction(tx_hex))
 396          return txid
 397  
 398      def create_self_transfer_chain(self, *, chain_length, utxo_to_spend=None):
 399          """
 400          Create a "chain" of chain_length transactions. The nth transaction in
 401          the chain is a child of the n-1th transaction and parent of the n+1th transaction.
 402          """
 403          chaintip_utxo = utxo_to_spend or self.get_utxo()
 404          chain = []
 405  
 406          for _ in range(chain_length):
 407              tx = self.create_self_transfer(utxo_to_spend=chaintip_utxo)
 408              chaintip_utxo = tx["new_utxo"]
 409              chain.append(tx)
 410  
 411          return chain
 412  
 413      def send_self_transfer_chain(self, *, from_node, **kwargs):
 414          """Create and send a "chain" of chain_length transactions. The nth transaction in
 415          the chain is a child of the n-1th transaction and parent of the n+1th transaction.
 416  
 417          Returns a list of objects for each tx (see create_self_transfer_multi).
 418          """
 419          chain = self.create_self_transfer_chain(**kwargs)
 420          for t in chain:
 421              self.sendrawtransaction(from_node=from_node, tx_hex=t["hex"])
 422          return chain
 423  
 424  
 425  class NodeSigner:
 426      """Simple wallet replacement that delegates signing of existing raw transactions to a node by
 427         using the `signrawtransactionwithkey` RPC. This can be used for spending from widespread
 428         output types (P2PKH, P2WPKH, P2SH-P2WPKH, P2TR) without having the wallet compiled in."""
 429      def __init__(self, node):
 430          self._node = node
 431          self._key_entries = []
 432  
 433      def getnewaddress(self, address_type='legacy'):
 434          (seckey, pubkey), spk, address = getnewdestination(address_type)
 435          redeem_script = key_to_p2wpkh_script(pubkey) if address_type == 'p2sh-segwit' else None
 436          self._key_entries.append({"seckey_wif": bytes_to_wif(seckey.get_bytes()), "output_script": spk, "redeem_script": redeem_script})
 437          return pubkey, spk, address
 438  
 439      def listunspent(self):
 440          needles = [descsum_create(f'raw({key_entry["output_script"].hex()})') for key_entry in self._key_entries]
 441          scan_res = self._node.scantxoutset(action="start", scanobjects=needles)
 442          spend_height = scan_res['height'] + 1  # coins would be spent in the next block
 443          unspents = []
 444          for u in scan_res['unspents']:
 445              if u["coinbase"] and (spend_height - u["height"]) < COINBASE_MATURITY:  # skip immature coins
 446                  continue
 447              unspent = { "txid": u["txid"], "vout": u["vout"], "scriptPubKey": u["scriptPubKey"], "amount": u["amount"] }
 448              key_entry = [ke for ke in self._key_entries if ke["output_script"] == bytes.fromhex(u["scriptPubKey"])][0]
 449              if key_entry["redeem_script"] is not None:
 450                  unspent["redeemScript"] = key_entry["redeem_script"].hex()
 451              unspents.append(unspent)
 452          return unspents
 453  
 454      def signrawtransaction(self, tx_hex, inputs):
 455          output_scripts_to_sign = {i["scriptPubKey"] for i in inputs}
 456          seckeys_wif = [ke["seckey_wif"] for ke in self._key_entries if ke["output_script"].hex() in output_scripts_to_sign]
 457          return self._node.signrawtransactionwithkey(tx_hex, seckeys_wif, inputs)
 458  
 459  
 460  def getnewdestination(address_type='bech32m'):
 461      """Generate a random destination of the specified type and return the
 462         corresponding key pair, scriptPubKey and address. Supported types are
 463         'legacy', 'p2sh-segwit', 'bech32' and 'bech32m'. Can be used when a random
 464         destination is needed, but no compiled wallet is available (e.g. as
 465         replacement to the getnewaddress/getaddressinfo RPCs)."""
 466      key, pubkey = generate_keypair()
 467      if address_type == 'legacy':
 468          scriptpubkey = key_to_p2pkh_script(pubkey)
 469          address = key_to_p2pkh(pubkey)
 470      elif address_type == 'p2sh-segwit':
 471          scriptpubkey = key_to_p2sh_p2wpkh_script(pubkey)
 472          address = key_to_p2sh_p2wpkh(pubkey)
 473      elif address_type == 'bech32':
 474          scriptpubkey = key_to_p2wpkh_script(pubkey)
 475          address = key_to_p2wpkh(pubkey)
 476      elif address_type == 'bech32m':
 477          tap = taproot_construct(compute_xonly_pubkey(key.get_bytes())[0])
 478          pubkey = tap.output_pubkey
 479          scriptpubkey = tap.scriptPubKey
 480          address = output_key_to_p2tr(pubkey)
 481      else:
 482          assert False
 483      return (key, pubkey), scriptpubkey, address
 484