chia-blockchain/tests/core/util/test_keychain.py

131 lines
6.4 KiB
Python

import json
import unittest
from secrets import token_bytes
from blspy import AugSchemeMPL, PrivateKey
from chia.util.keychain import Keychain, bytes_from_mnemonic, bytes_to_mnemonic, generate_mnemonic, mnemonic_to_seed
class TesKeychain(unittest.TestCase):
def test_basic_add_delete(self):
kc: Keychain = Keychain(testing=True)
kc.delete_all_keys()
assert kc._get_free_private_key_index() == 0
assert len(kc.get_all_private_keys()) == 0
assert kc.get_first_private_key() is None
assert kc.get_first_public_key() is None
mnemonic = generate_mnemonic()
entropy = bytes_from_mnemonic(mnemonic)
assert bytes_to_mnemonic(entropy) == mnemonic
mnemonic_2 = generate_mnemonic()
# misspelled words in the mnemonic
bad_mnemonic = mnemonic.split(" ")
bad_mnemonic[6] = "ZZZZZZ"
self.assertRaisesRegex(
ValueError,
"'ZZZZZZ' is not in the mnemonic dictionary; may be misspelled",
bytes_from_mnemonic,
" ".join(bad_mnemonic),
)
kc.add_private_key(mnemonic, "")
assert kc._get_free_private_key_index() == 1
assert len(kc.get_all_private_keys()) == 1
kc.add_private_key(mnemonic_2, "")
kc.add_private_key(mnemonic_2, "") # checks to not add duplicates
assert kc._get_free_private_key_index() == 2
assert len(kc.get_all_private_keys()) == 2
assert kc._get_free_private_key_index() == 2
assert len(kc.get_all_private_keys()) == 2
assert len(kc.get_all_public_keys()) == 2
assert kc.get_all_private_keys()[0] == kc.get_first_private_key()
assert kc.get_all_public_keys()[0] == kc.get_first_public_key()
assert len(kc.get_all_private_keys()) == 2
seed_2 = mnemonic_to_seed(mnemonic, "")
seed_key_2 = AugSchemeMPL.key_gen(seed_2)
kc.delete_key_by_fingerprint(seed_key_2.get_g1().get_fingerprint())
assert kc._get_free_private_key_index() == 0
assert len(kc.get_all_private_keys()) == 1
kc.delete_all_keys()
assert kc._get_free_private_key_index() == 0
assert len(kc.get_all_private_keys()) == 0
kc.add_private_key(bytes_to_mnemonic(token_bytes(32)), "my passphrase")
kc.add_private_key(bytes_to_mnemonic(token_bytes(32)), "")
kc.add_private_key(bytes_to_mnemonic(token_bytes(32)), "third passphrase")
assert len(kc.get_all_public_keys()) == 3
assert len(kc.get_all_private_keys()) == 1
assert len(kc.get_all_private_keys(["my passphrase", ""])) == 2
assert len(kc.get_all_private_keys(["my passphrase", "", "third passphrase", "another"])) == 3
assert len(kc.get_all_private_keys(["my passhrase wrong"])) == 0
assert kc.get_first_private_key() is not None
assert kc.get_first_private_key(["bad passphrase"]) is None
assert kc.get_first_public_key() is not None
kc.delete_all_keys()
kc.add_private_key(bytes_to_mnemonic(token_bytes(32)), "my passphrase")
assert kc.get_first_public_key() is not None
def test_bip39_eip2333_test_vector(self):
kc: Keychain = Keychain(testing=True)
kc.delete_all_keys()
mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
passphrase = "TREZOR"
print("entropy to seed:", mnemonic_to_seed(mnemonic, passphrase).hex())
master_sk = kc.add_private_key(mnemonic, passphrase)
tv_master_int = 5399117110774477986698372024995405256382522670366369834617409486544348441851
tv_child_int = 11812940737387919040225825939013910852517748782307378293770044673328955938106
assert master_sk == PrivateKey.from_bytes(tv_master_int.to_bytes(32, "big"))
child_sk = AugSchemeMPL.derive_child_sk(master_sk, 0)
assert child_sk == PrivateKey.from_bytes(tv_child_int.to_bytes(32, "big"))
def test_bip39_test_vectors_trezor(self):
with open("tests/util/bip39_test_vectors.json") as f:
all_vectors = json.loads(f.read())
for vector_list in all_vectors["english"]:
entropy_bytes = bytes.fromhex(vector_list[0])
mnemonic = vector_list[1]
seed = bytes.fromhex(vector_list[2])
assert bytes_from_mnemonic(mnemonic) == entropy_bytes
assert bytes_to_mnemonic(entropy_bytes) == mnemonic
assert mnemonic_to_seed(mnemonic, "TREZOR") == seed
def test_utf8_nfkd(self):
# Test code from trezor:
# Copyright (c) 2013 Pavol Rusnak
# Copyright (c) 2017 mruddy
# https://github.com/trezor/python-mnemonic/blob/master/test_mnemonic.py
# The same sentence in various UTF-8 forms
words_nfkd = "Pr\u030ci\u0301s\u030cerne\u030c z\u030clut\u030couc\u030cky\u0301 ku\u030an\u030c u\u0301pe\u030cl d\u030ca\u0301belske\u0301 o\u0301dy za\u0301ker\u030cny\u0301 uc\u030cen\u030c be\u030cz\u030ci\u0301 pode\u0301l zo\u0301ny u\u0301lu\u030a" # noqa: E501
words_nfc = "P\u0159\xed\u0161ern\u011b \u017elu\u0165ou\u010dk\xfd k\u016f\u0148 \xfap\u011bl \u010f\xe1belsk\xe9 \xf3dy z\xe1ke\u0159n\xfd u\u010de\u0148 b\u011b\u017e\xed pod\xe9l z\xf3ny \xfal\u016f" # noqa: E501
words_nfkc = "P\u0159\xed\u0161ern\u011b \u017elu\u0165ou\u010dk\xfd k\u016f\u0148 \xfap\u011bl \u010f\xe1belsk\xe9 \xf3dy z\xe1ke\u0159n\xfd u\u010de\u0148 b\u011b\u017e\xed pod\xe9l z\xf3ny \xfal\u016f" # noqa: E501
words_nfd = "Pr\u030ci\u0301s\u030cerne\u030c z\u030clut\u030couc\u030cky\u0301 ku\u030an\u030c u\u0301pe\u030cl d\u030ca\u0301belske\u0301 o\u0301dy za\u0301ker\u030cny\u0301 uc\u030cen\u030c be\u030cz\u030ci\u0301 pode\u0301l zo\u0301ny u\u0301lu\u030a" # noqa: E501
passphrase_nfkd = "Neuve\u030cr\u030citelne\u030c bezpec\u030cne\u0301 hesli\u0301c\u030cko"
passphrase_nfc = "Neuv\u011b\u0159iteln\u011b bezpe\u010dn\xe9 hesl\xed\u010dko"
passphrase_nfkc = "Neuv\u011b\u0159iteln\u011b bezpe\u010dn\xe9 hesl\xed\u010dko"
passphrase_nfd = "Neuve\u030cr\u030citelne\u030c bezpec\u030cne\u0301 hesli\u0301c\u030cko"
seed_nfkd = mnemonic_to_seed(words_nfkd, passphrase_nfkd)
seed_nfc = mnemonic_to_seed(words_nfc, passphrase_nfc)
seed_nfkc = mnemonic_to_seed(words_nfkc, passphrase_nfkc)
seed_nfd = mnemonic_to_seed(words_nfd, passphrase_nfd)
assert seed_nfkd == seed_nfc
assert seed_nfkd == seed_nfkc
assert seed_nfkd == seed_nfd