1import unittest
2from util import *
3
4class AddressCase(object):
5    def __init__(self, lines):
6        # https://github.com/ThePiachu/Bitcoin-Unit-Tests/blob/master/Address
7        self.ripemd_network = lines[4]
8        self.checksummed = lines[8]
9        self.base58 = lines[9]
10
11class Base58Tests(unittest.TestCase):
12
13    FLAG_CHECKSUM = 0x1
14    CHECKSUM_LEN = 4
15
16    def setUp(self):
17        if not hasattr(self, 'cases'):
18            # Test cases from https://github.com/ThePiachu/Bitcoin-Unit-Tests/
19            self.cases = []
20            cur = []
21            with open(root_dir + 'src/data/address_vectors.txt', 'r') as f:
22                for l in f.readlines():
23                    if len(l.strip()):
24                        cur.append(l.strip())
25                    else:
26                        self.cases.append(AddressCase(cur))
27                        cur = []
28
29    def encode(self, hex_in, flags):
30        buf, buf_len = make_cbuffer(hex_in)
31        ret, base58 = wally_base58_from_bytes(buf, buf_len, flags)
32        self.assertEqual(ret, WALLY_EINVAL if base58 is None else WALLY_OK)
33        return base58
34
35    def decode(self, str_in, flags):
36        buf, buf_len = make_cbuffer('00' * 1024)
37        ret, buf_len = wally_base58_to_bytes(utf8(str_in), flags, buf, buf_len)
38        self.assertEqual(ret, WALLY_OK)
39        self.assertNotEqual(buf_len, 0)
40        # Check that just computing the size returns us the actual size
41        ret, bin_len = wally_base58_get_length(utf8(str_in))
42        self.assertEqual(ret, WALLY_OK)
43        if flags == self.FLAG_CHECKSUM:
44            bin_len -= self.CHECKSUM_LEN
45        self.assertEqual(bin_len, buf_len)
46        return h(buf)[0:buf_len * 2].upper()
47
48
49    def test_address_vectors(self):
50        """Tests for encoding and decoding with and without checksums"""
51
52        for c in self.cases:
53            # Checksummed should match directly in base 58
54            base58 = self.encode(c.checksummed, 0)
55            self.assertEqual(base58, c.base58)
56            # Decode it and make sure it matches checksummed again
57            decoded = self.decode(c.base58, 0)
58            self.assertEqual(decoded, utf8(c.checksummed))
59
60            # Compute the checksum in the call
61            base58 = self.encode(c.ripemd_network, self.FLAG_CHECKSUM)
62            self.assertEqual(base58, c.base58)
63
64            # Decode without checksum validation/stripping, should match
65            # checksummed value
66            decoded = self.decode(c.base58, 0)
67            self.assertEqual(decoded, utf8(c.checksummed))
68
69            # Decode with checksum validation/stripping and compare
70            # to original ripemd + network
71            decoded = self.decode(c.base58, self.FLAG_CHECKSUM)
72            self.assertEqual(decoded, utf8(c.ripemd_network))
73
74
75    def test_to_bytes(self):
76        buf, buf_len = make_cbuffer('00' * 1024)
77
78        # Bad input base58 strings
79        for bad in [ '',        # Empty string can't be represented
80                     '0',       # Forbidden ASCII character
81                     'x0',      # Forbidden ASCII character, internal
82                     '\x80',    # High bit set
83                     'x\x80x',  # High bit set, internal
84                   ]:
85            ret, _ = wally_base58_to_bytes(utf8(bad), 0, buf, buf_len)
86            self.assertEqual(ret, WALLY_EINVAL)
87
88        # Bad checksummed base58 strings
89        for bad in [ # libbase58: decode-b58c-fail
90                    '19DXstMaV43WpYg4ceREiiTv2UntmoiA9a',
91                    # libbase58: decode-b58c-toolong
92                    '1119DXstMaV43WpYg4ceREiiTv2UntmoiA9a',
93                    # libbase58: decode-b58c-tooshort
94                    '111111111111111111114oLvT2']:
95            ret, _ = wally_base58_to_bytes(utf8(bad), self.FLAG_CHECKSUM, buf, buf_len)
96            self.assertEqual(ret, WALLY_EINVAL)
97
98        for base58 in ['BXvDbH', '16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM']:
99            ret, out_len = wally_base58_get_length(utf8(base58))
100            # Output buffer too small returns OK and the number of bytes required
101            ret, bin_len = wally_base58_to_bytes(utf8(base58), 0, buf, out_len - 1)
102            self.assertEqual((ret, bin_len), (WALLY_OK, out_len))
103            # Unknown flags
104            ret, _ = wally_base58_to_bytes(utf8(base58), 0x7, buf, buf_len)
105            self.assertEqual(ret, WALLY_EINVAL)
106
107        # If we ask for checksum validation/removal the output buffer
108        # must have room for a checksum.
109        ret, bin_len = wally_base58_to_bytes(utf8('1'), self.FLAG_CHECKSUM,
110                                             buf, self.CHECKSUM_LEN)
111        self.assertEqual(ret, WALLY_EINVAL)
112        # Also the input string must contain at least CHECKSUM_LEN + 1
113        # bytes worth of data
114        for i in range(self.CHECKSUM_LEN):
115            ret, bin_len = wally_base58_to_bytes(utf8('1'*i), self.FLAG_CHECKSUM,
116                                                 buf, buf_len)
117            self.assertEqual(ret, WALLY_EINVAL)
118
119        # Leading ones become zeros
120        for i in range(1, 10):
121            self.assertEqual(self.decode('1' * i, 0), utf8('00' * i))
122
123        # Vectors from https://github.com/bitcoinj/bitcoinj/
124        self.assertEqual(self.decode('16Ho7Hs', 0), utf8('00CEF022FA'))
125        self.assertEqual(self.decode('4stwEBjT6FYyVV', self.FLAG_CHECKSUM),
126                                     utf8('45046252208D'))
127        base58 = '93VYUMzRG9DdbRP72uQXjaWibbQwygnvaCu9DumcqDjGybD864T'
128        ret = self.decode(base58, self.FLAG_CHECKSUM)
129        expected = 'EFFB309E964684B54E6069F146E2CD6DA' \
130                   'E936B711A7A98DF4097156B9FC9B344EB'
131        self.assertEqual(ret, utf8(expected))
132
133
134    def test_from_bytes(self):
135
136        # Leading zeros become ones
137        for i in range(1, 10):
138            self.assertEqual(self.encode('00' * i, 0), '1' * i)
139
140        # Invalid flags
141        self.assertEqual(self.encode('00', 0x7), None)
142
143        buf, buf_len = make_cbuffer('00' * 8)
144
145        FAIL_RET = (WALLY_EINVAL, None)
146        # O length buffer, no checksum -> NULL
147        self.assertEqual(wally_base58_from_bytes(buf, 0, 0), FAIL_RET)
148
149        # O length buffer, append checksum -> NULL
150        self.assertEqual(wally_base58_from_bytes(buf, 0, self.FLAG_CHECKSUM), FAIL_RET)
151
152        # Vectors from https://github.com/bitcoinj/bitcoinj/
153        self.assertEqual(self.encode('00CEF022FA', 0), '16Ho7Hs')
154        self.assertEqual(self.encode('45046252208D', self.FLAG_CHECKSUM),
155                                     '4stwEBjT6FYyVV')
156
157
158
159if __name__ == '__main__':
160    unittest.main()
161