restore from xprv

[electrum-nvc.git] / lib / wallet.py

diff --git a/lib/wallet.py b/lib/wallet.py
index 467a11c..67676b5 100644 (file)
--- a/lib/wallet.py
+++ b/lib/wallet.py
@@ -1266,8 +1266,11 @@ class NewWallet(Deterministic_Wallet):
     def __init__(self, storage):
         Deterministic_Wallet.__init__(self, storage)
 
+    def is_watching_only(self):
+        return self.master_private_keys is {}
+
     def can_create_accounts(self):
-        return not self.is_watching_only()
+        return 'm/' in self.master_private_keys.keys()
 
     def get_master_public_key(self):
         return self.master_public_keys["m/"]
@@ -1291,19 +1294,29 @@ class NewWallet(Deterministic_Wallet):
         xpub = self.master_public_keys["m/"]
         assert deserialize_xkey(xpriv)[3] == deserialize_xkey(xpub)[3]
 
-    def create_watching_only_wallet(self, xpub):
+    def create_xprv_wallet(self, xprv, password):
+        xpub = bitcoin.xpub_from_xprv(xprv)
+        account = BIP32_Account({'xpub':xpub})
+        account_id = 'm/' + bitcoin.get_xkey_name(xpub)
         self.storage.put('seed_version', self.seed_version, True)
-        self.add_master_public_key("m/", xpub)
+        self.add_master_private_key(account_id, xprv, password)
+        self.add_master_public_key(account_id, xpub)
+        self.add_account(account_id, account)
+
+    def create_watching_only_wallet(self, xpub):
         account = BIP32_Account({'xpub':xpub})
-        self.add_account("m/", account)
+        account_id = 'm/' + bitcoin.get_xkey_name(xpub)
+        self.storage.put('seed_version', self.seed_version, True)
+        self.add_master_public_key(account_id, xpub)
+        self.add_account(account_id, account)
 
     def create_accounts(self, password):
         # First check the password is valid (this raises if it isn't).
         self.check_password(password)
         self.create_account('Main account', password)
 
-    def add_master_public_key(self, name, mpk):
-        self.master_public_keys[name] = mpk
+    def add_master_public_key(self, name, xpub):
+        self.master_public_keys[name] = xpub
         self.storage.put('master_public_keys', self.master_public_keys, True)
 
     def add_master_private_key(self, name, xpriv, password):
@@ -1576,21 +1589,31 @@ class Wallet(object):
             return False
 
     @classmethod
-    def is_mpk(self, mpk):
+    def is_old_mpk(self, mpk):
         try:
             int(mpk, 16)
-            old = True
+            assert len(mpk) == 128
+            return True
         except:
-            old = False
+            return False
 
-        if old:
-            return len(mpk) == 128
-        else:
-            try:
-                deserialize_xkey(mpk)
-                return True
-            except:
-                return False
+    @classmethod
+    def is_xpub(self, text):
+        try:
+            assert text[0:4] == 'xpub'
+            deserialize_xkey(text)
+            return True
+        except:
+            return False
+
+    @classmethod
+    def is_xprv(self, text):
+        try:
+            assert text[0:4] == 'xprv'
+            deserialize_xkey(text)
+            return True
+        except:
+            return False
 
     @classmethod
     def is_address(self, text):
@@ -1635,19 +1658,20 @@ class Wallet(object):
         return w
 
     @classmethod
-    def from_mpk(self, mpk, storage):
-        try:
-            int(mpk, 16)
-            old = True
-        except:
-            old = False
+    def from_old_mpk(self, mpk, storage):
+        w = OldWallet(storage)
+        w.seed = ''
+        w.create_watching_only_wallet(mpk)
+        return w
 
-        if old:
-            w = OldWallet(storage)
-            w.seed = ''
-            w.create_watching_only_wallet(mpk)
-        else:
-            w = NewWallet(storage)
-            w.create_watching_only_wallet(mpk)
+    @classmethod
+    def from_xpub(self, xpub, storage):
+        w = NewWallet(storage)
+        w.create_watching_only_wallet(xpub)
+        return w
 
+    @classmethod
+    def from_xprv(self, xprv, password, storage):
+        w = NewWallet(storage)
+        w.create_xprv_wallet(xprv, password)
         return w