From f63b35469c2fa2e5ddd8018947976bc813ca050f Mon Sep 17 00:00:00 2001 From: The MMGen Project Date: Sat, 2 Nov 2019 10:31:31 +0000 Subject: [PATCH] parse_wif(): additional cleanups --- mmgen/protocol.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/mmgen/protocol.py b/mmgen/protocol.py index b3049783..5f428fb8 100755 --- a/mmgen/protocol.py +++ b/mmgen/protocol.py @@ -136,21 +136,26 @@ class BitcoinProtocol(MMGenObject): @classmethod def parse_wif(cls,wif): - key = _b58chk_decode(wif) - pubkey_type = None - for k,v in list(cls.wif_ver_num.items()): + key = bytes.fromhex(_b58chk_decode(wif)) + + for k,v in cls.wif_ver_num.items(): + v = bytes.fromhex(v) if key[:len(v)] == v: pubkey_type = k key = key[len(v):] - assert pubkey_type,'invalid WIF version number' - if len(key) == 66: - assert key[-2:] == '01','invalid compressed key suffix' - compressed = True + break else: - assert len(key) == 64,'invalid key length' - compressed = False + raise ValueError('invalid WIF version number') - return parsed_wif(bytes.fromhex(key[:64]), pubkey_type, compressed) + if len(key) == cls.privkey_len + 1: + assert key[-1] == 0x01,'{!r}: invalid compressed key suffix byte'.format(key[-1]) + compressed = True + elif len(key) == cls.privkey_len: + compressed = False + else: + raise ValueError('{}: invalid key length'.format(len(key))) + + return parsed_wif(key[:cls.privkey_len], pubkey_type, compressed) @classmethod def verify_addr(cls,addr,hex_width,return_dict=False):