Inside Bitcoin

1#!/usr/bin/env python3
2# Copyright (c) 2022-present The Bitcoin Core developers
3# Distributed under the MIT software license, see the accompanying
4# file COPYING or http://www.opensource.org/licenses/mit-license.php.
5
6import base64
7import struct
8
9from io import BytesIO
10
11from .util import assert_equal
12from .messages import (
13    CTransaction,
14    deser_string,
15    deser_compact_size,
16    from_binary,
17    ser_compact_size,
18)
19
20
21# global types
22PSBT_GLOBAL_UNSIGNED_TX = 0x00
23PSBT_GLOBAL_XPUB = 0x01
24PSBT_GLOBAL_TX_VERSION = 0x02
25PSBT_GLOBAL_FALLBACK_LOCKTIME = 0x03
26PSBT_GLOBAL_INPUT_COUNT = 0x04
27PSBT_GLOBAL_OUTPUT_COUNT = 0x05
28PSBT_GLOBAL_TX_MODIFIABLE = 0x06
29PSBT_GLOBAL_VERSION = 0xfb
30PSBT_GLOBAL_PROPRIETARY = 0xfc
31
32# per-input types
33PSBT_IN_NON_WITNESS_UTXO = 0x00
34PSBT_IN_WITNESS_UTXO = 0x01
35PSBT_IN_PARTIAL_SIG = 0x02
36PSBT_IN_SIGHASH_TYPE = 0x03
37PSBT_IN_REDEEM_SCRIPT = 0x04
38PSBT_IN_WITNESS_SCRIPT = 0x05
39PSBT_IN_BIP32_DERIVATION = 0x06
40PSBT_IN_FINAL_SCRIPTSIG = 0x07
41PSBT_IN_FINAL_SCRIPTWITNESS = 0x08
42PSBT_IN_POR_COMMITMENT = 0x09
43PSBT_IN_RIPEMD160 = 0x0a
44PSBT_IN_SHA256 = 0x0b
45PSBT_IN_HASH160 = 0x0c
46PSBT_IN_HASH256 = 0x0d
47PSBT_IN_PREVIOUS_TXID = 0x0e
48PSBT_IN_OUTPUT_INDEX = 0x0f
49PSBT_IN_SEQUENCE = 0x10
50PSBT_IN_REQUIRED_TIME_LOCKTIME = 0x11
51PSBT_IN_REQUIRED_HEIGHT_LOCKTIME = 0x12
52PSBT_IN_TAP_KEY_SIG = 0x13
53PSBT_IN_TAP_SCRIPT_SIG = 0x14
54PSBT_IN_TAP_LEAF_SCRIPT = 0x15
55PSBT_IN_TAP_BIP32_DERIVATION = 0x16
56PSBT_IN_TAP_INTERNAL_KEY = 0x17
57PSBT_IN_TAP_MERKLE_ROOT = 0x18
58PSBT_IN_MUSIG2_PARTICIPANT_PUBKEYS = 0x1a
59PSBT_IN_MUSIG2_PUB_NONCE = 0x1b
60PSBT_IN_MUSIG2_PARTIAL_SIG = 0x1c
61PSBT_IN_PROPRIETARY = 0xfc
62
63# per-output types
64PSBT_OUT_REDEEM_SCRIPT = 0x00
65PSBT_OUT_WITNESS_SCRIPT = 0x01
66PSBT_OUT_BIP32_DERIVATION = 0x02
67PSBT_OUT_AMOUNT = 0x03
68PSBT_OUT_SCRIPT = 0x04
69PSBT_OUT_TAP_INTERNAL_KEY = 0x05
70PSBT_OUT_TAP_TREE = 0x06
71PSBT_OUT_TAP_BIP32_DERIVATION = 0x07
72PSBT_OUT_MUSIG2_PARTICIPANT_PUBKEYS = 0x08
73PSBT_OUT_PROPRIETARY = 0xfc
74
75
76class PSBTMap:
77    """Class for serializing and deserializing PSBT maps"""
78
79    def __init__(self, map=None):
80        self.map = map if map is not None else {}
81
82    def deserialize(self, f):
83        m = {}
84        while True:
85            k = deser_string(f)
86            if len(k) == 0:
87                break
88            v = deser_string(f)
89            if len(k) == 1:
90                k = k[0]
91            assert k not in m
92            m[k] = v
93        self.map = m
94
95    def serialize(self):
96        m = b""
97        for k,v in self.map.items():
98            if isinstance(k, int) and 0 <= k and k <= 255:
99                k = bytes([k])
100            if isinstance(v, list):
101                assert all(type(elem) is bytes for elem in v)
102                v = b"".join(v)  # simply concatenate the byte-strings w/o size prefixes
103            m += ser_compact_size(len(k)) + k
104            m += ser_compact_size(len(v)) + v
105        m += b"\x00"
106        return m
107
108class PSBT:
109    """Class for serializing and deserializing PSBTs"""
110
111    def __init__(self, *, g=None, i=None, o=None):
112        self.g = g if g is not None else PSBTMap()
113        self.i = i if i is not None else []
114        self.o = o if o is not None else []
115        self.version = None
116
117    def deserialize(self, f):
118        assert_equal(f.read(5), b"psbt\xff")
119        self.g = from_binary(PSBTMap, f)
120
121        self.version = 0
122        if PSBT_GLOBAL_VERSION in self.g.map:
123            self.version = struct.unpack("<I", self.g.map[PSBT_GLOBAL_VERSION])[0]
124            assert self.version in [0, 2]
125        if self.version == 2:
126            assert PSBT_GLOBAL_INPUT_COUNT in self.g.map
127            assert PSBT_GLOBAL_OUTPUT_COUNT in self.g.map
128            in_count = deser_compact_size(BytesIO(self.g.map[PSBT_GLOBAL_INPUT_COUNT]))
129            out_count = deser_compact_size(BytesIO(self.g.map[PSBT_GLOBAL_OUTPUT_COUNT]))
130        else:
131            assert PSBT_GLOBAL_UNSIGNED_TX in self.g.map
132            tx = from_binary(CTransaction, self.g.map[PSBT_GLOBAL_UNSIGNED_TX])
133            in_count = len(tx.vin)
134            out_count = len(tx.vout)
135
136        self.i = [from_binary(PSBTMap, f) for _ in range(in_count)]
137        self.o = [from_binary(PSBTMap, f) for _ in range(out_count)]
138        return self
139
140    def serialize(self):
141        assert isinstance(self.g, PSBTMap)
142        assert isinstance(self.i, list) and all(isinstance(x, PSBTMap) for x in self.i)
143        assert isinstance(self.o, list) and all(isinstance(x, PSBTMap) for x in self.o)
144        if self.version is not None and self.version == 2:
145            self.g.map[PSBT_GLOBAL_INPUT_COUNT] = ser_compact_size(len(self.i))
146            self.g.map[PSBT_GLOBAL_OUTPUT_COUNT] = ser_compact_size(len(self.o))
147        if self.version is None or (self.version is not None and self.version == 0):
148            assert PSBT_GLOBAL_UNSIGNED_TX in self.g.map
149            tx = from_binary(CTransaction, self.g.map[PSBT_GLOBAL_UNSIGNED_TX])
150            assert_equal(len(tx.vin), len(self.i))
151            assert_equal(len(tx.vout), len(self.o))
152
153        psbt = [x.serialize() for x in [self.g] + self.i + self.o]
154        return b"psbt\xff" + b"".join(psbt)
155
156    def make_blank(self):
157        """
158        Remove all fields except for required fields depending on version
159        """
160        if self.version == 0:
161            for m in self.i + self.o:
162                m.map.clear()
163
164            self.g = PSBTMap(map={PSBT_GLOBAL_UNSIGNED_TX: self.g.map[PSBT_GLOBAL_UNSIGNED_TX]})
165        elif self.version == 2:
166            self.g = PSBTMap(map={
167                PSBT_GLOBAL_TX_VERSION: self.g.map[PSBT_GLOBAL_TX_VERSION],
168                PSBT_GLOBAL_INPUT_COUNT: self.g.map[PSBT_GLOBAL_INPUT_COUNT],
169                PSBT_GLOBAL_OUTPUT_COUNT: self.g.map[PSBT_GLOBAL_OUTPUT_COUNT],
170                PSBT_GLOBAL_VERSION: self.g.map[PSBT_GLOBAL_VERSION],
171                PSBT_GLOBAL_FALLBACK_LOCKTIME: self.g.map[PSBT_GLOBAL_FALLBACK_LOCKTIME],
172            })
173
174            new_i = []
175            for m in self.i:
176                new_i.append(PSBTMap(map={
177                    PSBT_IN_PREVIOUS_TXID: m.map[PSBT_IN_PREVIOUS_TXID],
178                    PSBT_IN_OUTPUT_INDEX: m.map[PSBT_IN_OUTPUT_INDEX],
179                }))
180            self.i = new_i
181
182            new_o = []
183            for m in self.o:
184                new_o.append(PSBTMap(map={
185                    PSBT_OUT_SCRIPT: m.map[PSBT_OUT_SCRIPT],
186                    PSBT_OUT_AMOUNT: m.map[PSBT_OUT_AMOUNT],
187                }))
188            self.o = new_o
189        else:
190            assert False
191
192    def to_base64(self):
193        return base64.b64encode(self.serialize()).decode("utf8")
194
195    @classmethod
196    def from_base64(cls, b64psbt):
197        return from_binary(cls, base64.b64decode(b64psbt))
198