Repository: bitcoin/bitcoin
File: test/functional/test_framework/script_util.py (Python)

Code before:
    if isinstance(key, bytes) and (len(key) == 33 or len(key) == 65):
        return key
    assert False


def check_script(script):
    if isinstance(script, str):
        script = bytes.fromhex(script)  # Assuming this is hex string
    if isinstance(script, bytes) or isinstance(script, CScript):
        return script
    assert False


def valid_witness_malleate_tx(mini_wallet, node, parent_amount, child_amount):
    """
    Creates a valid witness malleation transaction:
    - Parent transaction with a script supporting 2 branches
    - 2 child transactions with the same txid but different wtxids
    """
    hashlock = hash160(b'Preimage')
    witness_script = CScript([OP_IF, OP_HASH160, hashlock, OP_EQUAL, OP_ELSE, OP_TRUE, OP_ENDIF])
    witness_program = sha256(witness_script)
    script_pubkey = CScript([OP_0, witness_program])

    # Create parent transaction with a script supporting 2 branches
    parent = mini_wallet.send_to(from_node=node, scriptPubKey=script_pubkey, amount=parent_amount, fee=1000)

    # Create 2 valid children that differ only in witness data.
    # 1. Create a new transaction with witness solving first branch
    child_witness_script = CScript([OP_TRUE])
    child_witness_program = sha256(child_witness_script)
    child_script_pubkey = CScript([OP_0, child_witness_program])
    child_one = CTransaction()

    child_one.vin.append(CTxIn(COutPoint(int(parent['txid'], 16), parent['sent_vout']), b""))
    child_one.vout.append(CTxOut(child_amount, child_script_pubkey))
    child_one.wit.vtxinwit.append(CTxInWitness())
    child_one.wit.vtxinwit[0].scriptWitness.stack = [b'Preimage', b'\x01', witness_script]
    # 2. Create another identical transaction with witness solving second branch
    child_two = deepcopy(child_one)
    child_two.wit.vtxinwit[0].scriptWitness.stack = [b'', witness_script]
    return parent, child_one, child_two


class TestFrameworkScriptUtil(unittest.TestCase):
    def test_multisig(self):
        fake_pubkey = bytes([0]*33)
        # check correct encoding of P2MS script with n,k <= 16
        normal_ms_script = keys_to_multisig_script([fake_pubkey]*16, k=15)
        self.assertEqual(len(normal_ms_script), 1 + 16*34 + 1 + 1)
        self.assertTrue(normal_ms_script.startswith(bytes([OP_15])))

Review comment (suggestion):
I feel there is scope for generalising this function further. Similar to how the previous class implementation used only the `CTransaction` class and the like. IMO, the earlier class implementation was fine as well but if the consensus is to prefer the function implementation then we should not lose the genericness.

Currently, this function is tied to the `MiniWallet` that I don't think is necessary. A more generic form of this function should be able to be used by the usual RPC wallets as well.

Consider the following diff that puts the responsibility of building the base parent transaction and submitting the final parent transaction to the caller, thereby allowing this function to not use any wallet or node at all.

<details>
<summary>Diff</summary>

