wallet.py raw

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