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