From 9e3d8d92ded1447f6a27522acca825a1e5b706d7 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Fri, 11 Feb 2022 13:41:03 +0000
Subject: [PATCH] mmgen-passchg: improve secure wallet deletion logic

---
 mmgen/data/version        |  2 +-
 mmgen/main_wallet.py      | 44 ++++++++++++++++++++++++++++++---------
 test/test_py_d/ts_main.py |  7 +++++--
 3 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/mmgen/data/version b/mmgen/data/version
index 2a5f20db..8dcbceb3 100644
--- a/mmgen/data/version
+++ b/mmgen/data/version
@@ -1 +1 @@
-13.1.dev19
+13.1.dev20
diff --git a/mmgen/main_wallet.py b/mmgen/main_wallet.py
index 07451c37..0c1fe5b6 100755
--- a/mmgen/main_wallet.py
+++ b/mmgen/main_wallet.py
@@ -220,16 +220,40 @@ if invoked_as == 'passchg':
 	if not ( opt.force_update or data_changed(('passwd','hash_preset','label')) ):
 		die(1,'Password, hash preset and label are unchanged.  Taking no action')
 
-if invoked_as == 'passchg' and ss_in.infile.dirname == g.data_dir:
-	m1 = yellow('Confirmation of default wallet update')
-	m2 = 'update the default wallet'
-	confirm_or_raise(m1,m2,exit_msg='Password not changed')
-	ss_out.write_to_file(desc='New wallet',outdir=g.data_dir)
-	bmsg('Securely deleting old wallet')
-	from .fileutil import shred_file
-	shred_file(
-		ss_in.infile.name,
-		verbose = opt.verbose )
+if invoked_as == 'passchg':
+
+	def secure_delete(fn):
+		bmsg('Securely deleting old wallet')
+		from .fileutil import shred_file
+		shred_file( fn, verbose = opt.verbose )
+
+	def rename_old_wallet_maybe(silent):
+		# though very unlikely, old and new wallets could have same Key ID and thus same filename.
+		# If so, rename old wallet file before deleting.
+		old_fn = ss_in.infile.name
+		if os.path.basename(old_fn) == ss_out._filename():
+			if not silent:
+				ymsg(
+					'Warning: diverting old wallet {old_fn!r} due to Key ID collision.  ' +
+					'Please securely delete or move the diverted file!' )
+			os.rename( old_fn, old_fn+'.divert' )
+			return old_fn+'.divert'
+		else:
+			return old_fn
+
+	if ss_in.infile.dirname == g.data_dir:
+		confirm_or_raise(
+			message  = yellow('Confirmation of default wallet update'),
+			action   = 'update the default wallet',
+			exit_msg = 'Password not changed' )
+		old_wallet = rename_old_wallet_maybe(silent=True)
+		ss_out.write_to_file( desc='New wallet', outdir=g.data_dir )
+		secure_delete( old_wallet )
+	else:
+		old_wallet = rename_old_wallet_maybe(silent=False)
+		ss_out.write_to_file()
+		if keypress_confirm(f'Securely delete old wallet {old_wallet!r}?'):
+			secure_delete( old_wallet )
 elif invoked_as == 'gen' and not opt.outdir and not opt.stdout:
 	from .filename import find_file_in_dir
 	if (
diff --git a/test/test_py_d/ts_main.py b/test/test_py_d/ts_main.py
index 1b471e13..ab6ecce2 100755
--- a/test/test_py_d/ts_main.py
+++ b/test/test_py_d/ts_main.py
@@ -290,7 +290,7 @@ class TestSuiteMain(TestSuiteBase,TestSuiteShared):
 		assert ext == ocls.ext,f'incorrect file extension: {ext}'
 		return t
 
-	def passchg(self,wf,pf,label_action='cmdline',dfl_wallet=False):
+	def passchg(self,wf,pf,label_action='cmdline',dfl_wallet=False,delete=False):
 		silence()
 		self.write_to_tmpfile(pwfile,get_data_from_file(pf))
 		end_silence()
@@ -319,10 +319,13 @@ class TestSuiteMain(TestSuiteBase,TestSuiteShared):
 			t.expect_getend('has been changed to ')
 		else:
 			t.written_to_file(capfirst(wcls.desc))
+			t.expect( 'Securely delete .*: ', ('y' if delete else 'n' ), regex=True )
+			if delete:
+				t.expect( 'Securely deleting' )
 		return t
 
 	def passchg_keeplabel(self,wf,pf):
-		return self.passchg(wf,pf,label_action='keep')
+		return self.passchg(wf,pf,label_action='keep',delete=True)
 
 	def passchg_usrlabel(self,wf,pf):
 		return self.passchg(wf,pf,label_action='user')