From 5488c51353e4c1b175ca950ba5c85bd943e9059b Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Thu, 26 May 2022 16:07:21 +0000
Subject: [PATCH] new command: `mmgen-tool rescan_address`

- updates an address balance using the `scantxoutset` RPC call
---
 mmgen/base_proto/bitcoin/tw/ctl.py | 67 +++++++++++++++++++++++++++++-
 mmgen/main_tool.py                 |  1 +
 mmgen/tool/rpc.py                  |  5 +++
 test/test_py_d/ts_regtest.py       |  8 ++++
 4 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/mmgen/base_proto/bitcoin/tw/ctl.py b/mmgen/base_proto/bitcoin/tw/ctl.py
index f4440001..a48957bf 100755
--- a/mmgen/base_proto/bitcoin/tw/ctl.py
+++ b/mmgen/base_proto/bitcoin/tw/ctl.py
@@ -14,7 +14,7 @@ base_proto.bitcoin.twctl: Bitcoin base protocol tracking wallet control class
 
 from ....globalvars import g
 from ....tw.ctl import TrackingWallet
-from ....util import msg,msg_r,rmsg,die,write_mode
+from ....util import msg,msg_r,rmsg,vmsg,die,suf,fmt_list,write_mode
 
 class BitcoinTrackingWallet(TrackingWallet):
 
@@ -88,3 +88,68 @@ class BitcoinTrackingWallet(TrackingWallet):
 				tip = await self.rpc.call('getblockcount')
 
 		msg('Done')
+
+	@write_mode
+	async def rescan_address(self,addrspec):
+		res = await self.resolve_address(addrspec,None)
+		if not res:
+			return False
+		return await self.rescan_addresses([res.coinaddr])
+
+	@write_mode
+	async def rescan_addresses(self,coin_addrs):
+
+		import asyncio
+
+		async def do_scan():
+			return await self.rpc.call(
+				'scantxoutset',
+				'start',
+				[f'addr({a})' for a in coin_addrs],
+				timeout = 720 ) # call may take several minutes to complete
+
+		async def do_status():
+			m = f'{CR}Scanning UTXO set: '
+			msg_r(m)
+			while True:
+				await asyncio.sleep(2)
+				res = await self.rpc.call('scantxoutset','status')
+				if res:
+					msg_r(m + f'{res["progress"]}% completed ')
+				if task1.done():
+					msg('')
+					return
+
+		CR = '\r'
+
+		if self.rpc.backend.name == 'aiohttp':
+			task1 = asyncio.create_task( do_scan() )
+			task2 = asyncio.create_task( do_status() )
+			res = await task1
+			await task2
+		else:
+			msg_r(f'Scanning UTXO set, this could take several minutes...')
+			res = await do_scan()
+			msg('done')
+
+		if not res['success']:
+			msg('UTXO scanning failed or was interrupted')
+			return False
+		elif res['unspents']:
+			blocks = sorted({ i['height'] for i in res['unspents'] })
+			msg('Found {} unspent output{} in {} block{}'.format(
+				len(res['unspents']),
+				suf(res['unspents']),
+				len(blocks),
+				suf(blocks) ))
+			vmsg(f'Blocks to rescan: {fmt_list(blocks,fmt="bare")}')
+			for n,block in enumerate(blocks):
+				msg_r(f'{CR}Rescanning block: {block} ({n+1}/{len(blocks)})')
+				# httplib seems to require fresh connection here, so specify timeout
+				await self.rpc.call('rescanblockchain',block,block,timeout=60)
+			msg('\nRescan completed OK')
+			return True
+		else:
+			msg('Imported address has no balance' if len(coin_addrs) == 1 else
+				'Imported addresses have no balances' )
+			return True
diff --git a/mmgen/main_tool.py b/mmgen/main_tool.py
index 1d927247..5c21c545 100755
--- a/mmgen/main_tool.py
+++ b/mmgen/main_tool.py
@@ -161,6 +161,7 @@ mods = {
 		'listaddresses',
 		'remove_address',
 		'remove_label',
+		'rescan_address',
 		'rescan_blockchain',
 		'resolve_address',
 		'twview',
diff --git a/mmgen/tool/rpc.py b/mmgen/tool/rpc.py
index 4de03297..5ecd1e81 100755
--- a/mmgen/tool/rpc.py
+++ b/mmgen/tool/rpc.py
@@ -179,6 +179,11 @@ class tool_cmd(tool_cmd_base):
 		else:
 			return False
 
+	async def rescan_address(self,mmgen_or_coin_addr:str):
+		"rescan an address in the tracking wallet to update its balance"
+		from ..tw.ctl import TrackingWallet
+		return await (await TrackingWallet(self.proto,mode='w')).rescan_address( mmgen_or_coin_addr )
+
 	async def rescan_blockchain(self,
 			start_block: int = None,
 			stop_block: int  = None ):
diff --git a/test/test_py_d/ts_regtest.py b/test/test_py_d/ts_regtest.py
index 35c03ead..2c810d87 100755
--- a/test/test_py_d/ts_regtest.py
+++ b/test/test_py_d/ts_regtest.py
@@ -203,6 +203,7 @@ class TestSuiteRegtest(TestSuiteBase,TestSuiteShared):
 		('bob_import_list',          'importing flat address list'),
 		('bob_import_list_rescan',   'importing flat address list with --rescan'),
 		('bob_resolve_addr',         'resolving an address in the tracking wallet'),
+		('bob_rescan_addr',          'rescanning an address'),
 		('bob_rescan_blockchain_all','rescanning the blockchain (full rescan)'),
 		('bob_rescan_blockchain_gb', 'rescanning the blockchain (Genesis block)'),
 		('bob_rescan_blockchain_one','rescanning the blockchain (single block)'),
@@ -931,6 +932,13 @@ class TestSuiteRegtest(TestSuiteBase,TestSuiteShared):
 		addrfile = joinpath(self.tmpdir,'non-mmgen.addrs')
 		return self.user_import('bob',['--quiet','--rescan','--addrlist',addrfile],nAddr=5)
 
+	def bob_rescan_addr(self):
+		sid = self._user_sid('bob')
+		t = self.spawn('mmgen-tool',['--bob','rescan_address',f'{sid}:C:1'])
+		t.expect('Found 1 unspent output')
+		t.expect('completed OK')
+		return t
+
 	def bob_rescan_blockchain(self,add_args,expect):
 		t = self.spawn('mmgen-tool',['--bob','rescan_blockchain'] + add_args)
 		t.expect(f'Scanning blocks {expect}')