From 99c2238c63b4da7e8e4e612de13dd0ebf3078342 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Mon, 4 Mar 2024 10:30:37 +0000
Subject: [PATCH] TwCtl: `resolve_address()`, `get_label_addr_pairs()`:
 cleanups

---
 mmgen/data/version              |  2 +-
 mmgen/proto/btc/tw/addresses.py | 23 ++++++----------
 mmgen/proto/btc/tw/ctl.py       |  4 +--
 mmgen/proto/btc/tw/rpc.py       | 12 +++-----
 mmgen/proto/btc/tw/txhistory.py |  4 +--
 mmgen/proto/eth/tw/addresses.py | 10 +++----
 mmgen/proto/eth/tw/ctl.py       |  9 +++---
 mmgen/proto/eth/tw/rpc.py       |  9 +-----
 mmgen/tw/ctl.py                 | 49 ++++++++++++++-------------------
 mmgen/tw/view.py                |  6 +++-
 test/cmdtest_py_d/ct_regtest.py |  4 +--
 11 files changed, 55 insertions(+), 77 deletions(-)

diff --git a/mmgen/data/version b/mmgen/data/version
index e469c7b2..59b4bb6c 100644
--- a/mmgen/data/version
+++ b/mmgen/data/version
@@ -1 +1 @@
-14.1.dev16
+14.1.dev17
diff --git a/mmgen/proto/btc/tw/addresses.py b/mmgen/proto/btc/tw/addresses.py
index efe8839d..aca0166c 100755
--- a/mmgen/proto/btc/tw/addresses.py
+++ b/mmgen/proto/btc/tw/addresses.py
@@ -53,20 +53,15 @@ Actions: [q]uit menu, r[e]draw, add [l]abel:
 		self.total = sum((v['amt'] for v in addrs.values()), start=amt0)
 
 		msg_r('Getting labels and associated addresses...')
-		pairs = await self.get_label_addr_pairs()
-
-		if pairs:
-			for label,addr in pairs:
-				if label and label.mmid not in addrs:
-					addrs[label.mmid] = {
-						'addr':   addr,
-						'amt':    amt0,
-						'recvd':  amt0,
-						'confs':  0,
-						'lbl':    label }
-			msg('done')
-		else:
-			msg('[none]')
+		for e in await self.get_label_addr_pairs():
+			if e.label and e.label.mmid not in addrs:
+				addrs[e.label.mmid] = {
+					'addr':   e.coinaddr,
+					'amt':    amt0,
+					'recvd':  amt0,
+					'confs':  0,
+					'lbl':    e.label}
+		msg('done')
 
 		msg_r('Getting received funds data...')
 		# args: 1:minconf, 2:include_empty, 3:include_watchonly, 4:include_immature_coinbase (>=v23.0.0)
diff --git a/mmgen/proto/btc/tw/ctl.py b/mmgen/proto/btc/tw/ctl.py
index f0f9713e..72ce1538 100755
--- a/mmgen/proto/btc/tw/ctl.py
+++ b/mmgen/proto/btc/tw/ctl.py
@@ -125,6 +125,6 @@ class BitcoinTwCtl(TwCtl):
 				'Addresses have no balances' )
 			return True
 
-	async def get_label_addr_pairs(self, twmmid):
+	async def get_label_addr_pairs(self):
 		from .rpc import TwRPC
-		return await TwRPC(proto=self.proto, rpc=self.rpc, twctl=self).get_label_addr_pairs(twmmid)
+		return await TwRPC(proto=self.proto, rpc=self.rpc, twctl=self).get_label_addr_pairs()
diff --git a/mmgen/proto/btc/tw/rpc.py b/mmgen/proto/btc/tw/rpc.py
index 4565395d..e8fabe76 100755
--- a/mmgen/proto/btc/tw/rpc.py
+++ b/mmgen/proto/btc/tw/rpc.py
@@ -16,10 +16,11 @@ from ....addr import CoinAddr
 from ....util import die,msg,rmsg
 from ....tw.shared import get_tw_label
 from ....tw.rpc import TwRPC
