wizard: display seed type. restore 2fa if needed

lib/base_wizard.py

@@ -24,6 +24,7 @@
 # SOFTWARE.
 
 import os
+import bitcoin
 import keystore
 from wallet import Wallet, Imported_Wallet, Standard_Wallet, Multisig_Wallet, WalletStorage, wallet_types
 from i18n import _
@@ -85,6 +86,11 @@ class BaseWizard(object):
         choices = [pair for pair in wallet_kinds if pair[0] in wallet_types]
         self.choice_dialog(title=title, message=message, choices=choices, run_next=self.on_wallet_type)
 
+    def load_2fa(self):
+        self.storage.put('wallet_type', '2fa')
+        self.storage.put('use_trustedcoin', True)
+        self.plugin = self.plugins.load_plugin('trustedcoin')
+
     def on_wallet_type(self, choice):
         self.wallet_type = choice
         if choice == 'standard':
@@ -92,9 +98,7 @@ class BaseWizard(object):
         elif choice == 'multisig':
             action = 'choose_multisig'
         elif choice == '2fa':
-            self.storage.put('wallet_type', '2fa')
-            self.storage.put('use_trustedcoin', True)
-            self.plugin = self.plugins.load_plugin('trustedcoin')
+            self.load_2fa()
             action = self.storage.get_action()
         elif choice == 'imported':
             action = 'import_addresses'
@@ -243,31 +247,46 @@ class BaseWizard(object):
         k = hardware_keystore(d)
         self.on_keystore(k)
 
+    def passphrase_dialog(self, run_next):
+        message = '\n'.join([
+            _('Your seed may be extended with a passphrase.'),
+            _('If that is the case, enter it here.'),
+        ])
+        warning = '\n'.join([
+            _('Note that this is NOT your encryption password.'),
+            _('If you do not know what this is, leave this field empty.'),
+        ])
+        self.line_dialog(title=_('Passphrase'), message=message, warning=warning, default='', test=lambda x:True, run_next=run_next)
+
     def restore_from_seed(self):
-        self.opt_bip39 = True
-        self.restore_seed_dialog(run_next=self.on_restore_seed, test=keystore.is_seed)
+        if self.wallet_type == 'standard':
+            self.opt_bip39 = True
+            test = bitcoin.is_seed
+        else:
+            self.opt_bip39 = False
+            test = bitcoin.is_new_seed
+        self.restore_seed_dialog(run_next=self.on_restore_seed, test=test)
 
     def on_restore_seed(self, seed, is_bip39):
-        if keystore.is_new_seed(seed) or is_bip39:
-            message = '\n'.join([
-                _('Your seed may be extended with a passphrase.'),
-                _('If that is the case, enter it here.'),
-            ])
-            warning = '\n'.join([
-                _('Note that this is NOT your encryption password.'),
-                _('If you do not know what this is, leave this field empty.'),
-            ])
-            f = lambda x: self.on_restore_passphrase(seed, x, is_bip39)
-            self.line_dialog(title=_('Passphrase'), message=message, warning=warning, default='', test=lambda x:True, run_next=f)
-        else:
-            self.on_restore_passphrase(seed, '', False)
-
-    def on_restore_passphrase(self, seed, passphrase, is_bip39):
         if is_bip39:
-            f = lambda x: self.run('on_bip44', seed, passphrase, int(x))
-            self.account_id_dialog(f)
+            f = lambda x: self.on_restore_bip39(seed, x)
+            self.passphrase_dialog(run_next=f)
         else:
-            self.run('create_keystore', seed, passphrase)
+            seed_type = bitcoin.seed_type(seed)
+            if seed_type == 'standard':
+                f = lambda x: self.run('create_keystore', seed, x)
+                self.passphrase_dialog(run_next=f)
+            elif seed_type == 'old':
+                self.run('create_keystore', seed, passphrase)
+            elif seed_type == '2fa':
+                self.load_2fa()
+                self.run('on_restore_seed', seed)
+            else:
+                raise
+
+    def on_restore_bip39(self, seed, passphrase):
+        f = lambda x: self.run('on_bip44', seed, passphrase, int(x))
+        self.account_id_dialog(f)
 
     def create_keystore(self, seed, passphrase):
         k = keystore.from_seed(seed, passphrase)

lib/bitcoin.py

@@ -173,16 +173,25 @@ def is_old_seed(seed):
         uses_electrum_words = True
     except Exception:
         uses_electrum_words = False
-
     try:
         seed.decode('hex')
         is_hex = (len(seed) == 32 or len(seed) == 64)
     except Exception:
         is_hex = False
-
     return is_hex or (uses_electrum_words and (len(words) == 12 or len(words) == 24))
 
 
+def seed_type(x):
+    if is_old_seed(x):
+        return 'old'
+    elif is_new_seed(x):
+        return 'standard'
+    elif is_new_seed(x, version.SEED_PREFIX_2FA):
+        return '2fa'
+    return ''
+
+is_seed = lambda x: bool(seed_type(x))
+
 # pywallet openssl private key implementation
 
 def i2o_ECPublicKey(pubkey, compressed=False):

lib/keystore.py

@@ -33,7 +33,7 @@ from bitcoin import pw_encode, pw_decode, bip32_root, bip32_private_derivation,
 from bitcoin import public_key_from_private_key, public_key_to_bc_address
 from bitcoin import *
 
-from bitcoin import is_old_seed, is_new_seed
+from bitcoin import is_old_seed, is_new_seed, is_seed
 from util import PrintError, InvalidPassword
 from mnemonic import Mnemonic
 
@@ -665,7 +665,7 @@ def is_private_key_list(text):
     parts = text.split()
     return bool(parts) and all(bitcoin.is_private_key(x) for x in parts)
 
-is_seed = lambda x: is_old_seed(x) or is_new_seed(x)
+
 is_mpk = lambda x: is_old_mpk(x) or is_xpub(x)
 is_private = lambda x: is_seed(x) or is_xprv(x) or is_private_key_list(x)
 is_any_key = lambda x: is_old_mpk(x) or is_xprv(x) or is_xpub(x) or is_address_list(x) or is_private_key_list(x)