From b5b05a969cab2e2c16ad589d7988f5bbf2ce3949 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Thu, 9 Jan 2020 17:27:12 +0300
Subject: [PATCH] keys: make NEP2Decrypt return a PrivateKey rather than WIF

There is no point in encoding the output of this function in a WIF format,
most of the users actually want the real key and those who need a WIF can
easily get if from the key (and it's simpler than getting the key from the
WIF).

It also fixes a severe bug in NEP2Decrypt, base58 decoding errors were not
processed correctly.
---
 pkg/consensus/consensus.go   |  9 ++-------
 pkg/crypto/keys/nep2.go      | 16 ++++++++--------
 pkg/crypto/keys/nep2_test.go |  7 +------
 pkg/wallet/account.go        |  4 ++--
 4 files changed, 13 insertions(+), 23 deletions(-)

diff --git a/pkg/consensus/consensus.go b/pkg/consensus/consensus.go
index 8faf2c080..5f8999644 100644
--- a/pkg/consensus/consensus.go
+++ b/pkg/consensus/consensus.go
@@ -14,7 +14,6 @@ import (
 	"github.com/CityOfZion/neo-go/pkg/smartcontract"
 	"github.com/CityOfZion/neo-go/pkg/util"
 	"github.com/CityOfZion/neo-go/pkg/vm/opcode"
-	"github.com/CityOfZion/neo-go/pkg/wallet"
 	"github.com/nspcc-dev/dbft"
 	"github.com/nspcc-dev/dbft/block"
 	"github.com/nspcc-dev/dbft/crypto"
@@ -178,16 +177,12 @@ func (s *service) validatePayload(p *Payload) bool {
 }
 
 func getKeyPair(cfg *config.WalletConfig) (crypto.PrivateKey, crypto.PublicKey) {
-	acc, err := wallet.DecryptAccount(cfg.Path, cfg.Password)
+	// TODO: replace with wallet opening from the given path (#588)
+	key, err := keys.NEP2Decrypt(cfg.Path, cfg.Password)
 	if err != nil {
 		return nil, nil
 	}
 
-	key := acc.PrivateKey()
-	if key == nil {
-		return nil, nil
-	}
-
 	return &privateKey{PrivateKey: key}, &publicKey{PublicKey: key.PublicKey()}
 }
 
diff --git a/pkg/crypto/keys/nep2.go b/pkg/crypto/keys/nep2.go
index abe478490..1f49298d2 100644
--- a/pkg/crypto/keys/nep2.go
+++ b/pkg/crypto/keys/nep2.go
@@ -77,13 +77,13 @@ func NEP2Encrypt(priv *PrivateKey, passphrase string) (s string, err error) {
 
 // NEP2Decrypt decrypts an encrypted key using a given passphrase
 // under the NEP-2 standard.
-func NEP2Decrypt(key, passphrase string) (s string, err error) {
+func NEP2Decrypt(key, passphrase string) (*PrivateKey, error) {
 	b, err := base58.CheckDecode(key)
 	if err != nil {
-		return s, nil
+		return nil, err
 	}
 	if err := validateNEP2Format(b); err != nil {
-		return s, err
+		return nil, err
 	}
 
 	addrHash := b[3:7]
@@ -91,7 +91,7 @@ func NEP2Decrypt(key, passphrase string) (s string, err error) {
 	phraseNorm := norm.NFC.Bytes([]byte(passphrase))
 	derivedKey, err := scrypt.Key(phraseNorm, addrHash, n, r, p, keyLen)
 	if err != nil {
-		return s, err
+		return nil, err
 	}
 
 	derivedKey1 := derivedKey[:32]
@@ -100,7 +100,7 @@ func NEP2Decrypt(key, passphrase string) (s string, err error) {
 
 	decrypted, err := aesDecrypt(encryptedBytes, derivedKey2)
 	if err != nil {
-		return s, err
+		return nil, err
 	}
 
 	privBytes := xor(decrypted, derivedKey1)
@@ -108,14 +108,14 @@ func NEP2Decrypt(key, passphrase string) (s string, err error) {
 	// Rebuild the private key.
 	privKey, err := NewPrivateKeyFromBytes(privBytes)
 	if err != nil {
-		return s, err
+		return nil, err
 	}
 
 	if !compareAddressHash(privKey, addrHash) {
-		return s, errors.New("password mismatch")
+		return nil, errors.New("password mismatch")
 	}
 
-	return privKey.WIF(), nil
+	return privKey, nil
 }
 
 func compareAddressHash(priv *PrivateKey, inhash []byte) bool {
diff --git a/pkg/crypto/keys/nep2_test.go b/pkg/crypto/keys/nep2_test.go
index 4c5e86c23..2a67b9b26 100644
--- a/pkg/crypto/keys/nep2_test.go
+++ b/pkg/crypto/keys/nep2_test.go
@@ -27,18 +27,13 @@ func TestNEP2Encrypt(t *testing.T) {
 
 func TestNEP2Decrypt(t *testing.T) {
 	for _, testCase := range keytestcases.Arr {
-
-		privKeyString, err := NEP2Decrypt(testCase.EncryptedWif, testCase.Passphrase)
+		privKey, err := NEP2Decrypt(testCase.EncryptedWif, testCase.Passphrase)
 		if testCase.Invalid {
 			assert.Error(t, err)
 			continue
 		}
 
 		assert.Nil(t, err)
-
-		privKey, err := NewPrivateKeyFromWIF(privKeyString)
-		assert.Nil(t, err)
-
 		assert.Equal(t, testCase.PrivateKey, privKey.String())
 
 		wif := privKey.WIF()
diff --git a/pkg/wallet/account.go b/pkg/wallet/account.go
index 02fb52b75..dc8e75b90 100644
--- a/pkg/wallet/account.go
+++ b/pkg/wallet/account.go
@@ -63,11 +63,11 @@ func NewAccount() (*Account, error) {
 // DecryptAccount decrypts the encryptedWIF with the given passphrase and
 // return the decrypted Account.
 func DecryptAccount(encryptedWIF, passphrase string) (*Account, error) {
-	wif, err := keys.NEP2Decrypt(encryptedWIF, passphrase)
+	key, err := keys.NEP2Decrypt(encryptedWIF, passphrase)
 	if err != nil {
 		return nil, err
 	}
-	return NewAccountFromWIF(wif)
+	return newAccountFromPrivateKey(key), nil
 }
 
 // Encrypt encrypts the wallet's PrivateKey with the given passphrase