1import unittest
2
3from pycoin.encoding.b58 import a2b_hashed_base58, b2a_hashed_base58
4from pycoin.networks.registry import network_for_netcode
5from pycoin.networks.registry import network_codes
6from pycoin.symbols.btc import network as BTC
7from pycoin.symbols.xtn import network as XTN
8
9NETCODES = "BTC XTN DOGE".split()
10
11
12def change_prefix(address, new_prefix):
13    return b2a_hashed_base58(new_prefix + a2b_hashed_base58(address)[1:])
14
15
16P2PKH_ADDRESSES = [
17    "1BgGZ9tcN4rm9KBzDn7KprQz87SZ26SAMH", "1EHNa6Q4Jz2uvNExL497mE43ikXhwF6kZm",
18    "1cMh228HTCiwS8ZsaakH8A8wze1JR5ZsP", "1LagHJk2FyCV2VzrNHVqg3gYG4TSYwDV4m",
19    "1CUNEBjYrCn2y1SdiUMohaKUi4wpP326Lb", "1NZUP3JAc9JkmbvmoTv7nVgZGtyJjirKV1"]
20
21P2SH_ADDRESSES = [
22    '3CNHUhP3uyB9EUtRLsmvFUmvGdjGdkTxJw', '3EyPVdtVrtMJ1XwPT9oiBrQysGpRY8LE9K',
23    '32JNcZWZqMX72bpzzgFLhkX56WviowgUtS', '3MGhCrETosWs7fhHVPAS6g3UQakA7Xz3wb',
24    '3DAP9jDzQ76R4B94qa2Q8CgQrbEXvUoghh', '3PFVJancA3d8rmdCvZaiD83VRRG2Em15Ge']
25
26
27class KeyUtilsTest(unittest.TestCase):
28
29    def test_address_valid_btc(self):
30        for address in P2PKH_ADDRESSES:
31            self.assertEqual(BTC.parse.p2pkh(address).address(), address)
32            a = address[:-1] + chr(ord(address[-1])+1)
33            self.assertIsNone(BTC.parse.address(a))
34
35        for address in P2PKH_ADDRESSES:
36            self.assertIsNone(BTC.parse.p2sh(address))
37            self.assertEqual(BTC.parse.p2pkh(address).address(), address)
38
39        for address in P2SH_ADDRESSES:
40            self.assertEqual(address[0], "3")
41            self.assertEqual(BTC.parse.p2sh(address).address(), address)
42            self.assertIsNone(BTC.parse.p2pkh(address))
43
44    def test_is_wif_valid(self):
45        WIFS = ["KwDiBf89QgGbjEhKnhXJuH7LrciVrZi3qYjgd9M7rFU73sVHnoWn",
46                "5HpHagT65TZzG1PH3CSu63k8DbpvD8s5ip4nEB3kEsreAnchuDf",
47                "KwDiBf89QgGbjEhKnhXJuH7LrciVrZi3qYjgd9M7rFU74NMTptX4",
48                "5HpHagT65TZzG1PH3CSu63k8DbpvD8s5ip4nEB3kEsreAvUcVfH"]
49
50        for wif in WIFS:
51            self.assertEqual(BTC.parse.wif(wif).wif(), wif)
52            a = wif[:-1] + chr(ord(wif[-1])+1)
53            self.assertIsNone(BTC.parse.wif(a))
54
55        NETWORK_NAMES = network_codes()
56        for netcode in NETWORK_NAMES:
57            network = network_for_netcode(netcode)
58            if not getattr(network, "Key", None):
59                continue
60            for se in range(1, 10):
61                key = network.Key(secret_exponent=se)
62                for tv in [True, False]:
63                    wif = key.wif(is_compressed=tv)
64                    self.assertEqual(network.parse.wif(wif).wif(), wif)
65                    a = wif[:-1] + chr(ord(wif[-1])+1)
66                    self.assertIsNone(network.parse.wif(a))
67
68    def test_is_public_private_bip32_valid(self):
69        from pycoin.networks.registry import network_for_netcode
70        WALLET_KEYS = ["foo", "1", "2", "3", "4", "5"]
71
72        # not all networks support BIP32 yet
73        for netcode in NETCODES:
74            network = network_for_netcode(netcode)
75            for wk in WALLET_KEYS:
76                wallet = network.keys.bip32_seed(wk.encode("utf8"))
77                text = wallet.hwif(as_private=True)
78                self.assertEqual(network.parse.bip32_prv(text).hwif(as_private=True), text)
79                self.assertIsNone(network.parse.bip32_pub(text))
80                a = text[:-1] + chr(ord(text[-1])+1)
81                self.assertIsNone(network.parse.bip32_prv(a))
82                self.assertIsNone(network.parse.bip32_pub(a))
83                text = wallet.hwif(as_private=False)
84                self.assertIsNone(network.parse.bip32_prv(text))
85                self.assertEqual(network.parse.bip32_pub(text).hwif(), text)
86                a = text[:-1] + chr(ord(text[-1])+1)
87                self.assertIsNone(network.parse.bip32_prv(a))
88                self.assertIsNone(network.parse.bip32_pub(a))
89
90    def test_key_limits(self):
91        nc = 'BTC'
92        cc = b'000102030405060708090a0b0c0d0e0f'
93        order = BTC.keys.private(1)._generator.order()
94
95        # BRAIN DAMAGE: hack
96        BIP32Node = BTC.keys.bip32_seed(b"foo").__class__
97        for k in -1, 0, order, order + 1:
98            self.assertRaises(BTC.keys.InvalidSecretExponentError, BTC.keys.private, secret_exponent=k)
99            self.assertRaises(BTC.keys.InvalidSecretExponentError, BIP32Node, nc, cc, secret_exponent=k)
100
101        for i in range(1, 512):
102            BTC.keys.private(secret_exponent=i)
103            BIP32Node(cc, secret_exponent=i)
104
105    def test_repr(self):
106        key = XTN.keys.private(secret_exponent=273)
107
108        address = key.address()
109        pub_k = XTN.parse(address)
110        self.assertEqual(repr(pub_k),  '<mhDVBkZBWLtJkpbszdjZRkH1o5RZxMwxca>')
111
112        wif = key.wif()
113        priv_k = XTN.parse.wif(wif)
114        self.assertEqual(
115            repr(priv_k),
116            'private_for <XTNSEC:0264e1b1969f9102977691a40431b0b672055dcf31163897d996434420e6c95dc9>')
117
118
119if __name__ == '__main__':
120    unittest.main()
121