From 61d79e28991854548e5dc415c74da0a339278c1d Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Mon, 3 Oct 2022 09:59:58 +0000
Subject: [PATCH] proto.decode_wif(): parse version bytes more efficiently

---
 mmgen/altcoin.py   |  2 +-
 mmgen/proto/btc.py | 26 ++++++++------------------
 mmgen/proto/xmr.py |  1 -
 mmgen/proto/zec.py | 15 +++++++++++++--
 mmgen/protocol.py  |  6 ++++++
 5 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/mmgen/altcoin.py b/mmgen/altcoin.py
index 9181675c..7d4a75e9 100755
--- a/mmgen/altcoin.py
+++ b/mmgen/altcoin.py
@@ -454,7 +454,7 @@ class CoinInfo(object):
 					test_equal(
 						'WIF version number',
 						e.wif_ver_num,
-						int.from_bytes(bytes.fromhex(proto.wif_ver_num['std']),'big'),
+						int.from_bytes(proto.wif_ver_bytes['std'],'big'),
 						*cdata )
 
 					test_equal(
diff --git a/mmgen/proto/btc.py b/mmgen/proto/btc.py
index 80d6d2f0..6f8b7d4e 100755
--- a/mmgen/proto/btc.py
+++ b/mmgen/proto/btc.py
@@ -52,36 +52,26 @@ class mainnet(CoinProtocol.Secp256k1): # chainparams.cpp
 
 	def encode_wif(self,privbytes,pubkey_type,compressed): # input is preprocessed hex
 		assert len(privbytes) == self.privkey_len, f'{len(privbytes)} bytes: incorrect private key length!'
-		assert pubkey_type in self.wif_ver_num, f'{pubkey_type!r}: invalid pubkey_type'
+		assert pubkey_type in self.wif_ver_bytes, f'{pubkey_type!r}: invalid pubkey_type'
 		return b58chk_encode(
-			bytes.fromhex(self.wif_ver_num[pubkey_type])
+			self.wif_ver_bytes[pubkey_type]
 			+ privbytes
 			+ (b'',b'\x01')[bool(compressed)])
 
 	def decode_wif(self,wif):
-		key = b58chk_decode(wif)
-
-		for k,v in self.wif_ver_num.items():
-			v = bytes.fromhex(v)
-			if key[:len(v)] == v:
-				pubkey_type = k
-				key = key[len(v):]
-				break
-		else:
-			raise ValueError('Invalid WIF version number')
+		key_data = b58chk_decode(wif)
+		vlen = self.wif_ver_bytes_len or self.get_wif_ver_bytes_len(key_data)
+		key = key_data[vlen:]
 
 		if len(key) == self.privkey_len + 1:
 			assert key[-1] == 0x01, f'{key[-1]!r}: invalid compressed key suffix byte'
-			compressed = True
-		elif len(key) == self.privkey_len:
-			compressed = False
-		else:
+		elif len(key) != self.privkey_len:
 			raise ValueError(f'{len(key)}: invalid key length')
 
 		return decoded_wif(
 			sec         = key[:self.privkey_len],
-			pubkey_type = pubkey_type,
-			compressed  = compressed )
+			pubkey_type = self.wif_ver_bytes_to_pubkey_type[key_data[:vlen]],
+			compressed  = len(key) == self.privkey_len + 1 )
 
 	def decode_addr(self,addr):
 
diff --git a/mmgen/proto/xmr.py b/mmgen/proto/xmr.py
index 48e3c8a1..cb54cb12 100755
--- a/mmgen/proto/xmr.py
+++ b/mmgen/proto/xmr.py
@@ -25,7 +25,6 @@ class mainnet(CoinProtocol.DummyWIF,CoinProtocol.Base):
 	base_coin      = 'XMR'
 	base_proto     = 'Monero'
 	addr_ver_info  = { '12': 'monero', '2a': 'monero_sub', '13': 'monero_integrated' }
-	wif_ver_num    = {}
 	pubkey_types   = ('monero',)
 	mmtypes        = ('M',)
 	dfl_mmtype     = 'M'
diff --git a/mmgen/proto/zec.py b/mmgen/proto/zec.py
index 11a5efd1..220a2585 100755
--- a/mmgen/proto/zec.py
+++ b/mmgen/proto/zec.py
@@ -13,7 +13,8 @@ Zcash protocol
 """
 
 from .btc import mainnet
-from ..protocol import decoded_addr
+from .common import b58chk_decode
+from ..protocol import decoded_wif,decoded_addr
 
 class mainnet(mainnet):
 	base_coin      = 'ZEC'
@@ -30,12 +31,22 @@ class mainnet(mainnet):
 		from ..opts import opt
 		self.coin_id = 'ZEC-Z' if opt.type in ('zcash_z','Z') else 'ZEC-T'
 
+	def get_wif_ver_bytes_len(self,key_data):
+		"""
+		vlen must be set dynamically since Zcash has variable-length version bytes
+		"""
+		for v in self.wif_ver_bytes.values():
+			if key_data[:len(v)] == v:
+				return len(v)
+		else:
+			raise ValueError('Invalid WIF version number')
+
 	def get_addr_len(self,addr_fmt):
 		return (20,64)[addr_fmt in ('zcash_z','viewkey')]
 
 	def decode_addr_bytes(self,addr_bytes):
 		"""
-		vlen must be set dynamically since Zcash has variable length ver_bytes
+		vlen must be set dynamically since Zcash has variable-length version bytes
 		"""
 		for ver_bytes,addr_fmt in self.addr_ver_bytes.items():
 			vlen = len(ver_bytes)
diff --git a/mmgen/protocol.py b/mmgen/protocol.py
index da0ab19a..b8508c10 100755
--- a/mmgen/protocol.py
+++ b/mmgen/protocol.py
@@ -68,6 +68,12 @@ class CoinProtocol(MMGenObject):
 				'regtest': '_rt',
 			}[network]
 
+			if hasattr(self,'wif_ver_num'):
+				self.wif_ver_bytes = {k:bytes.fromhex(v) for k,v in self.wif_ver_num.items()}
+				self.wif_ver_bytes_to_pubkey_type = {v:k for k,v in self.wif_ver_bytes.items()}
+				vbs = list(self.wif_ver_bytes.values())
+				self.wif_ver_bytes_len = len(vbs[0]) if len(set(len(b) for b in vbs)) == 1 else None
+
 			if hasattr(self,'addr_ver_info'):
 				self.addr_ver_bytes = {bytes.fromhex(k):v for k,v in self.addr_ver_info.items()}
 				self.addr_fmt_to_ver_bytes = {v:k for k,v in self.addr_ver_bytes.items()}