bitcoin.py: SCRIPT-related clean-up. transaction.py: construct_witness

lib/bitcoin.py

@@ -152,7 +152,7 @@ def int_to_hex(i, length=1):
     s = "0"*(2*length - len(s)) + s
     return rev_hex(s)
 
-def script_num_to_hex(i):
+def script_num_to_hex(i: int) -> str:
     """See CScriptNum in Bitcoin Core.
     Encodes an integer as hex, to be used in script.
 
@@ -176,7 +176,7 @@ def script_num_to_hex(i):
     return bh2u(result)
 
 
-def var_int(i):
+def var_int(i: int) -> str:
     # https://en.bitcoin.it/wiki/Protocol_specification#Variable_length_integer
     if i<0xfd:
         return int_to_hex(i)
@@ -188,14 +188,14 @@ def var_int(i):
         return "ff"+int_to_hex(i,8)
 
 
-def witness_push(item):
-    """ Returns data in the form it should be present in the witness.
+def witness_push(item: str) -> str:
+    """Returns data in the form it should be present in the witness.
     hex -> hex
     """
     return var_int(len(item) // 2) + item
 
 
-def op_push(i):
+def op_push(i: int) -> str:
     if i<0x4c:  # OP_PUSHDATA1
         return int_to_hex(i)
     elif i<=0xff:
@@ -206,36 +206,32 @@ def op_push(i):
         return '4e' + int_to_hex(i,4)
 
 
-def add_data_to_script(data):
+def push_script(data: str) -> str:
     """Returns pushed data to the script, automatically
     choosing canonical opcodes depending on the length of the data.
-    bytes -> bytes
+    hex -> hex
 
     ported from https://github.com/btcsuite/btcd/blob/fdc2bc867bda6b351191b5872d2da8270df00d13/txscript/scriptbuilder.go#L128
     """
-    assert_bytes(data)
+    data = bfh(data)
     from .transaction import opcodes
 
     data_len = len(data)
 
     # "small integer" opcodes
     if data_len == 0 or data_len == 1 and data[0] == 0:
-        return bytes([opcodes.OP_0])
+        return bh2u(bytes([opcodes.OP_0]))
     elif data_len == 1 and data[0] <= 16:
-        return bytes([opcodes.OP_1 - 1 + data[0]])
+        return bh2u(bytes([opcodes.OP_1 - 1 + data[0]]))
     elif data_len == 1 and data[0] == 0x81:
-        return bytes([opcodes.OP_1NEGATE])
-
-    return bfh(push_script(bh2u(data)))
+        return bh2u(bytes([opcodes.OP_1NEGATE]))
 
+    return op_push(data_len) + bh2u(data)
 
-def add_number_to_script(i):
-    """int -> bytes"""
-    return add_data_to_script(bfh(script_num_to_hex(i)))
 
+def add_number_to_script(i: int) -> bytes:
+    return bfh(push_script(script_num_to_hex(i)))
 
-def push_script(x):
-    return op_push(len(x)//2) + x
 
 def sha256(x):
     x = to_bytes(x, 'utf8')

lib/tests/test_bitcoin.py

@@ -12,9 +12,9 @@ from lib.bitcoin import (
     verify_message, deserialize_privkey, serialize_privkey, is_segwit_address,
     is_b58_address, address_to_scripthash, is_minikey, is_compressed, is_xpub,
     xpub_type, is_xprv, is_bip32_derivation, seed_type, EncodeBase58Check,
-    script_num_to_hex, add_data_to_script, add_number_to_script)
+    script_num_to_hex, push_script, add_number_to_script)
 from lib.transaction import opcodes
-from lib.util import bfh
+from lib.util import bfh, bh2u
 from lib import constants
 
 from . import TestCaseForTestnet
@@ -167,19 +167,19 @@ class Test_bitcoin(unittest.TestCase):
         self.assertEqual(script_num_to_hex(32768), '008000')
         self.assertEqual(script_num_to_hex(-32768), '008080')
 
-    def test_add_data_to_script(self):
+    def test_push_script(self):
         # https://github.com/bitcoin/bips/blob/master/bip-0062.mediawiki#push-operators
-        self.assertEqual(add_data_to_script(bfh('')), bytes([opcodes.OP_0]))
-        self.assertEqual(add_data_to_script(bfh('07')), bytes([opcodes.OP_7]))
-        self.assertEqual(add_data_to_script(bfh('10')), bytes([opcodes.OP_16]))
-        self.assertEqual(add_data_to_script(bfh('81')), bytes([opcodes.OP_1NEGATE]))
-        self.assertEqual(add_data_to_script(bfh('11')), bfh('0111'))
-        self.assertEqual(add_data_to_script(bfh(75 * '42')), bfh('4b' + 75 * '42'))
-        self.assertEqual(add_data_to_script(bfh(76 * '42')), bytes([opcodes.OP_PUSHDATA1]) + bfh('4c' + 76 * '42'))
-        self.assertEqual(add_data_to_script(bfh(100 * '42')), bytes([opcodes.OP_PUSHDATA1]) + bfh('64' + 100 * '42'))
-        self.assertEqual(add_data_to_script(bfh(255 * '42')), bytes([opcodes.OP_PUSHDATA1]) + bfh('ff' + 255 * '42'))
-        self.assertEqual(add_data_to_script(bfh(256 * '42')), bytes([opcodes.OP_PUSHDATA2]) + bfh('0001' + 256 * '42'))
-        self.assertEqual(add_data_to_script(bfh(520 * '42')), bytes([opcodes.OP_PUSHDATA2]) + bfh('0802' + 520 * '42'))
+        self.assertEqual(push_script(''), bh2u(bytes([opcodes.OP_0])))
+        self.assertEqual(push_script('07'), bh2u(bytes([opcodes.OP_7])))
+        self.assertEqual(push_script('10'), bh2u(bytes([opcodes.OP_16])))
+        self.assertEqual(push_script('81'), bh2u(bytes([opcodes.OP_1NEGATE])))
+        self.assertEqual(push_script('11'), '0111')
+        self.assertEqual(push_script(75 * '42'), '4b' + 75 * '42')
+        self.assertEqual(push_script(76 * '42'), bh2u(bytes([opcodes.OP_PUSHDATA1]) + bfh('4c' + 76 * '42')))
+        self.assertEqual(push_script(100 * '42'), bh2u(bytes([opcodes.OP_PUSHDATA1]) + bfh('64' + 100 * '42')))
+        self.assertEqual(push_script(255 * '42'), bh2u(bytes([opcodes.OP_PUSHDATA1]) + bfh('ff' + 255 * '42')))
+        self.assertEqual(push_script(256 * '42'), bh2u(bytes([opcodes.OP_PUSHDATA2]) + bfh('0001' + 256 * '42')))
+        self.assertEqual(push_script(520 * '42'), bh2u(bytes([opcodes.OP_PUSHDATA2]) + bfh('0802' + 520 * '42')))
 
     def test_add_number_to_script(self):
         # https://github.com/bitcoin/bips/blob/master/bip-0062.mediawiki#numbers

lib/transaction.py

@@ -27,6 +27,8 @@
 
 # Note: The deserialization code originally comes from ABE.
 
+from typing import Sequence, Union
+
 from .util import print_error, profiler
 
 from . import bitcoin
@@ -236,7 +238,7 @@ opcodes = Enumeration("Opcodes", [
 ])
 
 
-def script_GetOp(_bytes):
+def script_GetOp(_bytes : bytes):
     i = 0
     while i < len(_bytes):
         vch = None
@@ -382,7 +384,7 @@ def parse_scriptSig(d, _bytes):
                 bh2u(_bytes))
 
 
