1import unittest
2from util import *
3
4
5PRV_HEX = utf8('0C28FCA386C7A227600B2FE50B7CAE11EC86D3BF1FBE471BE89827E19D72AA1D')
6PRV_WIF_UNCOMPRESS = utf8('5HueCGU8rMjxEXxiPuD5BDku4MkFqeZyd4dZ1jvhTVqvbTLvyTJ')
7PRV_WIF_COMPRESS = utf8('KwdMAjGmerYanjeui5SHS7JkmpZvVipYvB2LJGU1ZxJwYvP98617')
8PREFIX = 0x80
9VERSION = 0x00
10
11
12class WIFTests(unittest.TestCase):
13
14    def test_wif_from_bytes(self):
15        prv, prv_len = make_cbuffer(PRV_HEX)
16
17        invalid_args = [
18            (None, prv_len, PREFIX, 0),  # Missing private key
19            (prv, 0, PREFIX, 0),  # Incorrect len
20            (prv, 0, 0x100, 0),  # Unsupported PREFIX
21            (prv, prv_len, PREFIX, 2),  # Unsupported flag
22        ]
23
24        for args in invalid_args:
25            ret, out = wally_wif_from_bytes(*args)
26            self.assertEqual(ret, WALLY_EINVAL)
27
28        for flag, expected_wif in [
29            (0, PRV_WIF_COMPRESS),
30            (1, PRV_WIF_UNCOMPRESS),
31        ]:
32            ret, wif = wally_wif_from_bytes(prv, prv_len, PREFIX, flag)
33            self.assertEqual(ret, WALLY_OK)
34            self.assertEqual(utf8(wif), expected_wif)
35
36    def test_wif_to_bytes(self):
37        buf, buf_len = make_cbuffer('00'*32)
38
39        # wif_to_bytes
40        invalid_args = [
41            (None, PREFIX, 0, buf, buf_len),  # Empty wif
42            (PRV_WIF_COMPRESS, 0x81, 0, buf, buf_len),  # Not matching PREFIX
43            (PRV_WIF_COMPRESS, 0x100, 0, buf, buf_len),  # Unsupported PREFIX
44            (PRV_WIF_COMPRESS, PREFIX, 2, buf, buf_len),  # Unsupported flag
45            (PRV_WIF_COMPRESS, PREFIX, 1, buf, buf_len),  # Inconsistent flag
46            (PRV_WIF_COMPRESS, PREFIX, 0, None, buf_len),  # Empty output
47            (PRV_WIF_COMPRESS, PREFIX, 0, buf, 31),  # Unsupported len
48        ]
49
50        for args in invalid_args:
51            self.assertEqual(wally_wif_to_bytes(*args), WALLY_EINVAL)
52
53        # wif_is_uncompressed
54        invalid_args = [
55            '',  # Empty
56            '11111',  # Incorrect checksum
57            'yNb7j1viLcZunrTHozyfJPTZJrprRSPpY485Lwzq1CFQPxF7A',  # Invalid length
58            'KwDiBf89QgGbjEhKnhXJuH7LrciVrZi3qYjgd9M7rFU73NUBByJr',  # Unexpected ending byte (not 0x01)
59        ]
60
61        for wif in invalid_args:
62            ret, _ = wally_wif_is_uncompressed(utf8(wif))
63            self.assertEqual(ret, WALLY_EINVAL)
64
65        # wif_to_public_key
66        pub, pub_len = make_cbuffer('00' * 65)
67
68        invalid_args = [
69            (None, PREFIX, pub, pub_len),  # Empty wif
70            (PRV_WIF_COMPRESS, 0x100, pub, pub_len),  # Empty wif
71            (PRV_WIF_COMPRESS, PREFIX, None, pub_len),  # Empty pubkey
72        ]
73
74        for args in invalid_args:
75            self.assertEqual(wally_wif_to_public_key(*args), (WALLY_EINVAL, 0))
76
77        # If the output length is incorrect, the correct one is returned
78        invalid_len = [
79            (PRV_WIF_COMPRESS, PREFIX, pub, 32),
80            (PRV_WIF_UNCOMPRESS, PREFIX, pub, 64),
81        ]
82
83        for args in invalid_len:
84            self.assertEqual(wally_wif_to_public_key(*args), (WALLY_OK, args[3] + 1))
85
86        # Valid args
87        for is_uncompressed, wif in [
88            (1, PRV_WIF_UNCOMPRESS),
89            (0, PRV_WIF_COMPRESS),
90        ]:
91            self.assertEqual(wally_wif_is_uncompressed(wif), (WALLY_OK, is_uncompressed))
92            self.assertEqual(wally_wif_to_bytes(wif, PREFIX, is_uncompressed, buf, buf_len), WALLY_OK)
93            self.assertEqual(h(buf).upper(), PRV_HEX)
94
95            pub, pub_len = make_cbuffer('00' * (1 + 32 * (is_uncompressed + 1)))
96            ret, written = wally_wif_to_public_key(wif, PREFIX, pub, pub_len)
97            self.assertEqual(ret, WALLY_OK)
98            self.assertEqual(written, pub_len)
99            self.assertEqual(pub, self.private_to_public(buf, is_uncompressed))
100
101            exp_addr = self.public_to_address(pub, VERSION)
102            ret, addr = wally_wif_to_address(wif, PREFIX, VERSION)
103            self.assertEqual(ret, WALLY_OK)
104            self.assertEqual(addr, exp_addr)
105
106    def private_to_public(self, prv, is_uncompressed):
107        pub, pub_len = make_cbuffer('00' * 33)
108        self.assertEqual(wally_ec_public_key_from_private_key(prv, 32, pub, pub_len), WALLY_OK)
109        if not is_uncompressed:
110            return pub
111        _pub, _pub_len = make_cbuffer('00' * 65)
112        self.assertEqual(wally_ec_public_key_decompress(pub, pub_len, _pub, _pub_len), WALLY_OK)
113        return _pub
114
115    def public_to_address(self, pub, version):
116        h, h_len = make_cbuffer('00' * 20)
117        wally_hash160(pub, len(pub), h, h_len)
118        return wally_base58_from_bytes(bytes(bytearray([version])) + h, 21, 1)[1]
119
120
121if __name__ == '__main__':
122    unittest.main()
123