+from ....tw.ctl import label_addr_pair
 
 class BitcoinTwRPC(TwRPC):
 
-	async def get_label_addr_pairs(self, twmmid=None):
+	async def get_label_addr_pairs(self):
 		"""
 		Get all the accounts in the tracking wallet and their associated addresses.
 		Returns list of (label,address) tuples.
@@ -49,9 +50,6 @@ class BitcoinTwRPC(TwRPC):
 
 		acct_labels = [get_tw_label(self.proto,a) for a in await get_acct_list()]
 
-		if twmmid:
-			acct_labels = [lbl for lbl in acct_labels if lbl.mmid == twmmid]
-
 		if not acct_labels:
 			return None
 
@@ -63,10 +61,8 @@ class BitcoinTwRPC(TwRPC):
 			if len(a) != 1:
 				raise ValueError(f'{a}: label {acct_labels[n]!r} has != 1 associated address!')
 
-		return [(
-			label,
-			CoinAddr(self.proto,addrs[0])
-		) for label,addrs in zip(acct_labels,acct_addrs)]
+		return [label_addr_pair(label, CoinAddr(self.proto,addrs[0]))
+			for label, addrs in zip(acct_labels, acct_addrs)]
 
 	async def get_unspent_by_mmid(self,minconf=1,mmid_filter=[]):
 		"""
diff --git a/mmgen/proto/btc/tw/txhistory.py b/mmgen/proto/btc/tw/txhistory.py
index 5009ba83..1c0b5c03 100755
--- a/mmgen/proto/btc/tw/txhistory.py
+++ b/mmgen/proto/btc/tw/txhistory.py
@@ -311,8 +311,8 @@ Filters/Actions: show [u]nconfirmed, [q]uit menu, r[e]draw:
 
 		if self.sinceblock: # mapping data may be incomplete for inputs, so update from 'listlabels'
 			mm_map.update(
-				{ addr: _mmp(label.mmid, label.comment) if label else _mmp(None,None)
-					for label, addr in await self.get_label_addr_pairs() }
+				{e.coinaddr: _mmp(e.label.mmid, e.label.comment) if e.label else _mmp(None, None)
+					for e in await self.get_label_addr_pairs()}
 			)
 
 		msg_r('Getting wallet transactions...')
diff --git a/mmgen/proto/eth/tw/addresses.py b/mmgen/proto/eth/tw/addresses.py
index e18f3f4e..e0baac3e 100755
--- a/mmgen/proto/eth/tw/addresses.py
+++ b/mmgen/proto/eth/tw/addresses.py
@@ -71,14 +71,14 @@ Actions: [q]uit menu, r[e]draw, [D]elete addr, add [l]abel:
 		self.minconf = None
 		addrs = {}
 
-		for label, addr in await self.get_label_addr_pairs():
-			bal = await self.twctl.get_balance(addr)
-			addrs[label.mmid] = {
-				'addr':  addr,
+		for e in await self.twctl.get_label_addr_pairs():
+			bal = await self.twctl.get_balance(e.coinaddr)
+			addrs[e.label.mmid] = {
+				'addr':  e.coinaddr,
 				'amt':   bal,
 				'recvd': amt0,
 				'confs': 0,
-				'lbl':   label }
+				'lbl':   e.label}
 			self.total += bal
 
 		return addrs
