test: MiniWallet: add P2PK support

Author: theStack

test/functional/test_framework/wallet.py

@@ -6,6 +6,7 @@
 
 from decimal import Decimal
 from test_framework.address import ADDRESS_BCRT1_P2WSH_OP_TRUE
+from test_framework.key import ECKey
 from test_framework.messages import (
     COIN,
     COutPoint,
@@ -16,8 +17,11 @@
 )
 from test_framework.script import (
     CScript,
+    LegacySignatureHash,
+    OP_CHECKSIG,
     OP_TRUE,
     OP_NOP,
+    SIGHASH_ALL,
 )
 from test_framework.util import (
     assert_equal,
@@ -27,12 +31,20 @@
 
 
 class MiniWallet:
-    def __init__(self, test_node, *, raw_script=False):
+    def __init__(self, test_node, *, raw_script=False, use_p2pk=False):
         self._test_node = test_node
         self._utxos = []
+        self._priv_key = None
+        self._address = None
+
         if raw_script:
-            self._address = None
             self._scriptPubKey = bytes(CScript([OP_TRUE]))
+        elif use_p2pk:
+            # use simple deterministic private key (k=1)
+            self._priv_key = ECKey()
+            self._priv_key.set((1).to_bytes(32, 'big'), True)
+            pub_key = self._priv_key.get_pubkey()
+            self._scriptPubKey = bytes(CScript([pub_key.get_bytes(), OP_CHECKSIG]))
         else:
             self._address = ADDRESS_BCRT1_P2WSH_OP_TRUE
             self._scriptPubKey = hex_str_to_bytes(self._test_node.validateaddress(self._address)['scriptPubKey'])
@@ -50,6 +62,13 @@ def scan_tx(self, tx):
             if out['scriptPubKey']['hex'] == self._scriptPubKey.hex():
                 self._utxos.append({'txid': tx['txid'], 'vout': out['n'], 'value': out['value']})
 
+    def sign_tx(self, tx):
+        """Sign tx that has been created by MiniWallet in P2PK mode"""
+        assert self._priv_key is not None
+        (sighash, err) = LegacySignatureHash(CScript(self._scriptPubKey), tx, 0, SIGHASH_ALL)
+        assert err is None
+        tx.vin[0].scriptSig = CScript([self._priv_key.sign_ecdsa(sighash) + bytes(bytearray([SIGHASH_ALL]))])
+
     def generate(self, num_blocks):
         """Generate blocks with coinbase outputs to the internal address, and append the outputs to the internal list"""
         blocks = self._test_node.generatetodescriptor(num_blocks, f'raw({self._scriptPubKey.hex()})')
@@ -99,7 +118,12 @@ def create_self_transfer(self, *, fee_rate=Decimal("0.003"), from_node, utxo_to_
         tx.vout = [CTxOut(int(send_value * COIN), self._scriptPubKey)]
         if not self._address:
             # raw script
-            tx.vin[0].scriptSig = CScript([OP_NOP] * 35)  # pad to identical size
+            if self._priv_key is not None:
+                # P2PK, need to sign
+                self.sign_tx(tx)
+            else:
+                # anyone-can-spend
+                tx.vin[0].scriptSig = CScript([OP_NOP] * 35)  # pad to identical size
         else:
             tx.wit.vtxinwit = [CTxInWitness()]
             tx.wit.vtxinwit[0].scriptWitness.stack = [CScript([OP_TRUE])]
@@ -108,7 +132,10 @@ def create_self_transfer(self, *, fee_rate=Decimal("0.003"), from_node, utxo_to_
         tx_info = from_node.testmempoolaccept([tx_hex])[0]
         assert_equal(mempool_valid, tx_info['allowed'])
         if mempool_valid:
-            assert_equal(tx_info['vsize'], vsize)
+            # TODO: for P2PK, vsize is not constant due to varying scriptSig length,
+            # so only check this for anyone-can-spend outputs right now
+            if self._priv_key is None:
+                assert_equal(tx_info['vsize'], vsize)
             assert_equal(tx_info['fees']['base'], fee)
         return {'txid': tx_info['txid'], 'wtxid': tx_info['wtxid'], 'hex': tx_hex, 'tx': tx}