```diff
diff --git a/test/functional/mempool_accept_wtxid.py b/test/functional/mempool_accept_wtxid.py
index a49a886d3a..f895772b08 100755
--- a/test/functional/mempool_accept_wtxid.py
+++ b/test/functional/mempool_accept_wtxid.py
@@ -11,7 +11,7 @@ from test_framework.messages import (
     COIN,
 )
 from test_framework.p2p import P2PTxInvStore
-from test_framework.script_util import valid_witness_malleate_tx
+from test_framework.script_util import build_malleated_tx_package
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
     assert_not_equal,
@@ -36,7 +36,9 @@ class MempoolWtxidTest(BitcoinTestFramework):
         self.log.info("Submit parent with multiple script branches to mempool")
         parent_amount = int(9.99998 * COIN)
         child_amount = int(9.99996 * COIN)
-        _, child_one, child_two = valid_witness_malleate_tx(mini_wallet, node, parent_amount, child_amount)
+        parent = mini_wallet.create_self_transfer(fee_rate=0)["tx"]
+        parent, child_one, child_two = build_malleated_tx_package(parent, parent_amount, child_amount)
+        mini_wallet.sendrawtransaction(from_node=node, tx_hex=parent.serialize().hex())
 
         self.generate(node, 1)
 
diff --git a/test/functional/test_framework/script_util.py b/test/functional/test_framework/script_util.py
index f3b3c803ed..862e30c088 100755
--- a/test/functional/test_framework/script_util.py
+++ b/test/functional/test_framework/script_util.py
@@ -6,6 +6,10 @@
 import unittest
 
 from copy import deepcopy
+from test_framework.util import (
+    assert_greater_than_or_equal,
+    assert_equal,
+)
 
 from test_framework.messages import (
     COutPoint,
@@ -150,19 +154,22 @@ def check_script(script):
     assert False
 
 
-def valid_witness_malleate_tx(mini_wallet, node, parent_amount, child_amount):
+def build_malleated_tx_package(parent: CTransaction, parent_amount, child_amount, fee=1000):
     """
-    Creates a valid witness malleation transaction:
+    Creates a malleated transaction package with valid witness:
     - Parent transaction with a script supporting 2 branches
-    - 2 child transactions with the same txid but different wtxids
+    - 2 child transactions with the same txid but different wtxids because of different witnesses
     """
     hashlock = hash160(b'Preimage')
     witness_script = CScript([OP_IF, OP_HASH160, hashlock, OP_EQUAL, OP_ELSE, OP_TRUE, OP_ENDIF])
     witness_program = sha256(witness_script)
     script_pubkey = CScript([OP_0, witness_program])
 
-    # Create parent transaction with a script supporting 2 branches
-    parent = mini_wallet.send_to(from_node=node, scriptPubKey=script_pubkey, amount=parent_amount, fee=1000)
+    # Append to the transaction the vout containing the script supporting 2 spending conditions
+    assert_equal(len(parent.vout), 1)
+    assert_greater_than_or_equal(parent.vout[0].nValue, parent_amount + fee)
+    parent.vout[0].nValue -= (parent_amount + fee)
+    parent.vout.append(CTxOut(parent_amount, script_pubkey))
 
     # Create 2 valid children that differ only in witness data.
     # 1. Create a new transaction with witness solving first branch
@@ -171,13 +178,15 @@ def valid_witness_malleate_tx(mini_wallet, node, parent_amount, child_amount):
     child_script_pubkey = CScript([OP_0, child_witness_program])
     child_one = CTransaction()
 
-    child_one.vin.append(CTxIn(COutPoint(int(parent['txid'], 16), parent['sent_vout']), b""))
+    child_one.vin.append(CTxIn(COutPoint(int(parent.txid_hex, 16), len(parent.vout) - 1), b""))
     child_one.vout.append(CTxOut(child_amount, child_script_pubkey))
     child_one.wit.vtxinwit.append(CTxInWitness())
     child_one.wit.vtxinwit[0].scriptWitness.stack = [b'Preimage', b'\x01', witness_script]
+
     # 2. Create another identical transaction with witness solving second branch
     child_two = deepcopy(child_one)
     child_two.wit.vtxinwit[0].scriptWitness.stack = [b'', witness_script]
+
     return parent, child_one, child_two
```
</details>

Code after:
        raise RuntimeError(f"target_vsize {target_vsize} is less than transaction virtual size {tx.get_vsize()}")
    # determine number of needed padding bytes
    dummy_vbytes = target_vsize - tx.get_vsize()
    # compensate for the increase of the compact-size encoded script length
    # (note that the length encoding of the unpadded output script needs one byte)
    dummy_vbytes -= len(ser_compact_size(dummy_vbytes)) - 1
    tx.vout[-1].scriptPubKey = CScript([OP_RETURN] + [OP_1] * dummy_vbytes)
    assert_equal(tx.get_vsize(), target_vsize)

def output_key_to_p2tr_script(key):
    assert len(key) == 32
    return program_to_witness_script(1, key)


def check_key(key):
    if isinstance(key, str):
        key = bytes.fromhex(key)  # Assuming this is hex string
    if isinstance(key, bytes) and (len(key) == 33 or len(key) == 65):
        return key
    assert False


def check_script(script):
    if isinstance(script, str):
        script = bytes.fromhex(script)  # Assuming this is hex string
    if isinstance(script, bytes) or isinstance(script, CScript):
        return script
    assert False


def build_malleated_tx_package(*, parent: CTransaction, rebalance_parent_output_amount, child_amount):
    """
    Returns a transaction package with valid witness:
    - Parent transaction whose last output contains a script that has two spending conditions
    - Two malleated child transactions with same txid but different wtxids because of different witnesses

    Args:
        parent: Transaction with modifiable outputs. Either unsigned (sign after
        calling this function) or anyone-can-spend (e.g., MiniWallet's OP_TRUE).
    """
    hashlock = hash160(b'Preimage')
    witness_script = CScript([OP_IF, OP_HASH160, hashlock, OP_EQUAL, OP_ELSE, OP_TRUE, OP_ENDIF])
    witness_program = sha256(witness_script)
    script_pubkey = CScript([OP_0, witness_program])

    # Append to the transaction the vout containing the script supporting 2 spending conditions
    assert_greater_than_or_equal(len(parent.vout), 1)
    last_output = parent.vout[len(parent.vout) - 1]
    assert_greater_than_or_equal(last_output.nValue, rebalance_parent_output_amount)
    last_output.nValue -= rebalance_parent_output_amount
    parent.vout.append(CTxOut(rebalance_parent_output_amount, script_pubkey))