diff --git a/mmgen/proto/eth/tw/ctl.py b/mmgen/proto/eth/tw/ctl.py
index 87ea894e..2ae2ca65 100755
--- a/mmgen/proto/eth/tw/ctl.py
+++ b/mmgen/proto/eth/tw/ctl.py
@@ -21,7 +21,7 @@ proto.eth.tw.ctl: Ethereum tracking wallet control class
 """
 
 from ....util import msg,ymsg,die
-from ....tw.ctl import TwCtl,write_mode
+from ....tw.ctl import TwCtl, write_mode, label_addr_pair
 from ....tw.shared import TwLabel
 from ....addr import is_coin_addr,is_mmgen_id,CoinAddr
 from ..contract import Token,ResolvedToken
@@ -166,12 +166,11 @@ class EthereumTwCtl(TwCtl):
 	def mmid_ordered_dict(self):
 		return dict((x['mmid'],{'addr':x['addr'],'comment':x['comment']}) for x in self.sorted_list)
 
-	async def get_label_addr_pairs(self, twmmid):
-		ret = [(
-				TwLabel(self.proto, mmid + ' ' + d['comment']),
+	async def get_label_addr_pairs(self):
+		return [label_addr_pair(
+				TwLabel(self.proto, f"{mmid} {d['comment']}"),
 				CoinAddr(self.proto, d['addr'])
 			) for mmid, d in self.mmid_ordered_dict.items()]
-		return [e for e in ret if e[0].mmid == twmmid] or None
 
 class EthereumTokenTwCtl(EthereumTwCtl):
 
diff --git a/mmgen/proto/eth/tw/rpc.py b/mmgen/proto/eth/tw/rpc.py
index 2d22f1d0..1c952baf 100755
--- a/mmgen/proto/eth/tw/rpc.py
+++ b/mmgen/proto/eth/tw/rpc.py
@@ -12,17 +12,10 @@
 proto.eth.tw.rpc: Ethereum base protocol tracking wallet RPC class
 """
 
-from ....addr import CoinAddr
-from ....tw.shared import TwLabel
 from ....tw.rpc import TwRPC
 
 class EthereumTwRPC(TwRPC):
-
-	async def get_label_addr_pairs(self):
-		return [(
-				TwLabel(self.proto, mmid + ' ' + d['comment']),
-				CoinAddr(self.proto, d['addr'])
-			) for mmid, d in self.twctl.mmid_ordered_dict.items()] or None
+	pass
 
 class EthereumTokenTwRPC(EthereumTwRPC):
 	pass
diff --git a/mmgen/tw/ctl.py b/mmgen/tw/ctl.py
index 5863ab5b..3967762f 100755
--- a/mmgen/tw/ctl.py
+++ b/mmgen/tw/ctl.py
@@ -31,7 +31,8 @@ from ..addr import CoinAddr,is_mmgen_id,is_coin_addr
 from ..rpc import rpc_init
 from .shared import TwMMGenID,TwLabel
 
-addr_info = namedtuple('addr_info',['twmmid','coinaddr'])
+twmmid_addr_pair = namedtuple('addr_info',['twmmid','coinaddr'])
+label_addr_pair = namedtuple('label_addr_pair',['label','coinaddr'])
 
 # decorator for TwCtl
 def write_mode(orig_func):
@@ -219,45 +220,38 @@ class TwCtl(MMGenObject,metaclass=AsyncInit):
 
 		twmmid,coinaddr = (None,None)
 
+		pairs = await self.get_label_addr_pairs()
+
 		if is_coin_addr(self.proto,addrspec):
 			coinaddr = get_obj(CoinAddr,proto=self.proto,addr=addrspec)
+			pair_data = [e for e in pairs if e.coinaddr == coinaddr]
 		elif is_mmgen_id(self.proto,addrspec):
 			twmmid = TwMMGenID(self.proto,addrspec)
+			pair_data = [e for e in pairs if e.label.mmid == twmmid]
 		else:
 			msg(f'{addrspec!r}: invalid address for this network')
 			return None
 
-		pairs = await self.get_label_addr_pairs(twmmid)
-
-		if not pairs:
-			msg(f'MMGen address {twmmid!r} not found in tracking wallet')
+		if not pair_data:
+			msg('{a} address {b!r} not found in tracking wallet'.format(
+				a = 'MMGen' if twmmid else 'Coin',
+				b = twmmid or coinaddr))
 			return None
 
