diff --git a/pybtc/hdwallet.py b/pybtc/hdwallet.py index 54d0383..e356189 100644 --- a/pybtc/hdwallet.py +++ b/pybtc/hdwallet.py @@ -11,9 +11,11 @@ from .constants import * # def create_passphrase(bits=256, language='english'): - entropy = os.urandom(bits // 8) - mnemonic = create_mnemonic(entropy, language) - return ' '.join(mnemonic[::-1]) + if bits in [128, 160, 192, 224, 256]: + entropy = os.urandom(bits // 8) + mnemonic = create_mnemonic(entropy, language) + return ' '.join(mnemonic[::-1]) + raise ValueError('Strength should be one of the following [128, 160, 192, 224, 256], but it is not (%d).' % bits) def create_mnemonic(entropy, language='english'): @@ -57,22 +59,24 @@ def add_checksum(data): def mnemonic2bytes(passphrase, language): mnemonic = passphrase.split() - wordlist = create_wordlist(language) - codes = dict() - for code, word in enumerate(wordlist): - codes[word] = code - word_count = len(mnemonic) - entropy_int = None - bit_size = word_count * 11 - chk_sum_bit_len = word_count * 11 % 32 - for word in mnemonic: - entropy_int = (entropy_int << 11) | codes[word] if entropy_int else codes[word] - chk_sum = entropy_int & (2 ** chk_sum_bit_len - 1) - entropy_int = entropy_int >> chk_sum_bit_len - entropy = entropy_int.to_bytes((bit_size - chk_sum_bit_len) // 8, byteorder="big") - ent_hash = hashlib.sha256(entropy).hexdigest() - fb = unhexlify(ent_hash)[0] - assert (fb >> (8 - chk_sum_bit_len)) & chk_sum - return entropy + if len(mnemonic) in [12, 15, 18, 21, 24]: + wordlist = create_wordlist(language) + codes = dict() + for code, word in enumerate(wordlist): + codes[word] = code + word_count = len(mnemonic) + entropy_int = None + bit_size = word_count * 11 + chk_sum_bit_len = word_count * 11 % 32 + for word in mnemonic: + entropy_int = (entropy_int << 11) | codes[word] if entropy_int else codes[word] + chk_sum = entropy_int & (2 ** chk_sum_bit_len - 1) + entropy_int = entropy_int >> chk_sum_bit_len + entropy = entropy_int.to_bytes((bit_size - chk_sum_bit_len) // 8, byteorder="big") + ent_hash = hashlib.sha256(entropy).hexdigest() + fb = unhexlify(ent_hash)[0] + assert (fb >> (8 - chk_sum_bit_len)) & chk_sum + return entropy + raise ValueError('Number of words must be one of the following: [12, 15, 18, 21, 24], but it is not (%d).' % len(mnemonic))