From 5e6737fc389db5aec0ebedb99230478b4a49dc09 Mon Sep 17 00:00:00 2001
From: Alrighttt <mzlot555@gmail.com>
Date: Mon, 1 Jul 2019 19:03:52 +0200
Subject: [PATCH] add additional param to importprivkey for other WIF formats

---
 src/key_io.cpp         | 16 ++++++++++++++++
 src/key_io.h           |  1 +
 src/rpc/client.cpp     |  1 +
 src/wallet/rpcdump.cpp | 12 ++++++++++--
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/key_io.cpp b/src/key_io.cpp
index 014159d65..dd4176fee 100644
--- a/src/key_io.cpp
+++ b/src/key_io.cpp
@@ -207,6 +207,22 @@ CKey DecodeSecret(const std::string& str)
     return key;
 }
 
+CKey DecodeCustomSecret(const std::string& str, uint8_t secret_key)
+{
+    CKey key;
+    std::vector<unsigned char> data;
+    if (DecodeBase58Check(str, data)) {
+        const std::vector<unsigned char>& privkey_prefix = std::vector<unsigned char>(1, secret_key);
+        if ((data.size() == 32 + privkey_prefix.size() || (data.size() == 33 + privkey_prefix.size() && data.back() == 1)) &&
+            std::equal(privkey_prefix.begin(), privkey_prefix.end(), data.begin())) {
+            bool compressed = data.size() == 33 + privkey_prefix.size();
+            key.Set(data.begin() + privkey_prefix.size(), data.begin() + privkey_prefix.size() + 32, compressed);
+        }
+    }
+    memory_cleanse(data.data(), data.size());
+    return key;
+}
+
 std::string EncodeSecret(const CKey& key)
 {
     assert(key.IsValid());
diff --git a/src/key_io.h b/src/key_io.h
index 72823d57e..013469ab6 100644
--- a/src/key_io.h
+++ b/src/key_io.h
@@ -17,6 +17,7 @@
 #include <string>
 
 CKey DecodeSecret(const std::string& str);
+CKey DecodeCustomSecret(const std::string& str, uint8_t secret_key);
 std::string EncodeSecret(const CKey& key);
 std::string EncodeCustomSecret(const CKey& key,uint8_t secret_key);
 
diff --git a/src/rpc/client.cpp b/src/rpc/client.cpp
index 3f0f9dea8..665887b3b 100644
--- a/src/rpc/client.cpp
+++ b/src/rpc/client.cpp
@@ -109,6 +109,7 @@ static const CRPCConvertParam vRPCConvertParams[] =
     { "lockunspent", 1 },
     { "importprivkey", 2 },
     { "importprivkey", 3 },
+    { "importprivkey", 4 },
     { "importaddress", 2 },
     { "verifychain", 0 },
     { "verifychain", 1 },
diff --git a/src/wallet/rpcdump.cpp b/src/wallet/rpcdump.cpp
index 05552b50b..10048c79e 100644
--- a/src/wallet/rpcdump.cpp
+++ b/src/wallet/rpcdump.cpp
@@ -94,7 +94,7 @@ UniValue importprivkey(const UniValue& params, bool fHelp)
     if (!EnsureWalletIsAvailable(fHelp))
         return NullUniValue;
 
-    if (fHelp || params.size() < 1 || params.size() > 4)
+    if (fHelp || params.size() < 1 || params.size() > 5)
         throw runtime_error(
             "importprivkey \"komodoprivkey\" ( \"label\" rescan height)\n"
             "\nAdds a private key (as returned by dumpprivkey) to your wallet.\n"
@@ -126,6 +126,7 @@ UniValue importprivkey(const UniValue& params, bool fHelp)
     string strSecret = params[0].get_str();
     string strLabel = "";
     int32_t height = 0;
+    uint8_t secret_key = 0;
     if (params.size() > 1)
         strLabel = params[1].get_str();
 
@@ -136,10 +137,17 @@ UniValue importprivkey(const UniValue& params, bool fHelp)
     if ( fRescan && params.size() == 4 )
         height = params[3].get_int();
 
+    CKey key = DecodeSecret(strSecret);
+
+    if (params.size() > 4)
+    {
+        secret_key = params[4].get_int();
+        CKey key = DecodeCustomSecret(strSecret, secret_key);
+    }
+
     if ( height < 0 || height > chainActive.Height() )
         throw JSONRPCError(RPC_WALLET_ERROR, "Rescan height is out of range.");
     
-    CKey key = DecodeSecret(strSecret);
     if (!key.IsValid()) throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid private key encoding");
 
     CPubKey pubkey = key.GetPubKey();