-		pairs_data = dict((label.mmid,addr) for label,addr in pairs)
-
-		if twmmid and not coinaddr:
-			coinaddr = pairs_data[twmmid]
-
-		# Allow for the possibility that BTC addr of MMGen addr was entered.
-		# Do reverse lookup, so that MMGen addr will not be marked as non-MMGen.
-		if not twmmid:
-			for mmid,addr in pairs_data.items():
-				if coinaddr == addr:
-					twmmid = mmid
-					break
-			else:
-				msg(f'Coin address {addrspec!r} not found in tracking wallet')
-				return None
-
-		return addr_info(twmmid,coinaddr)
+		return twmmid_addr_pair(
+			twmmid or pair_data[0].label.mmid,
+			coinaddr or pair_data[0].coinaddr)
 
 	# returns on failure
 	@write_mode
-	async def set_comment(self,addrspec,comment='',trusted_coinaddr=None,silent=False):
+	async def set_comment(
+			self,
+			addrspec,
+			comment      = '',
+			trusted_pair = None,
+			silent       = False):
 
-		res = (
-			addr_info(addrspec,trusted_coinaddr) if trusted_coinaddr
-			else await self.resolve_address(addrspec) )
+		res = twmmid_addr_pair(*trusted_pair) if trusted_pair else await self.resolve_address(addrspec)
 
 		if not res:
 			return False
@@ -276,9 +270,6 @@ class TwCtl(MMGenObject,metaclass=AsyncInit):
 			return False
 
 		if await self.set_label(res.coinaddr,lbl):
-			# redundant paranoia step:
-			pairs = await self.get_label_addr_pairs(res.twmmid)
-			assert pairs[0][0].comment == comment, f'{pairs[0][0].comment!r} != {comment!r}'
 			if not silent:
 				desc = '{} address {} in tracking wallet'.format(
 					res.twmmid.type.replace('mmgen','MMGen'),
diff --git a/mmgen/tw/view.py b/mmgen/tw/view.py
index f58d97e0..5d71fb62 100755
--- a/mmgen/tw/view.py
+++ b/mmgen/tw/view.py
@@ -723,7 +723,11 @@ class TwView(MMGenObject,metaclass=AsyncInit):
 
 			async def do_comment_add(comment):
 
-				if await parent.twctl.set_comment( entry.twmmid, comment, entry.addr, silent=parent.scroll ):
+				if await parent.twctl.set_comment(
+						addrspec     = None,
+						comment      = comment,
+						trusted_pair = (entry.twmmid, entry.addr),
+						silent       = parent.scroll):
 					entry.comment = comment
 					edited = cur_comment and comment
 					parent.oneshot_msg = (green if comment else yellow)('Label {a} {b}{c}'.format(
diff --git a/test/cmdtest_py_d/ct_regtest.py b/test/cmdtest_py_d/ct_regtest.py
index a2b6acb0..7f4eb594 100755
--- a/test/cmdtest_py_d/ct_regtest.py
+++ b/test/cmdtest_py_d/ct_regtest.py
@@ -1173,9 +1173,9 @@ class CmdTestRegtest(CmdTestBase,CmdTestShared):
 	def bob_resolve_addr(self):
 		mmaddr = '{}:C:1'.format(self._user_sid('bob'))
 		t = self.spawn('mmgen-tool',['--bob','resolve_address',mmaddr])
-		coinaddr = t.read().split()[0].strip()
+		coinaddr = re.search(r'[0-9A-Za-z]{30,}', t.read())[0]
 		t = self.spawn('mmgen-tool',['--bob','resolve_address',coinaddr],no_msg=True)
-		mmaddr_res = t.read().split()[0].strip()
+		mmaddr_res = re.search(r'[0-9A-F]{8}:C:1', t.read())[0]
 		assert mmaddr == mmaddr_res, f'{mmaddr} != {mmaddr_res}'
 		return t