Generalize txmaker to most supported policies

test_utils/taproot.py

@@ -0,0 +1,50 @@
+# from portions of BIP-0341
+# - https://github.com/bitcoin/bips/blob/b3701faef2bdb98a0d7ace4eedbeefa2da4c89ed/bip-0341.mediawiki
+# Distributed under the BSD-3-Clause license
+
+# fmt: off
+
+# If you want to print values on an individual basis, use
+# the pretty() function, e.g., print(pretty(foo)).
+import hashlib
+import struct
+
+
+# This implementation can be sped up by storing the midstate after hashing
+# tag_hash instead of rehashing it all the time.
+def tagged_hash(tag: str, msg: bytes) -> bytes:
+ tag_hash = hashlib.sha256(tag.encode()).digest()
+ return hashlib.sha256(tag_hash + tag_hash + msg).digest()
+
+
+def ser_compact_size(l):
+ r = b""
+ if l < 253:
+ r = struct.pack("B", l)
+ elif l < 0x10000:
+ r = struct.pack("<BH", 253, l)
+ elif l < 0x100000000:
+ r = struct.pack("<BI", 254, l)
+ else:
+ r = struct.pack("<BQ", 255, l)
+ return r
+
+def deser_compact_size(f):
+ nit = struct.unpack("<B", f.read(1))[0]
+ if nit == 253:
+ nit = struct.unpack("<H", f.read(2))[0]
+ elif nit == 254:
+ nit = struct.unpack("<I", f.read(4))[0]
+ elif nit == 255:
+ nit = struct.unpack("<Q", f.read(8))[0]
+ return nit
+
+def deser_string(f):
+ nit = deser_compact_size(f)
+ return f.read(nit)
+
+def ser_string(s):
+ return ser_compact_size(len(s)) + s
+
+def ser_script(s):
+ return ser_string(s)

test_utils/txmaker.py

@@ -1,8 +1,16 @@
+# This module contains a utility function to create test PSBTs spending from an arbitrary wallet policy.
+# It creates transactions spending non-existing UTXOs, and fills in the PSBTs with enough information to
+# satisfy the requirements of the Ledger bitcoin app.
+# It does not guarantee BIP-174 compliant PSBTs, as some fields that are not required in the
+# Ledger bitcoin app might not be filled in.
+
+
+from io import BytesIO
 from random import randint
 
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Union
 from bitcoin_client.ledger_bitcoin import WalletPolicy, WalletType
-from bitcoin_client.ledger_bitcoin.key import KeyOriginInfo, parse_path, get_taproot_output_key
+from bitcoin_client.ledger_bitcoin.key import ExtendedKey, KeyOriginInfo, parse_path, get_taproot_output_key
 from bitcoin_client.ledger_bitcoin.psbt import PSBT, PartiallySignedInput, PartiallySignedOutput
 from bitcoin_client.ledger_bitcoin.tx import CScriptWitness, CTransaction, CTxIn, CTxInWitness, CTxOut, COutPoint, CTxWitness, uint256_from_str
 
@@ -11,6 +19,10 @@
 from embit.bip32 import HDKey
 from embit.bip39 import mnemonic_to_seed
 
+from ledger_bitcoin.embit.descriptor.miniscript import Miniscript
+from test_utils import bip0340
+from test_utils.wallet_policy import DescriptorTemplate, KeyPlaceholder, PlainKeyPlaceholder, TrDescriptorTemplate, WshDescriptorTemplate, derive_plain_descriptor, tapleaf_hash
+
 
 SPECULOS_SEED = "glory promote mansion idle axis finger extra february uncover one trip resource lawn turtle enact monster seven myth punch hobby comfort wild raise skin"
 master_key = HDKey.from_seed(mnemonic_to_seed(SPECULOS_SEED))
@@ -40,6 +52,14 @@ def random_txid() -> bytes:
  return random_bytes(32)
 
 
