From f64be2b88355269857727a4290562f91b1e328dd Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Wed, 6 Oct 2021 13:22:33 +0000
Subject: [PATCH] txsign: support all address types with flat keylist

---
 mmgen/addr.py        | 55 ++++++++++++++++++++------------------------
 mmgen/main_txdo.py   |  2 --
 mmgen/main_txsign.py |  2 --
 mmgen/txsign.py      | 11 +++------
 4 files changed, 28 insertions(+), 42 deletions(-)

diff --git a/mmgen/addr.py b/mmgen/addr.py
index af36f085..8db737a0 100755
--- a/mmgen/addr.py
+++ b/mmgen/addr.py
@@ -596,39 +596,34 @@ Removed {{}} duplicate WIF key{{}} from keylist (also in {pnm} key-address file
 				pass
 		return d
 
-	def remove_dup_keys(self,cmplist):
-		assert self.has_keys
-		pop_list = []
-		for n,d in enumerate(self.data):
-			for e in cmplist.data:
-				if e.sec.wif == d.sec.wif:
-					pop_list.append(n)
-		for n in reversed(pop_list): self.data.pop(n)
-		if pop_list:
-			vmsg(self.msgs['removed_dup_keys'].format(len(pop_list),suf(removed)))
-
 	def add_wifs(self,key_list):
-		if not key_list: return
+		"""
+		Match WIF keys in a flat list to addresses in self by generating all
+		possible addresses for each key.
+		"""
+		def gen_addr(pk,t):
+			at = self.proto.addr_type(t)
+			kg = KeyGenerator(self.proto,at.pubkey_type)
+			ag = AddrGenerator(self.proto,at)
+			return ag.to_addr(kg.to_pubhex(pk))
+
+		compressed_types = set(self.proto.mmtypes) - {'L','E'}
+		uncompressed_types = set(self.proto.mmtypes) & {'L','E'}
+
+		def gen():
+			for wif in key_list:
+				pk = PrivKey(proto=self.proto,wif=wif)
+				for t in (compressed_types if pk.compressed else uncompressed_types):
+					yield ( gen_addr(pk,t), pk )
+
+		addrs4keys = dict(gen())
+
 		for d in self.data:
-			for e in key_list.data:
-				if e.addr and e.sec and e.addr == d.addr:
-					d.sec = e.sec
+			if d.addr in addrs4keys:
+				d.sec = addrs4keys[d.addr]
 
-	def list_missing(self,key):
-		return [d.addr for d in self.data if not getattr(d,key)]
-
-	def generate_addrs_from_keys(self):
-		# assume that the first listed mmtype is valid for flat key list
-		at = self.proto.addr_type(self.proto.mmtypes[0])
-		kg = KeyGenerator(self.proto,at.pubkey_type)
-		ag = AddrGenerator(self.proto,at)
-		d = self.data
-		for n,e in enumerate(d,1):
-			qmsg_r(f'\rGenerating addresses from keylist: {n}/{len(d)}')
-			e.addr = ag.to_addr(kg.to_pubhex(e.sec))
-			if g.debug_addrlist:
-				Msg(f'generate_addrs_from_keys():\n{e.pfmt()}')
-		qmsg(f'\rGenerated addresses from keylist: {n}/{len(d)} ')
+	def list_missing(self,attr):
+		return [d.addr for d in self.data if not getattr(d,attr)]
 
 	def make_label(self):
 		bc,mt = self.proto.base_coin,self.al_id.mmtype
diff --git a/mmgen/main_txdo.py b/mmgen/main_txdo.py
index 390a86ff..7997d6f0 100755
--- a/mmgen/main_txdo.py
+++ b/mmgen/main_txdo.py
@@ -139,8 +139,6 @@ async def main():
 
 	kal = get_keyaddrlist(proto,opt)
 	kl = get_keylist(proto,opt)
-	if kl and kal:
-		kl.remove_dup_keys(kal)
 
 	tx3 = await txsign(tx2,seed_files,kl,kal)
 
diff --git a/mmgen/main_txsign.py b/mmgen/main_txsign.py
index 147c3e3f..b7da9731 100755
--- a/mmgen/main_txsign.py
+++ b/mmgen/main_txsign.py
@@ -140,8 +140,6 @@ async def main():
 
 		kal = get_keyaddrlist(tx1.proto,opt)
 		kl = get_keylist(tx1.proto,opt)
-		if kl and kal:
-			kl.remove_dup_keys(kal)
 
 		tx2 = await txsign(tx1,seed_files,kl,kal,tx_num_disp)
 		if tx2:
diff --git a/mmgen/txsign.py b/mmgen/txsign.py
index c5dfb6b8..5708a6ba 100755
--- a/mmgen/txsign.py
+++ b/mmgen/txsign.py
@@ -133,13 +133,7 @@ def get_keyaddrlist(proto,opt):
 
 def get_keylist(proto,opt):
 	if opt.keys_from_file:
-		l = get_lines_from_file(opt.keys_from_file,'key-address data',trim_comments=True)
-		kal = KeyAddrList(
-			proto     = proto,
-			keylist   = [m.split()[0] for m in l], # accept coin daemon wallet dumps
-			skip_chksum = True )
-		kal.generate_addrs_from_keys()
-		return kal
+		return get_lines_from_file(opt.keys_from_file,'key-address data',trim_comments=True)
 	return None
 
 async def txsign(tx,seed_files,kl,kal,tx_num_str=''):
@@ -153,7 +147,8 @@ async def txsign(tx,seed_files,kl,kal,tx_num_str=''):
 			proto = tx.proto,
 			addrlist = non_mmaddrs,
 			skip_chksum = True )
-		tmp.add_wifs(kl)
+		if kl:
+			tmp.add_wifs(kl)
 		missing = tmp.list_missing('sec')
 		if missing:
 			sep = '\n    '