1"""Tests for PSBT wrappers"""
2import unittest
3from wallycore import *
4
5SAMPLE = "cHNidP8BAFICAAAAAZ38ZijCbFiZ/hvT3DOGZb/VXXraEPYiCXPfLTht7BJ2AQAAAAD/////AfA9zR0AAAAAFgAUezoAv9wU0neVwrdJAdCdpu8TNXkAAAAATwEENYfPAto/0AiAAAAAlwSLGtBEWx7IJ1UXcnyHtOTrwYogP/oPlMAVZr046QADUbdDiH7h1A3DKmBDck8tZFmztaTXPa7I+64EcvO8Q+IM2QxqT64AAIAAAACATwEENYfPAto/0AiAAAABuQRSQnE5zXjCz/JES+NTzVhgXj5RMoXlKLQH+uP2FzUD0wpel8itvFV9rCrZp+OcFyLrrGnmaLbyZnzB1nHIPKsM2QxqT64AAIABAACAAAEBKwBlzR0AAAAAIgAgLFSGEmxJeAeagU4TcV1l82RZ5NbMre0mbQUIZFuvpjIBBUdSIQKdoSzbWyNWkrkVNq/v5ckcOrlHPY5DtTODarRWKZyIcSEDNys0I07Xz5wf6l0F1EFVeSe+lUKxYusC4ass6AIkwAtSriIGAp2hLNtbI1aSuRU2r+/lyRw6uUc9jkO1M4NqtFYpnIhxENkMak+uAACAAAAAgAAAAAAiBgM3KzQjTtfPnB/qXQXUQVV5J76VQrFi6wLhqyzoAiTACxDZDGpPrgAAgAEAAIAAAAAAACICA57/H1R6HV+S36K6evaslxpL0DukpzSwMVaiVritOh75EO3kXMUAAACAAAAAgAEAAIAA"
6
7
8class PSBTTests(unittest.TestCase):
9
10    def _try_invalid(self, fn, psbt, *args):
11        with self.assertRaises(ValueError):
12            fn(None, 0, *args) # Null PSBT
13        with self.assertRaises(ValueError):
14            fn(psbt, 1, *args) # Invalid index
15
16    def _try_set(self, fn, psbt, valid_value, null_value=None):
17        fn(psbt, 0, valid_value) # Set
18        fn(psbt, 0, null_value) # Un-set
19        self._try_invalid(fn, psbt, valid_value)
20
21    def _try_get_set_b(self, setfn, getfn, lenfn, psbt, valid_value, null_value=None):
22        self._try_set(setfn, psbt, valid_value, null_value)
23        setfn(psbt, 0, valid_value) # Set
24        self._try_invalid(lenfn, psbt)
25        self._try_invalid(getfn, psbt)
26        ret = getfn(psbt, 0) # Get
27        self.assertEqual(valid_value, ret)
28
29    def _try_get_set_m(self, setfn, sizefn, lenfn, getfn, findfn, psbt, valid_value, valid_item):
30        self._try_set(setfn, psbt, valid_value, None)
31        self._try_invalid(sizefn, psbt)
32        self.assertEqual(sizefn(psbt, 0), 0)
33        setfn(psbt, 0, valid_value) # Set
34        self.assertEqual(sizefn(psbt, 0), 1) # 1 item in the map
35        self._try_invalid(lenfn, psbt, 0)
36        with self.assertRaises(ValueError):
37            lenfn(psbt, 0, 1) # Invalid subindex
38        map_val = getfn(psbt, 0, 0)
39        self.assertTrue(len(map_val) > 0)
40        self.assertEqual(lenfn(psbt, 0, 0), len(map_val))
41        self._try_invalid(findfn, psbt, map_val)
42        self.assertEqual(findfn(psbt, 0, valid_item), 1)
43
44
45    def test_psbt(self):
46        psbt = psbt_from_base64(SAMPLE)
47
48        # Roundtrip to/from bytes
49        psbt_bytes = psbt_to_bytes(psbt, 0)
50        psbt_tmp = psbt_from_bytes(psbt_bytes)
51        self.assertEqual(hex_from_bytes(psbt_bytes),
52                         hex_from_bytes(psbt_to_bytes(psbt_tmp, 0)))
53
54        self.assertIsNotNone(psbt_get_global_tx(psbt))
55
56        for fn, ret in [(psbt_get_version, 0),
57                        (psbt_get_num_inputs, 1),
58                        (psbt_get_num_outputs, 1)]:
59            self.assertEqual(fn(psbt), ret)
60            with self.assertRaises(ValueError):
61                fn(None) # Null PSBT
62
63        # Conversion to base64 should round trip
64        self.assertEqual(psbt_to_base64(psbt, 0), SAMPLE)
65
66        # Combining with ourselves shouldn't change the PSBT
67        psbt_combine(psbt, psbt)
68        self.assertEqual(psbt_to_base64(psbt, 0), SAMPLE)
69
70        # Test setters
71        dummy_tx = psbt_get_global_tx(psbt)
72        self.assertIsNotNone(dummy_tx)
73
74        dummy_txout = tx_output_init(1234567, bytearray(b'\x00' * 33))
75
76        dummy_witness = tx_witness_stack_init(5)
77        self.assertIsNotNone(dummy_witness)
78
79        dummy_bytes = bytearray(b'\x00' * 32)
80        dummy_pubkey = bytearray(b'\x02'* EC_PUBLIC_KEY_LEN)
81        dummy_fingerprint = bytearray(b'\x00' * BIP32_KEY_FINGERPRINT_LEN)
82        dummy_path = [1234, 1234, 1234]
83        dummy_sig = bytearray(b'\x00' * 72)
84        if is_elements_build():
85            dummy_nonce = bytearray(b'\x00' * WALLY_TX_ASSET_CT_NONCE_LEN)
86            dummy_bf = bytearray(b'\x00' * BLINDING_FACTOR_LEN)
87            dummy_commitment = bytearray(b'\x00' * ASSET_COMMITMENT_LEN)
88            dummy_asset = bytearray(b'\x00' * ASSET_TAG_LEN)
89
90        dummy_keypaths = map_init(0)
91        self.assertIsNotNone(dummy_keypaths)
92        map_add_keypath_item(dummy_keypaths, dummy_pubkey, dummy_fingerprint, dummy_path)
93        self.assertEqual(map_find(dummy_keypaths, dummy_pubkey), 1)
94
95        dummy_signatures = map_init(0)
96        self.assertIsNotNone(dummy_signatures)
97        map_add(dummy_signatures, dummy_pubkey, dummy_sig)
98        self.assertEqual(map_find(dummy_signatures, dummy_pubkey), 1)
99
100        dummy_unknowns = map_init(1)
101        self.assertIsNotNone(dummy_unknowns)
102        map_add(dummy_unknowns, dummy_pubkey, dummy_fingerprint)
103        self.assertEqual(map_find(dummy_unknowns, dummy_pubkey), 1)
104
105        #
106        # Inputs
107        #
108        self._try_set(psbt_set_input_utxo, psbt, dummy_tx)
109        self._try_invalid(psbt_get_input_utxo, psbt)
110        self._try_set(psbt_set_input_witness_utxo, psbt, dummy_txout)
111        self._try_invalid(psbt_get_input_witness_utxo, psbt)
112        self._try_get_set_b(psbt_set_input_redeem_script,
113                            psbt_get_input_redeem_script,
114                            psbt_get_input_redeem_script_len, psbt, dummy_bytes)
115        self._try_get_set_b(psbt_set_input_witness_script,
116                            psbt_get_input_witness_script,
117                            psbt_get_input_witness_script_len, psbt, dummy_bytes)
118        self._try_get_set_b(psbt_set_input_final_scriptsig,
119                            psbt_get_input_final_scriptsig,
120                            psbt_get_input_final_scriptsig_len, psbt, dummy_bytes)
121        self._try_set(psbt_set_input_final_witness, psbt, dummy_witness)
122        self._try_invalid(psbt_get_input_final_witness, psbt)
123        self._try_get_set_m(psbt_set_input_keypaths,
124                            psbt_get_input_keypaths_size,
125                            psbt_get_input_keypath_len,
126                            psbt_get_input_keypath,
127                            psbt_find_input_keypath,
128                            psbt, dummy_keypaths, dummy_pubkey)
129        self._try_get_set_m(psbt_set_input_signatures,
130                            psbt_get_input_signatures_size,
131                            psbt_get_input_signature_len,
132                            psbt_get_input_signature,
133                            psbt_find_input_signature,
134                            psbt, dummy_signatures, dummy_pubkey)
135        self._try_get_set_m(psbt_set_input_unknowns,
136                            psbt_get_input_unknowns_size,
137                            psbt_get_input_unknown_len,
138                            psbt_get_input_unknown,
139                            psbt_find_input_unknown,
140                            psbt, dummy_unknowns, dummy_pubkey)
141        self._try_set(psbt_set_input_sighash, psbt, 0xff, 0x0)
142        self.assertEqual(psbt_get_input_sighash(psbt, 0), 0)
143        self._try_invalid(psbt_get_input_sighash, psbt)
144
145        if is_elements_build():
146            self._try_set(psbt_set_input_value, psbt, 1234567, 0)
147            self._try_invalid(psbt_has_input_value, psbt)
148            self._try_invalid(psbt_get_input_value, psbt)
149            self._try_invalid(psbt_clear_input_value, psbt)
150            self.assertEqual(psbt_has_input_value(psbt, 0), 1)
151            psbt_clear_input_value(psbt, 0)
152            self.assertEqual(psbt_has_input_value(psbt, 0), 0)
153            self._try_get_set_b(psbt_set_input_vbf,
154                                psbt_get_input_vbf,
155                                psbt_get_input_vbf_len, psbt, dummy_bf)
156            self._try_get_set_b(psbt_set_input_asset,
157                                psbt_get_input_asset,
158                                psbt_get_input_asset_len, psbt, dummy_asset)
159            self._try_get_set_b(psbt_set_input_abf,
160                                psbt_get_input_abf,
161                                psbt_get_input_abf_len, psbt, dummy_bf)
162            self._try_set(psbt_set_input_pegin_tx, psbt, dummy_tx)
163            self._try_invalid(psbt_get_input_pegin_tx, psbt)
164            self._try_get_set_b(psbt_set_input_txoutproof,
165                                psbt_get_input_txoutproof,
166                                psbt_get_input_txoutproof_len, psbt, dummy_bytes)
167            self._try_get_set_b(psbt_set_input_genesis_blockhash,
168                                psbt_get_input_genesis_blockhash,
169                                psbt_get_input_genesis_blockhash_len, psbt, dummy_bytes)
170            self._try_get_set_b(psbt_set_input_claim_script,
171                                psbt_get_input_claim_script,
172                                psbt_get_input_claim_script_len, psbt, dummy_bytes)
173
174        #
175        # Outputs
176        #
177        self._try_get_set_b(psbt_set_output_redeem_script,
178                            psbt_get_output_redeem_script,
179                            psbt_get_output_redeem_script_len, psbt, dummy_bytes)
180        self._try_get_set_b(psbt_set_output_witness_script,
181                            psbt_get_output_witness_script,
182                            psbt_get_output_witness_script_len, psbt, dummy_bytes)
183        self._try_get_set_m(psbt_set_output_keypaths,
184                            psbt_get_output_keypaths_size,
185                            psbt_get_output_keypath_len,
186                            psbt_get_output_keypath,
187                            psbt_find_output_keypath,
188                            psbt, dummy_keypaths, dummy_pubkey)
189        self._try_get_set_m(psbt_set_output_unknowns,
190                            psbt_get_output_unknowns_size,
191                            psbt_get_output_unknown_len,
192                            psbt_get_output_unknown,
193                            psbt_find_output_unknown,
194                            psbt, dummy_unknowns, dummy_pubkey)
195        if is_elements_build():
196            self._try_get_set_b(psbt_set_output_blinding_pubkey,
197                                psbt_get_output_blinding_pubkey,
198                                psbt_get_output_blinding_pubkey_len, psbt, dummy_pubkey)
199            self._try_get_set_b(psbt_set_output_value_commitment,
200                                psbt_get_output_value_commitment,
201                                psbt_get_output_value_commitment_len, psbt, dummy_commitment)
202            self._try_get_set_b(psbt_set_output_vbf,
203                                psbt_get_output_vbf,
204                                psbt_get_output_vbf_len, psbt, dummy_bf)
205            self._try_get_set_b(psbt_set_output_asset_commitment,
206                                psbt_get_output_asset_commitment,
207                                psbt_get_output_asset_commitment_len, psbt, dummy_commitment)
208            self._try_get_set_b(psbt_set_output_abf,
209                                psbt_get_output_abf,
210                                psbt_get_output_abf_len, psbt, dummy_bf)
211            self._try_get_set_b(psbt_set_output_nonce,
212                                psbt_get_output_nonce,
213                                psbt_get_output_nonce_len, psbt, dummy_nonce)
214            self._try_get_set_b(psbt_set_output_rangeproof,
215                                psbt_get_output_rangeproof,
216                                psbt_get_output_rangeproof_len, psbt, dummy_bytes)
217            self._try_get_set_b(psbt_set_output_surjectionproof,
218                                psbt_get_output_surjectionproof,
219                                psbt_get_output_surjectionproof_len, psbt, dummy_bytes)
220
221
222if __name__ == '__main__':
223    unittest.main()
224