Skip to content

Commit

Permalink
test: refactor: introduce generate_keypair helper with WIF support
Browse files Browse the repository at this point in the history
In functional tests it is a quite common scenario to generate fresh
elliptic curve keypairs, which is currently a bit cumbersome as it
involves multiple steps, e.g.:

    privkey = ECKey()
    privkey.generate()
    privkey_wif = bytes_to_wif(privkey.get_bytes())
    pubkey = privkey.get_pubkey().get_bytes()

Simplify this by providing a new `generate_keypair` helper function that
returns the private key either as `ECKey` object or as WIF-string
(depending on the boolean `wif` parameter) and the public key as
byte-string; these formats are what we mostly need (currently we don't
use `ECPubKey` objects from generated keypairs anywhere).

With this, most of the affected code blocks following the pattern above
can be replaced by one-liners, e.g.:

    privkey, pubkey = generate_keypair(wif=True)

Note that after this commit, the only direct uses of `ECKey` remain in
situations where we want to set the private key explicitly, e.g. in
MiniWallet (test/functional/test_framework/wallet.py) or the test for
the signet miner script (test/functional/tool_signet_miner.py).
  • Loading branch information
theStack committed Jun 19, 2023
1 parent 7f0b79e commit 1a572ce
Show file tree
Hide file tree
Showing 18 changed files with 79 additions and 130 deletions.
11 changes: 6 additions & 5 deletions test/functional/feature_assumevalid.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@
create_block,
create_coinbase,
)
from test_framework.key import ECKey
from test_framework.messages import (
CBlockHeader,
COutPoint,
Expand All @@ -46,9 +45,13 @@
msg_headers,
)
from test_framework.p2p import P2PInterface
from test_framework.script import (CScript, OP_TRUE)
from test_framework.script import (
CScript,
OP_TRUE,
)
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import assert_equal
from test_framework.wallet_util import generate_keypair


class BaseNode(P2PInterface):
Expand Down Expand Up @@ -90,9 +93,7 @@ def run_test(self):
self.blocks = []

# Get a pubkey for the coinbase TXO
coinbase_key = ECKey()
coinbase_key.generate()
coinbase_pubkey = coinbase_key.get_pubkey().get_bytes()
_, coinbase_pubkey = generate_keypair()

# Create the first block with a coinbase output to our key
height = 1
Expand Down
6 changes: 2 additions & 4 deletions test/functional/feature_block.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
get_legacy_sigopcount_block,
MAX_BLOCK_SIGOPS,
)
from test_framework.key import ECKey
from test_framework.messages import (
CBlock,
COIN,
Expand Down Expand Up @@ -55,6 +54,7 @@
assert_equal,
assert_greater_than,
)
from test_framework.wallet_util import generate_keypair
from data import invalid_txs


Expand Down Expand Up @@ -98,9 +98,7 @@ def run_test(self):
self.bootstrap_p2p() # Add one p2p connection to the node

self.block_heights = {}
self.coinbase_key = ECKey()
self.coinbase_key.generate()
self.coinbase_pubkey = self.coinbase_key.get_pubkey().get_bytes()
self.coinbase_key, self.coinbase_pubkey = generate_keypair()
self.tip = None
self.blocks = {}
self.genesis_hash = int(self.nodes[0].getbestblockhash(), 16)
Expand Down
12 changes: 4 additions & 8 deletions test/functional/feature_nulldummy.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,7 @@
assert_raises_rpc_error,
)
from test_framework.wallet import getnewdestination
from test_framework.key import ECKey
from test_framework.wallet_util import bytes_to_wif
from test_framework.wallet_util import generate_keypair

NULLDUMMY_ERROR = "non-mandatory-script-verify-flag (Dummy CHECKMULTISIG argument must be zero)"

Expand Down Expand Up @@ -71,12 +70,9 @@ def create_transaction(self, *, txid, input_details=None, addr, amount, privkey)
return tx_from_hex(signedtx["hex"])