+def random_p2tr() -> bytes:
+ """Returns 32 random bytes. Not cryptographically secure."""
+ privkey = random_bytes(32)
+ pubkey = bip0340.point_mul(bip0340.G, int.from_bytes(privkey, 'big'))
+
+ return b'\x51\x20' + (pubkey[0]).to_bytes(32, 'big')
+
+
 def getScriptPubkeyFromWallet(wallet: WalletPolicy, change: bool, address_index: int) -> Script:
  descriptor_str = wallet.descriptor_template
 
@@ -116,27 +136,109 @@ def createFakeWalletTransaction(n_inputs: int, n_outputs: int, output_amount: in
  return tx, selected_output_index, selected_output_change, selected_output_address_index
 
 
-def createPsbt(wallet: WalletPolicy, input_amounts: List[int], output_amounts: List[int], output_is_change: List[bool], output_wallet: Optional[List[Optional[WalletPolicy]]] = None) -> PSBT:
- if output_wallet is None:
- output_wallet = [None] * len(output_amounts)
+def get_placeholder_root_key(placeholder: KeyPlaceholder, keys_info: List[str]) -> Tuple[ExtendedKey, Optional[KeyOriginInfo]]:
+ if isinstance(placeholder, PlainKeyPlaceholder):
+ key_info = keys_info[placeholder.key_index]
+ key_origin_end_pos = key_info.find("]")
+ if key_origin_end_pos == -1:
+ xpub = key_info
+ root_key_origin = None
+ else:
+ xpub = key_info[key_origin_end_pos+1:]
+ root_key_origin = KeyOriginInfo.from_string(
+ key_info[1:key_origin_end_pos])
+ root_pubkey = ExtendedKey.deserialize(xpub)
+ else:
+ raise ValueError("Unsupported placeholder type")
+
+ return root_pubkey, root_key_origin
 
- assert len(output_amounts) == len(output_is_change)
- assert len(output_amounts) == len(output_wallet)
- assert sum(output_amounts) <= sum(input_amounts)
 
- # TODO: add support for wrapped segwit wallets
+def fill_inout(wallet_policy: WalletPolicy, inout: Union[PartiallySignedInput, PartiallySignedOutput], is_change: bool, address_index: int):
+ desc_tmpl = DescriptorTemplate.from_string(
+ wallet_policy.descriptor_template)
 
- if wallet.n_keys != 1:
- raise NotImplementedError("Only 1-key wallets supported")
- if wallet.version == WalletType.WALLET_POLICY_V1:
- if wallet.descriptor_template not in ["pkh(@0)", "wpkh(@0)", "tr(@0)"]:
- raise NotImplementedError("Unsupported policy type")
- elif wallet.version == WalletType.WALLET_POLICY_V2:
- if wallet.descriptor_template not in ["pkh(@0/**)", "wpkh(@0/**)", "tr(@0/**)"]:
- raise NotImplementedError("Unsupported policy type")
+ if isinstance(desc_tmpl, TrDescriptorTemplate):
+ keypath_der_subpath = [
+ desc_tmpl.key.num1 if not is_change else desc_tmpl.key.num2,
+ address_index
+ ]
+
+ keypath_pubkey, _ = get_placeholder_root_key(
+ desc_tmpl.key, wallet_policy.keys_info)
+
+ inout.tap_internal_key = keypath_pubkey.derive_pub_path(
+ keypath_der_subpath).pubkey[1:]
+
+ if desc_tmpl.tree is not None:
+ inout.tap_merkle_root = desc_tmpl.get_taptree_hash(
+ wallet_policy.keys_info, is_change, address_index)
+
+ for placeholder, tapleaf_desc in desc_tmpl.placeholders():
+ root_pubkey, root_pubkey_origin = get_placeholder_root_key(
+ placeholder, wallet_policy.keys_info)
+
+ placeholder_der_subpath = [
+ placeholder.num1 if not is_change else placeholder.num2,
+ address_index
+ ]
+
+ leaf_script = None
+ if tapleaf_desc is not None:
+ leaf_desc = derive_plain_descriptor(
+ tapleaf_desc, wallet_policy.keys_info, is_change, address_index)
+ s = BytesIO(leaf_desc.encode())
+ desc: Miniscript = Miniscript.read_from(s, taproot=True)
+ leaf_script = desc.compile()
+
+ derived_pubkey = root_pubkey.derive_pub_path(
+ placeholder_der_subpath)
+
+ if root_pubkey_origin is not None:
+ derived_key_origin = KeyOriginInfo(
+ root_pubkey_origin.fingerprint, root_pubkey_origin.path + placeholder_der_subpath)
+
+ leaf_hashes = []
+ if leaf_script is not None:
+ # In BIP-388 compliant wallet policies, there will be only one tapleaf with a given key
+ leaf_hashes = [tapleaf_hash(leaf_script)]
+
+ inout.tap_bip32_paths[derived_pubkey.pubkey[1:]] = (
+ leaf_hashes, derived_key_origin)
  else:
- raise ValueError(
- f"Unknown wallet policy version: {wallet.version}")
+ if isinstance(desc_tmpl, WshDescriptorTemplate):
+ # add witnessScript
+ desc_str = derive_plain_descriptor(
+ wallet_policy.descriptor_template, wallet_policy.keys_info, is_change, address_index)
+ s = BytesIO(desc_str.encode())
+ desc: Descriptor = Descriptor.read_from(s)
+ inout.witness_script = desc.witness_script().data
+
+ for placeholder, _ in desc_tmpl.placeholders():
+ root_pubkey, root_pubkey_origin = get_placeholder_root_key(
+ placeholder, wallet_policy.keys_info)
+
+ placeholder_der_subpath = [
+ placeholder.num1 if not is_change else placeholder.num2,
+ address_index
+ ]
+
+ derived_pubkey = root_pubkey.derive_pub_path(
+ placeholder_der_subpath)
+
+ if root_pubkey_origin is not None:
+ derived_key_origin = KeyOriginInfo(
+ root_pubkey_origin.fingerprint, root_pubkey_origin.path + placeholder_der_subpath)
+
+ inout.hd_keypaths[derived_pubkey.pubkey] = derived_key_origin
+
+
+def createPsbt(wallet_policy: WalletPolicy, input_amounts: List[int], output_amounts: List[int], output_is_change: List[bool]) -> PSBT:
+ assert output_is_change.count(
+ True) <= 1, "At most one change output is supported"
+
+ assert len(output_amounts) == len(output_is_change)
+ assert sum(output_amounts) <= sum(input_amounts)
 
  vin: List[CTxIn] = [CTxIn() for _ in input_amounts]
  vout: List[CTxOut] = [CTxOut() for _ in output_amounts]
@@ -150,7 +252,7 @@ def createPsbt(wallet: WalletPolicy, input_amounts: List[int], output_amounts: L
  n_inputs = randint(1, 10)
  n_outputs = randint(1, 10)
  prevout, idx, is_change, addr_idx = createFakeWalletTransaction(
- n_inputs, n_outputs, prevout_amount, wallet)
+ n_inputs, n_outputs, prevout_amount, wallet_policy)
  prevouts.append(prevout)
  prevout_ns.append(idx)
  prevout_path_change.append(is_change)
@@ -168,68 +270,43 @@ def createPsbt(wallet: WalletPolicy, input_amounts: List[int], output_amounts: L
  tx.vout = vout
  tx.wit = CTxWitness()
 
+ change_address_index = randint(0, 10_000)
+ for i, output_amount in enumerate(output_amounts):
+ tx.vout[i].nValue = output_amount
+ if output_is_change[i]:
+ script = getScriptPubkeyFromWallet(
+ wallet_policy, output_is_change[i], change_address_index)
+
+ tx.vout[i].scriptPubKey = script.data
+ else:
+ # a random P2TR output
+ tx.vout[i].scriptPubKey = random_p2tr()
+
  psbt.inputs = [PartiallySignedInput(0) for _ in input_amounts]
  psbt.outputs = [PartiallySignedOutput(0) for _ in output_amounts]
 
- # simplification; good enough for the scripts we support now, but will need more work
- is_legacy = wallet.descriptor_template.startswith("pkh(")
- is_segwitv0 = wallet.descriptor_template.startswith(
- "wpkh(") or wallet.descriptor_template.startswith("sh(wpkh(")
- is_taproot = wallet.descriptor_template.startswith("tr(")
-
- key_origin = wallet.keys_info[0][1:wallet.keys_info[0].index("]")]
+ desc_tmpl = DescriptorTemplate.from_string(
+ wallet_policy.descriptor_template)
 
- for i in range(len(input_amounts)):
- if is_legacy or is_segwitv0:
- # add non-witness UTXO
- psbt.inputs[i].non_witness_utxo = prevouts[i]
- if is_segwitv0 or is_taproot:
+ for input_index, input in enumerate(psbt.inputs):
+ if desc_tmpl.is_segwit():
  # add witness UTXO
- psbt.inputs[i].witness_utxo = prevouts[i].vout[prevout_ns[i]]
-
- path_str = f"m{key_origin[8:]}/{prevout_path_change[i]}/{prevout_path_addr_idx[i]}"
- path = parse_path(path_str)
- input_key: bytes = master_key.derive(path_str).key.sec()
-
- assert len(input_key) == 33
-
- # add key and path info
- if is_legacy or is_segwitv0:
- psbt.inputs[i].hd_keypaths[input_key] = KeyOriginInfo(
- master_key_fpr, path)
- elif is_taproot:
- internal_key = input_key[1:]
- psbt.inputs[i].tap_bip32_paths[internal_key] = (
- {}, KeyOriginInfo(master_key_fpr, path))
- else:
- raise RuntimeError("Unexpected state: unknown transaction type")
+ input.witness_utxo = prevouts[input_index].vout[prevout_ns[input_index]]
 
- for i, output_amount in enumerate(output_amounts):
- wallet_i = output_wallet[i]
- if output_is_change[i] or wallet_i is None:
- script = getScriptPubkeyFromWallet(wallet, output_is_change[i], i)
- else:
- script = getScriptPubkeyFromWallet(wallet_i, 0, i)
+ if desc_tmpl.is_legacy() or (desc_tmpl.is_segwit() and not desc_tmpl.is_taproot()):
+ # add non_witness_utxo for legacy or segwitv0
+ input.non_witness_utxo = prevouts[input_index]
 
- tx.vout[i].scriptPubKey = script.data
- tx.vout[i].nValue = output_amount
+ is_change = bool(prevout_path_change[input_index])
+ address_index = prevout_path_addr_idx[input_index]
 
- if output_is_change[i]:
- path_str = f"m{key_origin[8:]}/1/{i}"
- path = parse_path(path_str)
- output_key: bytes = master_key.derive(path_str).key.sec()
-
- # add key and path information for change output
- if is_legacy or is_segwitv0:
- psbt.outputs[i].hd_keypaths[output_key] = KeyOriginInfo(
- master_key_fpr, path)
- elif is_taproot:
- internal_key = output_key[1:]
- psbt.outputs[i].tap_bip32_paths[internal_key] = (
- {}, KeyOriginInfo(master_key_fpr, path))
-
- psbt.outputs[i].tap_bip32_paths[internal_key] = (
- {}, KeyOriginInfo(master_key_fpr, path))
+ fill_inout(wallet_policy, input, is_change, address_index)
+
+ # only for the change output, we need to do the same
+ for output_index, output in enumerate(psbt.outputs):
+ if output_is_change[output_index]:
+ fill_inout(wallet_policy, output, is_change=True,
+ address_index=change_address_index)
 
  psbt.tx = tx