From 9e1c1bd0ad1769c25b175037059982d31b51c77e Mon Sep 17 00:00:00 2001
From: Sander van Grieken <sander@outrightsolutions.nl>
Date: Tue, 8 Oct 2024 00:04:20 +0200
Subject: [PATCH] trustedcoin: fix continuation of 2fa wallet file with
 keystore-only encryption

---
 electrum/gui/qt/__init__.py                 |  3 ++
 electrum/plugins/trustedcoin/qml.py         |  3 ++
 electrum/plugins/trustedcoin/qt.py          | 46 ++++++++++++++++++++-
 electrum/plugins/trustedcoin/trustedcoin.py |  6 ++-
 4 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/electrum/gui/qt/__init__.py b/electrum/gui/qt/__init__.py
index 35853ea56..8400823d9 100644
--- a/electrum/gui/qt/__init__.py
+++ b/electrum/gui/qt/__init__.py
@@ -72,6 +72,7 @@ from electrum.simple_config import SimpleConfig
 from electrum.storage import WalletStorage
 from electrum.wizard import WizardViewState
 from electrum.keystore import load_keystore
+from electrum.bip32 import is_xprv
 
 from electrum.gui.common_qt.i18n import ElectrumTranslator
 
@@ -468,6 +469,8 @@ class ElectrumGui(BaseElectrumGui, Logger):
                 xprv = k1.get_master_private_key(d['password'])
             else:
                 xprv = db.get('x1')['xprv']
+                if not is_xprv(xprv):
+                    xprv = k1
             _wiz_data_updates = {
                 'wallet_name': wallet_file,
                 'xprv1': xprv,
diff --git a/electrum/plugins/trustedcoin/qml.py b/electrum/plugins/trustedcoin/qml.py
index a9b30e75d..5bc6b330c 100644
--- a/electrum/plugins/trustedcoin/qml.py
+++ b/electrum/plugins/trustedcoin/qml.py
@@ -75,6 +75,9 @@ class Plugin(TrustedCoinPlugin):
             'trustedcoin_tos': {
                 'gui': '../../../../plugins/trustedcoin/qml/Terms',
             },
+            'trustedcoin_keystore_unlock': {
+                # TODO when QML can import external wallet files
+            },
             'trustedcoin_show_confirm_otp': {
                 'gui': '../../../../plugins/trustedcoin/qml/ShowConfirmOTP',
             }
diff --git a/electrum/plugins/trustedcoin/qt.py b/electrum/plugins/trustedcoin/qt.py
index 69c1233af..0f98f6c8f 100644
--- a/electrum/plugins/trustedcoin/qt.py
+++ b/electrum/plugins/trustedcoin/qt.py
@@ -35,13 +35,13 @@ from PyQt6.QtWidgets import (QTextEdit, QVBoxLayout, QLabel, QGridLayout, QHBoxL
 
 from electrum.i18n import _
 from electrum.plugin import hook
-from electrum.util import is_valid_email
+from electrum.util import is_valid_email, InvalidPassword
 from electrum.logging import Logger, get_logger
 from electrum import keystore
 
 from electrum.gui.qt.util import (read_QIcon, WindowModalDialog, WaitingDialog, OkButton,
                                   CancelButton, Buttons, icon_path, WWLabel, CloseButton, ColorScheme,
-                                  ChoiceWidget)
+                                  ChoiceWidget, PasswordLineEdit, char_width_in_lineedit)
 from electrum.gui.qt.qrcodewidget import QRCodeWidget
 from electrum.gui.qt.amountedit import AmountEdit
 from electrum.gui.qt.main_window import StatusBarButton
@@ -261,6 +261,10 @@ class Plugin(TrustedCoinPlugin):
                 'gui': WCTerms,
                 'params': {'icon': icon_path('trustedcoin-wizard.png')},
             },
+            'trustedcoin_keystore_unlock': {
+                'gui': WCKeystorePassword,
+                'params': {'icon': icon_path('trustedcoin-wizard.png')},
+            },
             'trustedcoin_show_confirm_otp': {
                 'gui': WCShowConfirmOTP,
                 'params': {'icon': icon_path('trustedcoin-wizard.png')},
@@ -600,3 +604,41 @@ class WCContinueOnline(WizardComponent):
 
     def apply(self):
         self.wizard_data['trustedcoin_go_online'] = self.cb_online.isChecked()
+
+
+class WCKeystorePassword(WizardComponent):
+    def __init__(self, parent, wizard):
+        WizardComponent.__init__(self, parent, wizard, title=_('Unlock Keystore'))
+
+    def on_ready(self):
+        self.layout().addStretch(1)
+
+        hbox2 = QHBoxLayout()
+        hbox2.addStretch(1)
+        self.pw_e = PasswordLineEdit('', self)
+        self.pw_e.setFixedWidth(17 * char_width_in_lineedit())
+        self.pw_e.textEdited.connect(self.on_text)
+        pw_label = QLabel(_('Password') + ':')
+        hbox2.addWidget(pw_label)
+        hbox2.addWidget(self.pw_e)
+        hbox2.addStretch(1)
+        self.layout().addLayout(hbox2)
+
+        self.layout().addStretch(1)
+
+        self._valid = False
+
+        self.ks = self.wizard_data['xprv1']
+
+    def on_text(self):
+        try:
+            self.ks.check_password(self.pw_e.text())
+        except InvalidPassword:
+            self.valid = False
+            return
+        self.valid = True
+
+    def apply(self):
+        if self.valid:
+            self.wizard_data['xprv1'] = self.ks.get_master_private_key(self.pw_e.text())
+            self.wizard_data['password'] = self.pw_e.text()
diff --git a/electrum/plugins/trustedcoin/trustedcoin.py b/electrum/plugins/trustedcoin/trustedcoin.py
index 4304ce1cd..199591d51 100644
--- a/electrum/plugins/trustedcoin/trustedcoin.py
+++ b/electrum/plugins/trustedcoin/trustedcoin.py
@@ -35,7 +35,7 @@ from aiohttp import ClientResponse
 import electrum_ecc as ecc
 
 from electrum import constants, keystore, version, bip32, bitcoin
-from electrum.bip32 import BIP32Node, xpub_type
+from electrum.bip32 import BIP32Node, xpub_type, is_xprv
 from electrum.crypto import sha256
 from electrum.transaction import PartialTxOutput, PartialTxInput, PartialTransaction, Transaction
 from electrum.mnemonic import Mnemonic, calc_seed_type, is_any_2fa_seed_type
@@ -598,6 +598,10 @@ class TrustedCoinPlugin(BasePlugin):
                 'last': lambda d: wizard.is_single_password() and d['trustedcoin_keepordisable'] == 'disable'
             },
             'trustedcoin_tos': {
+                'next': lambda d: 'trustedcoin_show_confirm_otp' if is_xprv(d['xprv1'])
+                        else 'trustedcoin_keystore_unlock'
+            },
+            'trustedcoin_keystore_unlock': {
                 'next': 'trustedcoin_show_confirm_otp'
             },
             'trustedcoin_show_confirm_otp': {