From ddb818f77cb0e2d9810f76bff470d21eec0f43ad Mon Sep 17 00:00:00 2001
From: Qiao Jin <43407364+Qiao-Jin@users.noreply.github.com>
Date: Fri, 3 Apr 2020 13:14:48 +0800
Subject: [PATCH] Support changing password for nep6 wallet (#1504)

---
 src/neo/Wallets/NEP6/NEP6Account.cs           | 46 ++++++++++++++++++-
 src/neo/Wallets/NEP6/NEP6Wallet.cs            | 30 ++++++++++++
 .../Wallets/NEP6/UT_NEP6Account.cs            | 20 ++++++++
 .../Wallets/NEP6/UT_NEP6Wallet.cs             | 20 ++++++++
 4 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/src/neo/Wallets/NEP6/NEP6Account.cs b/src/neo/Wallets/NEP6/NEP6Account.cs
index 990612d542..58ed806fb1 100644
--- a/src/neo/Wallets/NEP6/NEP6Account.cs
+++ b/src/neo/Wallets/NEP6/NEP6Account.cs
@@ -1,12 +1,14 @@
 using Neo.IO.Json;
 using System;
+using System.Threading;
 
 namespace Neo.Wallets.NEP6
 {
     internal class NEP6Account : WalletAccount
     {
         private readonly NEP6Wallet wallet;
-        private readonly string nep2key;
+        private string nep2key;
+        private string nep2KeyNew = null;
         private KeyPair key;
         public JObject Extra;
 
@@ -83,5 +85,47 @@ public bool VerifyPassword(string password)
                 return false;
             }
         }
+
+        /// <summary>
+        /// Cache draft nep2key during wallet password changing process. Should not be called alone for a single account
+        /// </summary>
+        internal bool ChangePasswordPrepare(string password_old, string password_new)
+        {
+            if (WatchOnly) return true;
+            KeyPair keyTemplate = key;
+            if (nep2key == null)
+            {
+                if (keyTemplate == null)
+                {
+                    return true;
+                }
+            }
+            else
+            {
+                try
+                {
+                    keyTemplate = new KeyPair(Wallet.GetPrivateKeyFromNEP2(nep2key, password_old, wallet.Scrypt.N, wallet.Scrypt.R, wallet.Scrypt.P));
+                }
+                catch
+                {
+                    return false;
+                }
+            }
+            nep2KeyNew = keyTemplate.Export(password_new, wallet.Scrypt.N, wallet.Scrypt.R, wallet.Scrypt.P);
+            return true;
+        }
+
+        internal void ChangePasswordCommit()
+        {
+            if (nep2KeyNew != null)
+            {
+                nep2key = Interlocked.Exchange(ref nep2KeyNew, null);
+            }
+        }
+
+        internal void ChangePasswordRoolback()
+        {
+            nep2KeyNew = null;
+        }
     }
 }
diff --git a/src/neo/Wallets/NEP6/NEP6Wallet.cs b/src/neo/Wallets/NEP6/NEP6Wallet.cs
index daf1b9158d..2db4d87e77 100644
--- a/src/neo/Wallets/NEP6/NEP6Wallet.cs
+++ b/src/neo/Wallets/NEP6/NEP6Wallet.cs
@@ -6,6 +6,7 @@
 using System.Linq;
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading.Tasks;
 using UserWallet = Neo.Wallets.SQLite.UserWallet;
 
 namespace Neo.Wallets.NEP6
@@ -300,5 +301,34 @@ public override bool VerifyPassword(string password)
                 }
             }
         }
+
+        public bool ChangePassword(string password_old, string password_new)
+        {
+            bool succeed = true;
+            lock (accounts)
+            {
+                Parallel.ForEach(accounts.Values, (account, state) =>
+                {
+                    if (!account.ChangePasswordPrepare(password_old, password_new))
+                    {
+                        state.Stop();
+                        succeed = false;
+                    }
+                });
+            }
+            if (succeed)
+            {
+                foreach (NEP6Account account in accounts.Values)
+                    account.ChangePasswordCommit();
+                if (password != null)
+                    password = password_new;
+            }
+            else
+            {
+                foreach (NEP6Account account in accounts.Values)
+                    account.ChangePasswordRoolback();
+            }
+            return succeed;
+        }
     }
 }
diff --git a/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Account.cs b/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Account.cs
index f1c719c853..addfc5b54e 100644
--- a/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Account.cs
+++ b/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Account.cs
@@ -35,6 +35,26 @@ public void TestSetup()
             _account = new NEP6Account(_wallet, _hash);
         }
 
+        [TestMethod]
+        public void TestChangePassword()
+        {
+            _account = new NEP6Account(_wallet, _hash, _nep2);
+            _account.ChangePasswordPrepare("b", "Satoshi").Should().BeTrue();
+            _account.ChangePasswordCommit();
+            _account.Contract = new Contract();
+            _wallet.Unlock("Satoshi");
+            _account.ChangePasswordPrepare("b", "Satoshi").Should().BeFalse();
+            _account.ChangePasswordPrepare("Satoshi", "b").Should().BeTrue();
+            _account.ChangePasswordCommit();
+            _account.VerifyPassword("b").Should().BeTrue();
+            _account.ChangePasswordPrepare("b", "Satoshi").Should().BeTrue();
+            _account.ChangePasswordCommit();
+            _account.ChangePasswordPrepare("Satoshi", "b").Should().BeTrue();
+            _account.ChangePasswordRoolback();
+            _account.VerifyPassword("Satoshi").Should().BeTrue();
+            _wallet.Lock();
+        }
+
         [TestMethod]
         public void TestConstructorWithNep2Key()
         {
diff --git a/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Wallet.cs b/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Wallet.cs
index dcf436e993..0e8ac02e05 100644
--- a/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Wallet.cs
+++ b/tests/neo.UnitTests/Wallets/NEP6/UT_NEP6Wallet.cs
@@ -70,6 +70,26 @@ public void TestCleanUp()
             if (Directory.Exists(rootPath)) Directory.Delete(rootPath);
         }
 
+        [TestMethod]
+        public void TestChangePassword()
+        {
+            JObject wallet = new JObject();
+            wallet["name"] = "name";
+            wallet["version"] = new System.Version("3.0").ToString();
+            wallet["scrypt"] = new ScryptParameters(0, 0, 0).ToJson();
+            wallet["accounts"] = new JArray();
+            wallet["extra"] = new JObject();
+            File.WriteAllText(wPath, wallet.ToString());
+            uut = new NEP6Wallet(wPath);
+            uut.Unlock("123");
+            uut.CreateAccount(keyPair.PrivateKey);
+            uut.ChangePassword("456", "123").Should().BeFalse();
+            uut.ChangePassword("123", "456").Should().BeTrue();
+            uut.VerifyPassword("456").Should().BeTrue();
+            uut.ChangePassword("456", "123").Should().BeTrue();
+            uut.Lock();
+        }
+
         [TestMethod]
         public void TestConstructorWithPathAndName()
         {