-def parse_redeemScript_multisig(redeem_script):
+def parse_redeemScript_multisig(redeem_script: bytes):
     dec2 = [ x for x in script_GetOp(redeem_script) ]
     try:
         m = dec2[0][0] - opcodes.OP_1 + 1
@@ -464,6 +466,18 @@ def parse_input(vds):
     return d
 
 
+def construct_witness(items: Sequence[Union[str, int, bytes]]) -> str:
+    """Constructs a witness from the given stack items."""
+    witness = var_int(len(items))
+    for item in items:
+        if type(item) is int:
+            item = bitcoin.script_num_to_hex(item)
+        elif type(item) is bytes:
+            item = bh2u(item)
+        witness += bitcoin.witness_push(item)
+    return witness
+
+
 def parse_witness(vds, txin):
     n = vds.read_compact_size()
     if n == 0:
@@ -474,7 +488,7 @@ def parse_witness(vds, txin):
     # now 'n' is the number of items in the witness
     w = list(bh2u(vds.read_bytes(vds.read_compact_size())) for i in range(n))
 
-    txin['witness'] = var_int(n) + ''.join(witness_push(i) for i in w)
+    txin['witness'] = construct_witness(w)
 
     # FIXME: witness version > 0 will probably fail here.
     # For native segwit, we would need the scriptPubKey of the parent txn
@@ -750,11 +764,10 @@ class Transaction:
         if witness is None:
             pubkeys, sig_list = self.get_siglist(txin, estimate_size)
             if txin['type'] in ['p2wpkh', 'p2wpkh-p2sh']:
-                witness = var_int(2) + witness_push(sig_list[0]) + witness_push(pubkeys[0])
+                witness = construct_witness([sig_list[0], pubkeys[0]])
             elif txin['type'] in ['p2wsh', 'p2wsh-p2sh']:
-                n = len(sig_list) + 2
                 witness_script = multisig_script(pubkeys, txin['num_sig'])
-                witness = var_int(n) + '00' + ''.join(witness_push(x) for x in sig_list) + witness_push(witness_script)
+                witness = construct_witness([0, *sig_list, witness_script])
             else:
                 raise Exception('wrong txin type:', txin['type'])
         if self.is_txin_complete(txin) or estimate_size: