From f23bd334518d41fdc4c66598d3c1f7a5e81c1531 Mon Sep 17 00:00:00 2001
From: Sander van Grieken <sander@outrightsolutions.nl>
Date: Thu, 24 Aug 2023 11:28:57 +0200
Subject: [PATCH] wizard: add safe_t, except for wallet initial setup support.

---
 electrum/plugins/safe_t/qt.py     | 21 ++++++++++++++++++++-
 electrum/plugins/safe_t/safe_t.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/electrum/plugins/safe_t/qt.py b/electrum/plugins/safe_t/qt.py
index 07d36ffd1..73e21a595 100644
--- a/electrum/plugins/safe_t/qt.py
+++ b/electrum/plugins/safe_t/qt.py
@@ -1,5 +1,5 @@
 from functools import partial
-import threading
+from typing import TYPE_CHECKING
 
 from PyQt5.QtCore import Qt, pyqtSignal, QRegExp
 from PyQt5.QtGui import QRegExpValidator
@@ -16,7 +16,10 @@ from electrum.plugin import hook
 from ..hw_wallet.qt import QtHandlerBase, QtPluginBase
 from ..hw_wallet.plugin import only_hook_if_libraries_available
 from .safe_t import SafeTPlugin, TIM_NEW, TIM_RECOVER, TIM_MNEMONIC
+from electrum.gui.qt.wizard.wallet import WCScriptAndDerivation, WCHWUninitialized, WCHWUnlock, WCHWXPub
 
+if TYPE_CHECKING:
+    from electrum.gui.qt.wizard.wallet import QENewWalletWizard
 
 PASSPHRASE_HELP_SHORT =_(
     "Passphrases allow you to access new wallets, each "
@@ -184,6 +187,22 @@ class Plugin(SafeTPlugin, QtPlugin):
         from safetlib.qt.pinmatrix import PinMatrixWidget
         return PinMatrixWidget
 
+    @hook
+    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
+        self.extend_wizard(wizard)
+
+    # insert safe_t pages in new wallet wizard
+    def extend_wizard(self, wizard: 'QENewWalletWizard'):
+        # TODO: device initialization not ported yet to new wizard
+        super().extend_wizard(wizard)
+        views = {
+            'safet_start': {'gui': WCScriptAndDerivation},
+            'safet_xpub': {'gui': WCHWXPub},
+            'safet_not_initialized': {'gui': WCHWUninitialized},
+            'safet_unlock': {'gui': WCHWUnlock}
+        }
+        wizard.navmap_merge(views)
+
 
 class SettingsDialog(WindowModalDialog):
     '''This dialog doesn't require a device be paired with a wallet.
diff --git a/electrum/plugins/safe_t/safe_t.py b/electrum/plugins/safe_t/safe_t.py
index 850c00588..dc26f9b6f 100644
--- a/electrum/plugins/safe_t/safe_t.py
+++ b/electrum/plugins/safe_t/safe_t.py
@@ -18,6 +18,8 @@ from ..hw_wallet.plugin import is_any_tx_output_on_change_branch, trezor_validat
 
 if TYPE_CHECKING:
     from .client import SafeTClient
+    from electrum.plugin import DeviceInfo
+    from electrum.wizard import NewWalletWizard
 
 # Safe-T mini initialization methods
 TIM_NEW, TIM_RECOVER, TIM_MNEMONIC, TIM_PRIVKEY = range(0, 4)
@@ -460,3 +462,29 @@ class SafeTPlugin(HW_PluginBase):
     def get_tx(self, tx_hash):
         tx = self.prev_tx[tx_hash]
         return self.electrum_tx_to_txtype(tx)
+
+    # new wizard
+
+    def wizard_entry_for_device(self, device_info: 'DeviceInfo', *, new_wallet=True) -> str:
+        if new_wallet:
+            return 'safet_start' if device_info.initialized else 'safet_not_initialized'
+        else:
+            return 'safet_unlock'
+
+    # insert safe_t pages in new wallet wizard
+    def extend_wizard(self, wizard: 'NewWalletWizard'):
+        views = {
+            'safet_start': {
+                'next': 'safet_xpub',
+            },
+            'safet_xpub': {
+                'next': lambda d: wizard.wallet_password_view(d) if wizard.last_cosigner(d) else 'multisig_cosigner_keystore',
+                'accept': wizard.maybe_master_pubkey,
+                'last': lambda d: wizard.is_single_password() and wizard.last_cosigner(d)
+            },
+            'safet_not_initialized': {},
+            'safet_unlock': {
+                'last': True
+            },
+        }
+        wizard.navmap_merge(views)