Merge pull request #10123 from accumulator/keystorewizard_scriptandderivation

wizard: add script and derivation to keystorewizard flow. fixes #10063

electrum/gui/qt/wallet_info_dialog.py

@@ -14,6 +14,7 @@ from PyQt6.QtWidgets import (
 from electrum.plugin import run_hook
 from electrum.i18n import _
 from electrum.wallet import Multisig_Wallet
+from electrum.wizard import WizardViewState
 
 from .main_window import protected
 from electrum.gui.qt.wizard.wallet import QEKeystoreWizard
@@ -192,7 +193,9 @@ class WalletInfoDialog(WindowModalDialog):
         self.window.gui_object.reload_windows()
 
     def enable_keystore(self, b: bool):
-        dialog = QEKeystoreWizard(self.window.config, self.window.wallet.wallet_type, self.window.gui_object.app, self.window.gui_object.plugins)
+        v = WizardViewState('keystore_type', {'wallet_type': self.window.wallet.wallet_type}, {})
+        dialog = QEKeystoreWizard(config=self.window.config, app=self.window.gui_object.app,
+                                  plugins=self.window.gui_object.plugins, start_viewstate=v)
         result = dialog.run()
         if not result:
             return

electrum/gui/qt/wizard/wallet.py

@@ -23,7 +23,7 @@ from .wizard import QEAbstractWizard, WizardComponent
 from electrum.logging import get_logger, Logger
 from electrum import WalletStorage, mnemonic, keystore
 from electrum.wallet_db import WalletDB
-from electrum.wizard import NewWalletWizard, KeystoreWizard
+from electrum.wizard import NewWalletWizard, KeystoreWizard, WizardViewState
 
 from electrum.gui.qt.bip39_recovery_dialog import Bip39RecoveryDialog
 from electrum.gui.qt.password_dialog import PasswordLayout, PW_NEW, MSG_ENTER_PASSWORD, PasswordLayoutForHW
@@ -48,14 +48,21 @@ MSG_HW_STORAGE_ENCRYPTION = _("Set wallet file encryption.") + '\n'\
                           + _("It also contains your master public key that allows watching your addresses.")
 
 
-
 class QEKeystoreWizard(KeystoreWizard, QEAbstractWizard, MessageBoxMixin):
     _logger = get_logger(__name__)
 
-    def __init__(self, config: 'SimpleConfig', wallet_type: str, app: 'QElectrumApplication', plugins: 'Plugins', *, start_viewstate=None):
+    def __init__(
+            self,
+            *,
+            config: 'SimpleConfig',
+            app: 'QElectrumApplication',
+            plugins: 'Plugins',
+            start_viewstate: WizardViewState = None
+    ):
+        assert 'wallet_type' in start_viewstate.wizard_data, 'wallet_type required'
+
         QEAbstractWizard.__init__(self, config, app, start_viewstate=start_viewstate)
         KeystoreWizard.__init__(self, plugins)
-        self._wallet_type = wallet_type
         self.window_title = _('Extend wallet keystore')
         # attach gui classes to views
         self.navmap_merge({
@@ -418,9 +425,7 @@ class WCKeystoreType(WalletWizardComponent):
         self.wizard_data['keystore_type'] = self.choice_w.selected_key
 
 
-
 class WCExtendKeystore(WalletWizardComponent):
-
     def __init__(self, parent, wizard):
         WalletWizardComponent.__init__(self, parent, wizard, title=_('Keystore'))
         message = _('What type of signing method do you want to add?')
@@ -432,13 +437,10 @@ class WCExtendKeystore(WalletWizardComponent):
         self.layout().addWidget(self.choice_w)
         self.layout().addStretch(1)
         self._valid = True
-        self.wizard_data['wallet_type'] = self._wallet_type = wizard._wallet_type
 
     def apply(self):
-        self.wizard_data['wallet_type'] = self._wallet_type
         self.wizard_data['keystore_type'] = self.choice_w.selected_key
-        if multisig_type(self._wallet_type):
+        if multisig_type(self.wizard_data['wallet_type']):
-            self.wizard_data['wallet_type'] = self._wallet_type = 'multisig'
             self.wizard_data['multisig_participants'] = 2
             self.wizard_data['multisig_signatures'] = 2
             self.wizard_data['multisig_cosigner_data'] = {}

electrum/keystore.py

@@ -640,7 +640,11 @@ class BIP32_KeyStore(Xpub, Deterministic_KeyStore):
         self.xprv = d.get('xprv')
 
     def watching_only_keystore(self):
-        return BIP32_KeyStore({'xpub':self.xpub})
+        return BIP32_KeyStore({
+            'xpub': self.xpub,
+            'root_fingerprint': self.get_root_fingerprint(),
+            'derivation_prefix': self.get_derivation_prefix(),
+        })
 
     def format_seed(self, seed):
         return ' '.join(seed.split())
@@ -895,6 +899,13 @@ class Hardware_KeyStore(Xpub, KeyStore):
         self.handler = None  # type: Optional[HardwareHandlerBase]
         run_hook('init_keystore', self)
 
+    def watching_only_keystore(self):
+        return BIP32_KeyStore({
+            'xpub': self.xpub,
+            'root_fingerprint': self.get_root_fingerprint(),
+            'derivation_prefix': self.get_derivation_prefix(),
+        })
+
     def set_label(self, label: Optional[str]) -> None:
         self.label = label
 
@@ -914,13 +925,13 @@ class Hardware_KeyStore(Xpub, KeyStore):
             'xpub': self.xpub,
             'derivation': self.get_derivation_prefix(),
             'root_fingerprint': self.get_root_fingerprint(),
-            'label':self.label,
+            'label': self.label,
             'soft_device_id': self.soft_device_id,
         }
 
     def is_watching_only(self):
-        '''The wallet is not watching-only; the user will be prompted for
+        """The wallet is not watching-only; the user will be prompted for
-        pin and passphrase as appropriate when needed.'''
+        pin and passphrase as appropriate when needed."""
         assert not self.has_seed()
         return False
 

electrum/wizard.py

@@ -210,17 +210,28 @@ class KeystoreWizard(AbstractWizard):
                 'next': self.on_keystore_type
             },
             'enter_seed': {
-                'next': 'enter_ext',
+                'next': lambda d: 'enter_ext' if self.wants_ext(d) else 'script_and_derivation',
-                'accept': lambda d: None if self.wants_ext(d) else self.update_keystore(d),
+                'accept': lambda d: None if (self.wants_ext(d) or self.needs_derivation_path(d)) else self.update_keystore(d),
-                'last': lambda d: not self.wants_ext(d),
+                'last': lambda d: not self.wants_ext(d) and not self.needs_derivation_path(d),
             },
             'enter_ext': {
+                'next': 'script_and_derivation',
+                'accept': lambda d: None if self.needs_derivation_path(d) else self.update_keystore(d),
+                'last': lambda d: not self.needs_derivation_path(d)
+            },
+            'script_and_derivation': {
                 'accept': self.update_keystore,
                 'last': True
             },
             'choose_hardware_device': {
                 'next': self.on_hardware_device,
             },
+            'wallet_password': {
+                'last': True
+            },
+            'wallet_password_hardware': {
+                'last': True
+            },
         }
 
     def maybe_master_pubkey(self, wizard_data):

tests/test_wizard.py

@@ -2,9 +2,10 @@ import os
 
 from electrum import SimpleConfig
 from electrum.interface import ServerAddr
+from electrum.keystore import bip44_derivation, Hardware_KeyStore
 from electrum.network import NetworkParameters, ProxySettings
 from electrum.plugin import Plugins, DeviceInfo, Device
-from electrum.wizard import ServerConnectWizard, NewWalletWizard, WizardViewState
+from electrum.wizard import ServerConnectWizard, NewWalletWizard, WizardViewState, KeystoreWizard
 from electrum.daemon import Daemon
 from electrum.wallet import Abstract_Wallet
 from electrum import util
@@ -121,7 +122,167 @@ class ServerConnectWizardTestCase(WizardTestCase):
         self.assertTrue(w._daemon.network.run_called)
         self.assertEqual(NetworkParameters(server=serverobj, proxy=None, auto_connect=False, oneserver=False), w._daemon.network.parameters)
 
-# TODO KeystoreWizard ("enable keystore")
+
+class KeystoreWizardTestCase(WizardTestCase):
+    class TKeystoreWizard(KeystoreWizard):
+        def is_single_password(self):
+            """impl abstract reqd"""
+            return True
+
+    class TNewWalletWizard(NewWalletWizard):
+        def is_single_password(self):
+            """impl abstract reqd"""
+            return True
+
+    def _wizard_for(self, *, wallet_type: str = 'standard', hww: bool = False) -> tuple[KeystoreWizard, WizardViewState]:
+        w = KeystoreWizardTestCase.TKeystoreWizard(self.plugins)
+        v = w.start({'wallet_type': wallet_type})
+        self.assertEqual('keystore_type', v.view)
+        d = v.wizard_data
+        if hww:
+            d.update({'keystore_type': 'hardware'})
+            v = w.resolve_next(v.view, d)
+            self.assertEqual('choose_hardware_device', v.view)
+        else:
+            d.update({'keystore_type': 'haveseed'})
+            v = w.resolve_next(v.view, d)
+            self.assertEqual('enter_seed', v.view)
+
+        return w, v
+
+    def _create_xpub_keystore_wallet(self, *, wallet_type: str = 'standard', xpub):
+        w = KeystoreWizardTestCase.TNewWalletWizard(DaemonMock(self.config), self.plugins)
+        wallet_path = self.wallet_path
+        d = {
+            'wallet_type': wallet_type,
+            'keystore_type': 'masterkey',
+            'master_key': xpub,
+            'password': None,
+            'encrypt': False,
+        }
+        w.create_storage(wallet_path, d)
+        self.assertTrue(os.path.exists(wallet_path))
+        wallet = Daemon._load_wallet(wallet_path, password=None, config=self.config)
+        return wallet
+
+    async def test_haveseed_electrum(self):
+        w, v = self._wizard_for()
+        d = v.wizard_data
+        d.update({
+            'seed': '9dk', 'seed_type': 'segwit', 'seed_extend': False, 'seed_variant': 'electrum',
+        })
+        self.assertTrue(w.is_last_view(v.view, d))
+        w.resolve_next(v.view, d)
+        ks, ishww = w._result
+        self.assertFalse(ishww)
+        self.assertEqual(ks.xpub, 'zpub6nAZodjgiMNf9zzX1pTqd6ZVX61ax8azhUDnWRumKVUr1VYATVoqAuqv3qKsb8WJXjxei4wei2p4vnMG9RnpKnen2kmgdhvZUmug2NnHNsr')
+
+        wallet = self._create_xpub_keystore_wallet(xpub='zpub6nAZodjgiMNf9zzX1pTqd6ZVX61ax8azhUDnWRumKVUr1VYATVoqAuqv3qKsb8WJXjxei4wei2p4vnMG9RnpKnen2kmgdhvZUmug2NnHNsr')
+        self.assertTrue(wallet.get_keystore().is_watching_only())
+        wallet.enable_keystore(ks, ishww, None)
+        self.assertFalse(wallet.get_keystore().is_watching_only())
+
+    async def test_haveseed_ext_electrum(self):
+        w, v = self._wizard_for()
+        d = v.wizard_data
+        d.update({
+            'seed': '9dk', 'seed_type': 'segwit', 'seed_extend': True, 'seed_variant': 'electrum',
+        })
+        self.assertFalse(w.is_last_view(v.view, d))
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('enter_ext', v.view)
+        d.update({'seed_extra_words': 'abc'})
+        self.assertTrue(w.is_last_view(v.view, d))
+        w.resolve_next(v.view, d)
+        ks, ishww = w._result
+        self.assertFalse(ishww)
+        self.assertEqual(ks.xpub, 'zpub6oLFCUpqxT8BUzy8g5miUuRofPZ46ZjjvZfcfH7qJanRM7aRYGpNX4uBGtcJRbgcKbi7dYkiiPw1GB2sc3SufyDcZskuQEWp5jBwbNcj1VL')
+
+        wallet = self._create_xpub_keystore_wallet(xpub='zpub6oLFCUpqxT8BUzy8g5miUuRofPZ46ZjjvZfcfH7qJanRM7aRYGpNX4uBGtcJRbgcKbi7dYkiiPw1GB2sc3SufyDcZskuQEWp5jBwbNcj1VL')
+        self.assertTrue(wallet.get_keystore().is_watching_only())
+        wallet.enable_keystore(ks, ishww, None)
+        self.assertFalse(wallet.get_keystore().is_watching_only())
+
+    async def test_haveseed_bip39(self):
+        w, v = self._wizard_for()
+        d = v.wizard_data
+        d.update({
+            'seed': '9dk', 'seed_type': 'bip39', 'seed_extend': False, 'seed_variant': 'bip39',
+        })
+        self.assertFalse(w.is_last_view(v.view, d))
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('script_and_derivation', v.view)
+        d.update({'script_type': 'p2wpkh', 'derivation_path': 'm'})
+        v = w.resolve_next(v.view, d)
+        ks, ishww = w._result
+        self.assertFalse(ishww)
+        self.assertEqual(ks.xpub, 'zpub6jftahH18ngZwMBBp7epRdBwPMPphfdy9gM6P4n5zFUXdfQJmsYfMNZoBnQMkAoBAiQYRyDQKdpxLYp6QuTrWbgmt6v1cxnFdesyiDSocAs')
+
+        wallet = self._create_xpub_keystore_wallet(xpub='zpub6jftahH18ngZwMBBp7epRdBwPMPphfdy9gM6P4n5zFUXdfQJmsYfMNZoBnQMkAoBAiQYRyDQKdpxLYp6QuTrWbgmt6v1cxnFdesyiDSocAs')
+        self.assertTrue(wallet.get_keystore().is_watching_only())
+        wallet.enable_keystore(ks, ishww, None)
+        self.assertFalse(wallet.get_keystore().is_watching_only())
+
+    async def test_haveseed_ext_bip39(self):
+        w, v = self._wizard_for()
+        d = v.wizard_data
+        d.update({
+            'seed': '9dk', 'seed_type': 'bip39', 'seed_extend': True, 'seed_variant': 'bip39',
+        })
+        self.assertFalse(w.is_last_view(v.view, d))
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('enter_ext', v.view)
+        d.update({'seed_extra_words': 'abc'})
+        v = w.resolve_next(v.view, d)
+
+        self.assertEqual('script_and_derivation', v.view)
+        d.update({'script_type': 'p2wpkh', 'derivation_path': 'm'})
+        v = w.resolve_next(v.view, d)
+        ks, ishww = w._result
+        self.assertFalse(ishww)
+        self.assertEqual(ks.xpub, 'zpub6jftahH18ngZwVNQQqNX9vgARaQRs5X89bPzjruSH2hgEBr1LRZN8reopYDALiKngTd8j5jUeGDipb68BXqjP6qMFsReLGwP6naDRvzVHxy')
+
+        wallet = self._create_xpub_keystore_wallet(xpub='zpub6jftahH18ngZwVNQQqNX9vgARaQRs5X89bPzjruSH2hgEBr1LRZN8reopYDALiKngTd8j5jUeGDipb68BXqjP6qMFsReLGwP6naDRvzVHxy')
+        self.assertTrue(wallet.get_keystore().is_watching_only())
+        wallet.enable_keystore(ks, ishww, None)
+        self.assertFalse(wallet.get_keystore().is_watching_only())
+
+    async def test_hww(self):
+        w, v = self._wizard_for(hww=True)
+        d = v.wizard_data
+        d.update({
+            'hardware_device': (
+                'trezor',
+                DeviceInfo(
+                    device=Device(path='webusb:002:1', interface_number=-1, id_='webusb:002:1', product_key='Trezor', usage_page=0, transport_ui_string='webusb:002:1'),
+                    label='trezor_unittests', initialized=True, exception=None, plugin_name='trezor', soft_device_id='088C3F260B66F60E15DE0FA5', model_name='Trezor T'))})
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('trezor_start', v.view)
+        d.update({
+            'script_type': 'p2wpkh',
+            'derivation_path': bip44_derivation(0, bip43_purpose=84)
+        })
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('trezor_xpub', v.view)
+        d.update({
+            'hw_type': 'trezor',
+            'master_key': 'zpub6rakEaM5ps5UiQ2yhbWiEkd6ceJfmuzegwc62G4itMz8L7rRFRqh6y8bTCScXV6NfTMUhANYQnfqfBd9dYfBRKf4LD1Yyfc8UvwY1MtNKWs',
+            'root_fingerprint': 'b3569ff0',
+            'label': 'test',
+            'soft_device_id': '1',
+        })
+        self.assertTrue(w.is_last_view(v.view, d))
+        v = w.resolve_next(v.view, d)
+
+        ks, ishww = w._result
+        self.assertTrue(ishww)
+
+        wallet = self._create_xpub_keystore_wallet(xpub='zpub6rakEaM5ps5UiQ2yhbWiEkd6ceJfmuzegwc62G4itMz8L7rRFRqh6y8bTCScXV6NfTMUhANYQnfqfBd9dYfBRKf4LD1Yyfc8UvwY1MtNKWs')
+        self.assertTrue(wallet.get_keystore().is_watching_only())
+        wallet.enable_keystore(ks, ishww, None)
+        self.assertFalse(wallet.get_keystore().is_watching_only())
+        self.assertTrue(isinstance(wallet.get_keystore(), Hardware_KeyStore))
+
 
 class WalletWizardTestCase(WizardTestCase):