def run_test(self):
eckey = ECKey()
eckey.generate()
self.privkey = bytes_to_wif(eckey.get_bytes())
self.pubkey = eckey.get_pubkey().get_bytes().hex()
cms = self.nodes[0].createmultisig(1, [self.pubkey])
wms = self.nodes[0].createmultisig(1, [self.pubkey], 'p2sh-segwit')
self.privkey, self.pubkey = generate_keypair(wif=True)
cms = self.nodes[0].createmultisig(1, [self.pubkey.hex()])
wms = self.nodes[0].createmultisig(1, [self.pubkey.hex()], 'p2sh-segwit')
self.ms_address = cms["address"]
ms_unlock_details = {"scriptPubKey": address_to_scriptpubkey(self.ms_address).hex(),
"redeemScript": cms["redeemScript"]}
Expand Down
8 changes: 3 additions & 5 deletions test/functional/feature_taproot.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@
assert_equal,
random_bytes,
)
from test_framework.wallet_util import generate_keypair
from test_framework.key import (
generate_privkey,
compute_xonly_pubkey,
Expand Down Expand Up @@ -1186,11 +1187,8 @@ def predict_sigops_ratio(n, dummy_size):

# Also add a few legacy spends into the mix, so that transactions which combine taproot and pre-taproot spends get tested too.
for compressed in [False, True]:
eckey1 = ECKey()
eckey1.set(generate_privkey(), compressed)
pubkey1 = eckey1.get_pubkey().get_bytes()
eckey2 = ECKey()
eckey2.set(generate_privkey(), compressed)
eckey1, pubkey1 = generate_keypair(compressed=compressed)
eckey2, _ = generate_keypair(compressed=compressed)
for p2sh in [False, True]:
for witv0 in [False, True]:
for hashtype in VALID_SIGHASHES_ECDSA + [random.randrange(0x04, 0x80), random.randrange(0x84, 0x100)]:
Expand Down
6 changes: 2 additions & 4 deletions test/functional/mempool_accept.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import math

from test_framework.test_framework import BitcoinTestFramework
from test_framework.key import ECKey
from test_framework.messages import (
MAX_BIP125_RBF_SEQUENCE,
COIN,
Expand Down Expand Up @@ -44,6 +43,7 @@
assert_raises_rpc_error,
)
from test_framework.wallet import MiniWallet
from test_framework.wallet_util import generate_keypair


class MempoolAcceptanceTest(BitcoinTestFramework):
Expand Down Expand Up @@ -283,9 +283,7 @@ def run_test(self):
rawtxs=[tx.serialize().hex()],
)
tx = tx_from_hex(raw_tx_reference)
key = ECKey()
key.generate()
pubkey = key.get_pubkey().get_bytes()
_, pubkey = generate_keypair()
tx.vout[0].scriptPubKey = keys_to_multisig_script([pubkey] * 3, k=2) # Some bare multisig script (2-of-3)
self.check_mempool_result(
result_expected=[{'txid': tx.rehash(), 'allowed': False, 'reject-reason': 'bare-multisig'}],
Expand Down
9 changes: 3 additions & 6 deletions test/functional/mempool_dust.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
"""Test dust limit mempool policy (`-dustrelayfee` parameter)"""
from decimal import Decimal

from test_framework.key import ECKey
from test_framework.messages import (
COIN,
CTxOut,
Expand All @@ -32,6 +31,7 @@
get_fee,
)
from test_framework.wallet import MiniWallet
from test_framework.wallet_util import generate_keypair


DUST_RELAY_TX_FEE = 3000 # default setting [sat/kvB]
Expand Down Expand Up @@ -74,11 +74,8 @@ def run_test(self):
self.wallet = MiniWallet(self.nodes[0])

# prepare output scripts of each standard type
key = ECKey()
key.generate(compressed=False)
uncompressed_pubkey = key.get_pubkey().get_bytes()
key.generate(compressed=True)
pubkey = key.get_pubkey().get_bytes()
_, uncompressed_pubkey = generate_keypair(compressed=False)
_, pubkey = generate_keypair(compressed=True)

output_scripts = (
(key_to_p2pk_script(uncompressed_pubkey), "P2PK (uncompressed)"),
Expand Down
12 changes: 3 additions & 9 deletions test/functional/p2p_segwit.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
create_block,
create_coinbase,
)
from test_framework.key import ECKey
from test_framework.messages import (
MAX_BIP125_RBF_SEQUENCE,
CBlockHeader,
Expand Down Expand Up @@ -89,6 +88,7 @@
assert_raises_rpc_error,
)
from test_framework.wallet import MiniWallet
from test_framework.wallet_util import generate_keypair


MAX_SIGOP_COST = 80000
Expand Down Expand Up @@ -1448,9 +1448,7 @@ def test_uncompressed_pubkey(self):

# Segwit transactions using uncompressed pubkeys are not accepted
# under default policy, but should still pass consensus.
key = ECKey()
key.generate(False)
pubkey = key.get_pubkey().get_bytes()
key, pubkey = generate_keypair(compressed=False)
assert_equal(len(pubkey), 65) # This should be an uncompressed pubkey

utxo = self.utxo.pop(0)
Expand Down Expand Up @@ -1544,11 +1542,7 @@ def test_uncompressed_pubkey(self):

@subtest
def test_signature_version_1(self):

key = ECKey()
key.generate()
pubkey = key.get_pubkey().get_bytes()

key, pubkey = generate_keypair()
witness_script = key_to_p2pk_script(pubkey)
script_pubkey = script_to_p2wsh_script(witness_script)

Expand Down
11 changes: 5 additions & 6 deletions test/functional/rpc_createmultisig.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,13 @@
from test_framework.blocktools import COINBASE_MATURITY
from test_framework.authproxy import JSONRPCException
from test_framework.descriptors import descsum_create, drop_origins
from test_framework.key import ECPubKey, ECKey
from test_framework.key import ECPubKey
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import (
assert_raises_rpc_error,
assert_equal,
)
from test_framework.wallet_util import bytes_to_wif
from test_framework.wallet_util import generate_keypair
from test_framework.wallet import (
MiniWallet,
getnewdestination,
Expand All @@ -38,10 +38,9 @@ def get_keys(self):
self.priv = []
node0, node1, node2 = self.nodes
for _ in range(self.nkeys):
k = ECKey()
k.generate()
self.pub.append(k.get_pubkey().get_bytes().hex())
self.priv.append(bytes_to_wif(k.get_bytes(), k.is_compressed))
privkey, pubkey = generate_keypair(wif=True)
self.pub.append(pubkey.hex())
self.priv.append(privkey)
if self.is_bdb_compiled():
self.final = node2.getnewaddress()
else:
Expand Down
22 changes: 8 additions & 14 deletions test/functional/rpc_psbt.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from itertools import product

from test_framework.descriptors import descsum_create
from test_framework.key import ECKey, H_POINT
from test_framework.key import H_POINT
from test_framework.messages import (
COutPoint,
CTransaction,
Expand Down Expand Up @@ -43,8 +43,8 @@
random_bytes,
)
from test_framework.wallet_util import (
bytes_to_wif,
get_generate_key
generate_keypair,
get_generate_key,
)

import json
Expand Down Expand Up @@ -710,9 +710,7 @@ def test_psbt_input_keys(psbt_input, keys):

self.log.info("Test that we can fund psbts with external inputs specified")

eckey = ECKey()
eckey.generate()
privkey = bytes_to_wif(eckey.get_bytes())
privkey, _ = generate_keypair(wif=True)

self.nodes[1].createwallet("extfund")
wallet = self.nodes[1].get_wallet_rpc("extfund")
Expand Down Expand Up @@ -825,11 +823,9 @@ def test_psbt_input_keys(psbt_input, keys):
self.nodes[1].createwallet(wallet_name="scriptwatchonly", disable_private_keys=True)
watchonly = self.nodes[1].get_wallet_rpc("scriptwatchonly")

eckey = ECKey()
eckey.generate()
privkey = bytes_to_wif(eckey.get_bytes())
privkey, pubkey = generate_keypair(wif=True)

desc = descsum_create("wsh(pkh({}))".format(eckey.get_pubkey().get_bytes().hex()))
desc = descsum_create("wsh(pkh({}))".format(pubkey.hex()))
if self.options.descriptors:
res = watchonly.importdescriptors([{"desc": desc, "timestamp": "now"}])
else:
Expand All @@ -846,11 +842,9 @@ def test_psbt_input_keys(psbt_input, keys):

# Same test but for taproot
if self.options.descriptors:
eckey = ECKey()
eckey.generate()
privkey = bytes_to_wif(eckey.get_bytes())
privkey, pubkey = generate_keypair(wif=True)

desc = descsum_create("tr({},pk({}))".format(H_POINT, eckey.get_pubkey().get_bytes().hex()))
desc = descsum_create("tr({},pk({}))".format(H_POINT, pubkey.hex()))
res = watchonly.importdescriptors([{"desc": desc, "timestamp": "now"}])
assert res[0]["success"]
addr = self.nodes[0].deriveaddresses(desc)[0]
Expand Down
21 changes: 7 additions & 14 deletions test/functional/rpc_signrawtransactionwithkey.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
address_to_scriptpubkey,
script_to_p2sh,
)
from test_framework.key import ECKey
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import (
assert_equal,
Expand All @@ -23,16 +22,16 @@
script_to_p2sh_p2wsh_script,
script_to_p2wsh_script,
)
from test_framework.wallet import (
getnewdestination,
)
from test_framework.wallet_util import (
bytes_to_wif,
generate_keypair,
)

from decimal import (
Decimal,
)
from test_framework.wallet import (
getnewdestination,
)


class SignRawTransactionWithKeyTest(BitcoinTestFramework):
Expand Down Expand Up @@ -80,11 +79,8 @@ def successful_signing_test(self):
def witness_script_test(self):
self.log.info("Test signing transaction to P2SH-P2WSH addresses without wallet")
# Create a new P2SH-P2WSH 1-of-1 multisig address:
eckey = ECKey()
eckey.generate()
embedded_privkey = bytes_to_wif(eckey.get_bytes())
embedded_pubkey = eckey.get_pubkey().get_bytes().hex()
p2sh_p2wsh_address = self.nodes[1].createmultisig(1, [embedded_pubkey], "p2sh-segwit")
embedded_privkey, embedded_pubkey = generate_keypair(wif=True)
p2sh_p2wsh_address = self.nodes[1].createmultisig(1, [embedded_pubkey.hex()], "p2sh-segwit")
# send transaction to P2SH-P2WSH 1-of-1 multisig address
self.block_hash = self.generate(self.nodes[0], COINBASE_MATURITY + 1)
self.blk_idx = 0
Expand All @@ -109,10 +105,7 @@ def witness_script_test(self):

def verify_txn_with_witness_script(self, tx_type):
self.log.info("Test with a {} script as the witnessScript".format(tx_type))
eckey = ECKey()
eckey.generate()
embedded_privkey = bytes_to_wif(eckey.get_bytes())
embedded_pubkey = eckey.get_pubkey().get_bytes().hex()
embedded_privkey, embedded_pubkey = generate_keypair(wif=True)
witness_script = {
'P2PKH': key_to_p2pkh_script(embedded_pubkey).hex(),
'P2PK': key_to_p2pk_script(embedded_pubkey).hex()
Expand Down
7 changes: 3 additions & 4 deletions test/functional/test_framework/wallet.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
key_to_p2wpkh,
output_key_to_p2tr,
)
from test_framework.blocktools import COINBASE_MATURITY
from test_framework.descriptors import descsum_create
from test_framework.key import (
ECKey,
Expand Down Expand Up @@ -53,7 +54,7 @@
assert_equal,
assert_greater_than_or_equal,
)
from test_framework.blocktools import COINBASE_MATURITY
from test_framework.wallet_util import generate_keypair

DEFAULT_FEE = Decimal("0.0001")

Expand Down Expand Up @@ -395,9 +396,7 @@ def getnewdestination(address_type='bech32m'):
'legacy', 'p2sh-segwit', 'bech32' and 'bech32m'. Can be used when a random
destination is needed, but no compiled wallet is available (e.g. as
replacement to the getnewaddress/getaddressinfo RPCs)."""
key = ECKey()
key.generate()
pubkey = key.get_pubkey().get_bytes()
key, pubkey = generate_keypair()
if address_type == 'legacy':
scriptpubkey = key_to_p2pkh_script(pubkey)
address = key_to_p2pkh(pubkey)
Expand Down
Loading

0 comments on commit 1a572ce

Please sign in to comment.