1import unittest
2from util import *
3
4K_MAIN, K_TEST, K_COMP, K_EC, K_CHECK, K_RAW, K_ORDER = 0, 7, 256, 512, 1024, 2048, 4096
5
6# BIP38 Vectors from
7# https://github.com/bitcoin/bips/blob/master/bip-0038.mediawiki
8cases = [
9    [ 'CBF4B9F70470856BB4F40F80B87EDB90865997FFEE6DF315AB166D713AF433A5',
10      'TestingOneTwoThree',
11      K_MAIN,
12      '6PRVWUbkzzsbcVac2qwfssoUJAN1Xhrg6bNk8J7Nzm5H7kxEbn2Nh2ZoGg' ],
13    [ '09C2686880095B1A4C249EE3AC4EEA8A014F11E6F986D0B5025AC1F39AFBD9AE',
14      'Satoshi',
15      K_MAIN,
16      '6PRNFFkZc2NZ6dJqFfhRoFNMR9Lnyj7dYGrzdgXXVMXcxoKTePPX1dWByq' ],
17    [ '64EEAB5F9BE2A01A8365A579511EB3373C87C40DA6D2A25F05BDA68FE077B66E',
18      unhexlify('cf9300f0909080f09f92a9'),
19      K_MAIN,
20      '6PRW5o9FLp4gJDDVqJQKJFTpMvdsSGJxMYHtHaQBF3ooa8mwD69bapcDQn' ],
21    [ 'CBF4B9F70470856BB4F40F80B87EDB90865997FFEE6DF315AB166D713AF433A5',
22      'TestingOneTwoThree',
23      K_MAIN + K_COMP,
24      '6PYNKZ1EAgYgmQfmNVamxyXVWHzK5s6DGhwP4J5o44cvXdoY7sRzhtpUeo' ],
25    [ '09C2686880095B1A4C249EE3AC4EEA8A014F11E6F986D0B5025AC1F39AFBD9AE',
26      'Satoshi',
27      K_MAIN + K_COMP,
28      '6PYLtMnXvfG3oJde97zRyLYFZCYizPU5T3LwgdYJz1fRhh16bU7u6PPmY7' ],
29    # Raw vectors
30    [ '09C2686880095B1A4C249EE3AC4EEA8A014F11E6F986D0B5025AC1F39AFBD9AE',
31      'Satoshi',
32      K_MAIN + K_COMP + K_RAW,
33      '0142E00B76EA60B62F66F0AF93D8B5380652AF51D1A3902EE00726CCEB70CA636B5B57CE6D3E2F' ],
34    [ '3CBC4D1E5C5248F81338596C0B1EE025FBE6C112633C357D66D2CE0BE541EA18',
35      'jon',
36      K_MAIN + K_COMP + K_RAW + K_ORDER,
37      '0142E09F8EE6E3A2FFCB13A99AA976AEDA5A2002ED3DF97FCB9957CD863357B55AA2072D3EB2F9' ],
38]
39
40
41class BIP38Tests(unittest.TestCase):
42
43    def from_priv(self, priv_key, passwd, flags):
44        priv, p_len = make_cbuffer(priv_key)
45        if flags > K_RAW:
46            out_buf, out_len = make_cbuffer('00' * 39)
47            ret = bip38_raw_from_private_key(priv, p_len, passwd, len(passwd),
48                                             flags, out_buf, out_len)
49            return ret, h(out_buf).upper()
50        else:
51            return bip38_from_private_key(priv, p_len, passwd, len(passwd), flags)
52
53    def to_priv(self, bip38, passwd, flags):
54        priv, priv_len = make_cbuffer('00' * 32)
55        bip38 = utf8(bip38)
56        if flags > K_RAW:
57            raw, raw_len = make_cbuffer(bip38)
58            ret = bip38_raw_to_private_key(raw, raw_len, passwd, len(passwd),
59                                           flags, priv, priv_len)
60        else:
61            ret = bip38_to_private_key(bip38, passwd, len(passwd), flags,
62                                       priv, priv_len)
63        return ret, priv
64
65
66    def test_bip38(self):
67
68        for case in cases:
69            priv_key, passwd, flags, expected = case
70            passwd = utf8(passwd) if type(passwd) is not bytes else passwd
71            ret, bip38 = self.from_priv(priv_key, passwd, flags)
72            self.assertEqual(ret, WALLY_OK)
73            bip38 = bip38.decode('utf-8') if type(bip38) is bytes else bip38
74            self.assertEqual(bip38, expected)
75
76            ret, new_priv_key = self.to_priv(bip38, passwd, flags)
77            self.assertEqual(ret, WALLY_OK)
78            self.assertEqual(h(new_priv_key).upper(), utf8(priv_key))
79            ret, new_priv_key = self.to_priv(bip38, '', flags + K_CHECK)
80            self.assertEqual(ret, WALLY_OK)
81
82
83    def test_bip38_invalid(self):
84        priv_key = 'CBF4B9F70470856BB4F40F80B87EDB90865997FFEE6DF315AB166D713AF433A5'
85        passwd = utf8('TestingInvalidFlags')
86        K_RES1 = 0x10 # BIP38_FLAG_RESERVED1
87
88        for flags, expected in [(0,            WALLY_OK),
89                                (K_RES1,       WALLY_EINVAL),
90                                (K_RAW,        WALLY_OK),
91                                (K_RAW+K_RES1, WALLY_EINVAL)]:
92            ret, _ = self.from_priv(priv_key, passwd, K_MAIN + flags)
93            self.assertEqual(ret, expected)
94
95    def test_bip38_flags(self):
96        priv_key = "6PYTh1Jgj3caimSrFjsfR5wJ8zUgWNDiPoNVZapSy8BwkF4NaKa1R32CaN"
97        ret, flags = bip38_get_flags(utf8(priv_key))
98        self.assertEqual(ret, WALLY_OK)
99        flags = int(flags)
100        self.assertEqual(flags & 256, 256) # BIP38_COMPRESSED
101        self.assertEqual(flags & 512, 0) # BIP38_KEY_EC_MULT
102
103if __name__ == '__main__':
104    unittest.main()
105