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