From 2856c8f39d9ae757c6bbfdf70480b7f10a194e42 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Sun, 27 Jun 2021 20:58:02 +0000
Subject: [PATCH] mmgen-xmrwallet: add --rescan-blockchain option

---
 mmgen/main_xmrwallet.py        |  3 +++
 mmgen/xmrwallet.py             | 16 ++++++++++++++--
 test/test_py_d/ts_xmrwallet.py |  2 +-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/mmgen/main_xmrwallet.py b/mmgen/main_xmrwallet.py
index 65875473..924eba7e 100755
--- a/mmgen/main_xmrwallet.py
+++ b/mmgen/main_xmrwallet.py
@@ -38,6 +38,7 @@ opts_data = {
 -h, --help                       Print this help message
 --, --longhelp                   Print help message for long options (common
                                  options)
+-b, --rescan-blockchain          Rescan the blockchain if wallet fails to sync
 -d, --outdir=D                   Output or operate on wallets in directory 'D'
                                  instead of working dir
 -D, --daemon=H:P                 Connect to monerod at {D}
@@ -150,6 +151,7 @@ uo = namedtuple('uopts',[
 	'daemon',
 	'tx_relay_daemon',
 	'restore_height',
+	'rescan_blockchain',
 	'no_start_wallet_daemon',
 	'no_stop_wallet_daemon',
 ])
@@ -159,6 +161,7 @@ uopts = uo(
 	opt.daemon or '',
 	opt.tx_relay_daemon or '',
 	opt.restore_height or 0,
+	opt.rescan_blockchain,
 	opt.no_start_wallet_daemon,
 	opt.no_stop_wallet_daemon,
 )
diff --git a/mmgen/xmrwallet.py b/mmgen/xmrwallet.py
index f33371bb..6e601e28 100755
--- a/mmgen/xmrwallet.py
+++ b/mmgen/xmrwallet.py
@@ -391,6 +391,7 @@ class MoneroWalletOps:
 		name    = 'sync'
 		desc    = 'Sync'
 		past    = 'synced'
+		opts    = ('rescan_blockchain',)
 
 		def __init__(self,uarg_tuple,uopt_tuple):
 
@@ -432,6 +433,17 @@ class MoneroWalletOps:
 			if ret['received_money']:
 				msg('  Wallet has received funds')
 
+			while True:
+				wallet_height = (await self.c.call('get_height'))['height']
+				if wallet_height >= chain_height:
+					break
+				ymsg(f'  Wallet failed to sync (wallet height [{wallet_height}] < chain height [{chain_height}])')
+				if not uopt.rescan_blockchain:
+					break
+				msg('  Rescanning blockchain...')
+				await self.c.call('rescan_blockchain')
+				await self.c.call('refresh')
+
 			t_elapsed = int(time.time() - t_start)
 
 			bn = os.path.basename(fn)
@@ -445,11 +457,11 @@ class MoneroWalletOps:
 
 			self.accts_data[bn] = { 'accts': a, 'addrs': b }
 
-			msg('  Wallet height: {}'.format( (await self.c.call('get_height'))['height'] ))
+			msg('  Wallet height: {}'.format(wallet_height))
 			msg('  Sync time: {:02}:{:02}'.format( t_elapsed//60, t_elapsed%60 ))
 
 			await self.c.call('close_wallet')
-			return True
+			return wallet_height >= chain_height
 
 		def post_main(self):
 			d = self.accts_data
diff --git a/test/test_py_d/ts_xmrwallet.py b/test/test_py_d/ts_xmrwallet.py
index 5b251a2f..b008d87f 100755
--- a/test/test_py_d/ts_xmrwallet.py
+++ b/test/test_py_d/ts_xmrwallet.py
@@ -325,7 +325,7 @@ class TestSuiteXMRWallet(TestSuiteBase):
 		return 'ok'
 
 	def sync_wallets_selected(self):
-		return self.sync_wallets(wallets='1-2,4')
+		return self.sync_wallets(wallets='1-2,4',add_opts=['--rescan-blockchain'])
 
 	def sync_wallets(self,wallets=None,add_opts=None):
 		data = self.users['alice']