Keymanager fixes for bad file writes (#12284)

* WIP changes for keymanager

* WIP fix

* WIP needs unit tests

* fixing order

* adding unit test

* fixing linter

* updating unit tests and creating more reusable functions

* making accountStore copy method part of struct

* Update validator/keymanager/local/delete_test.go

* Update validator/keymanager/local/delete.go

* Update validator/keymanager/local/delete.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* Update validator/keymanager/local/import.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* Update validator/keymanager/local/delete.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* Update validator/keymanager/local/delete.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* Update validator/keymanager/local/import.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* Update validator/keymanager/local/import.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* Update validator/keymanager/local/delete.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* addressing suggestion of not reinitializing from reading the file but instead update the information based on memory on hand

* Update validator/accounts/wallet_create.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

* adding changes based on suggestions

* making logs more consistent

* fixing linting

---------

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>
Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

validator/accounts/testing/mock.go

@@ -25,6 +25,7 @@ type Wallet struct {
 	WalletPassword    string
 	UnlockAccounts    bool
 	lock              sync.RWMutex
+	HasWriteFileError bool
 }
 
 // AccountNames --
@@ -57,6 +58,11 @@ func (w *Wallet) Password() string {
 func (w *Wallet) WriteFileAtPath(_ context.Context, pathName, fileName string, data []byte) error {
 	w.lock.Lock()
 	defer w.lock.Unlock()
+	if w.HasWriteFileError {
+		// reset the flag to not contaminate other tests
+		w.HasWriteFileError = false
+		return errors.New("could not write keystore file for accounts")
+	}
 	if w.Files[pathName] == nil {
 		w.Files[pathName] = make(map[string][]byte)
 	}

validator/keymanager/local/BUILD.bazel

@@ -66,6 +66,7 @@ go_test(
         "//testing/require:go_default_library",
         "//validator/accounts/testing:go_default_library",
         "//validator/keymanager:go_default_library",
+        "@com_github_ethereum_go_ethereum//common/hexutil:go_default_library",
         "@com_github_google_uuid//:go_default_library",
         "@com_github_sirupsen_logrus//hooks/test:go_default_library",
         "@com_github_wealdtech_go_eth2_wallet_encryptor_keystorev4//:go_default_library",

validator/keymanager/local/delete.go

@@ -3,10 +3,7 @@ package local
 import (
 	"bytes"
 	"context"
-	"encoding/json"
-	"fmt"
 
-	"github.com/pkg/errors"
 	fieldparams "github.com/prysmaticlabs/prysm/v4/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v4/encoding/bytesutil"
 	ethpbservice "github.com/prysmaticlabs/prysm/v4/proto/eth/service"
@@ -16,17 +13,22 @@ import (
 // DeleteKeystores takes in public keys and removes the accounts from the wallet.
 // This includes their disk keystore and cached keystore, but maintains the slashing
 // protection history in the database.
+// 1) Copy the in memory keystore
+// 2) Delete the keys from copied in memory keystore
+// 3) Save the copy to disk
+// 4) Reinitialize account store and updating the keymanager
+// 5) Return API response
 func (km *Keymanager) DeleteKeystores(
 	ctx context.Context, publicKeys [][]byte,
 ) ([]*ethpbservice.DeletedKeystoreStatus, error) {
 	// Check for duplicate keys and filter them out.
 	trackedPublicKeys := make(map[[fieldparams.BLSPubkeyLength]byte]bool)
 	statuses := make([]*ethpbservice.DeletedKeystoreStatus, 0, len(publicKeys))
-	var store *AccountsKeystoreRepresentation
-	var err error
 	deletedKeys := make([][]byte, 0, len(publicKeys))
+	// 1) Copy the in memory keystore
+	storeCopy := km.accountsStore.Copy()
 	for _, publicKey := range publicKeys {
-		// Check if the key in the request is a duplicate.
+		// Check if the key in the request is a duplicate or not found
 		if _, ok := trackedPublicKeys[bytesutil.ToBytes48(publicKey)]; ok {
 			statuses = append(statuses, &ethpbservice.DeletedKeystoreStatus{
 				Status: ethpbservice.DeletedKeystoreStatus_NOT_ACTIVE,
@@ -35,7 +37,7 @@ func (km *Keymanager) DeleteKeystores(
 		}
 		var index int
 		var found bool
-		for j, pubKey := range km.accountsStore.PublicKeys {
+		for j, pubKey := range storeCopy.PublicKeys {
 			if bytes.Equal(pubKey, publicKey) {
 				index = j
 				found = true
@@ -48,14 +50,11 @@ func (km *Keymanager) DeleteKeystores(
 			})
 			continue
 		}
-		deletedPublicKey := km.accountsStore.PublicKeys[index]
+		// 2) Delete the keys from copied in memory keystore
+		deletedPublicKey := storeCopy.PublicKeys[index]
 		deletedKeys = append(deletedKeys, deletedPublicKey)
-		km.accountsStore.PrivateKeys = append(km.accountsStore.PrivateKeys[:index], km.accountsStore.PrivateKeys[index+1:]...)
-		km.accountsStore.PublicKeys = append(km.accountsStore.PublicKeys[:index], km.accountsStore.PublicKeys[index+1:]...)
-		store, err = km.CreateAccountsKeystore(ctx, km.accountsStore.PrivateKeys, km.accountsStore.PublicKeys)
-		if err != nil {
-			return nil, errors.Wrap(err, "could not rewrite accounts keystore")
-		}
+		storeCopy.PrivateKeys = append(storeCopy.PrivateKeys[:index], storeCopy.PrivateKeys[index+1:]...)
+		storeCopy.PublicKeys = append(storeCopy.PublicKeys[:index], storeCopy.PublicKeys[index+1:]...)
 		statuses = append(statuses, &ethpbservice.DeletedKeystoreStatus{
 			Status: ethpbservice.DeletedKeystoreStatus_DELETED,
 		})
@@ -64,35 +63,14 @@ func (km *Keymanager) DeleteKeystores(
 	if len(deletedKeys) == 0 {
 		return statuses, nil
 	}
-	var deletedKeysStr string
-	for i, k := range deletedKeys {
-		if i == 0 {
-			deletedKeysStr += fmt.Sprintf("%#x", bytesutil.Trunc(k))
-		} else if i == len(deletedKeys)-1 {
-			deletedKeysStr += fmt.Sprintf("%#x", bytesutil.Trunc(k))
-		} else {
-			deletedKeysStr += fmt.Sprintf(",%#x", bytesutil.Trunc(k))
-		}
-	}
-
-	log.WithFields(logrus.Fields{
-		"publicKeys": deletedKeysStr,
-	}).Info("Successfully deleted validator key(s)")
-
-	// Write the encoded keystore.
-	encoded, err := json.MarshalIndent(store, "", "\t")
-	if err != nil {
+	// 3 & 4) save to disk and re-initializes keystore
+	if err := km.SaveStoreAndReInitialize(ctx, storeCopy); err != nil {
 		return nil, err
 	}
-	if err := km.wallet.WriteFileAtPath(ctx, AccountsPath, AccountsKeystoreFileName, encoded); err != nil {
-		return nil, errors.Wrap(err, "could not write keystore file for accounts")
-	}
-	err = km.initializeKeysCachesFromKeystore()
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to initialize key caches")
-	}
+
 	log.WithFields(logrus.Fields{
-		"publicKeys": deletedKeysStr,
+		"publicKeys": CreatePrintoutOfKeys(deletedKeys),
 	}).Info("Successfully deleted validator key(s)")
+	// 5) Return API response
 	return statuses, nil
 }

validator/keymanager/local/delete_test.go

@@ -50,6 +50,17 @@ func TestLocalKeymanager_DeleteKeystores(t *testing.T) {
 		require.Equal(t, ethpbservice.DeletedKeystoreStatus_NOT_FOUND, statuses[0].Status)
 		require.Equal(t, ethpbservice.DeletedKeystoreStatus_NOT_FOUND, statuses[1].Status)
 	})
+	t.Run("file write errors should not lead to updated local keystore or cache", func(t *testing.T) {
+		wallet.HasWriteFileError = true
+		accountToRemove := uint64(2)
+		accountPubKey := accounts[accountToRemove]
+		require.NotEqual(t, len(dr.accountsStore.PublicKeys), 0)
+		copyStore := dr.accountsStore.Copy()
+		statuses, err := dr.DeleteKeystores(ctx, [][]byte{accountPubKey[:]})
+		require.ErrorContains(t, "could not write keystore file for accounts", err)
+		require.Equal(t, len(statuses), 0)
+		require.DeepEqual(t, dr.accountsStore, copyStore)
+	})
 	t.Run("deletes properly", func(t *testing.T) {
 		accountToRemove := uint64(2)
 		accountPubKey := accounts[accountToRemove]
@@ -83,6 +94,7 @@ func TestLocalKeymanager_DeleteKeystores(t *testing.T) {
 		require.LogsContain(t, hook, fmt.Sprintf("%#x", bytesutil.Trunc(accountPubKey[:])))
 		require.LogsContain(t, hook, "Successfully deleted validator key(s)")
 	})
+
 	t.Run("returns NOT_ACTIVE status for duplicate public key in request", func(t *testing.T) {
 		accountToRemove := uint64(3)
 		accountPubKey := accounts[accountToRemove]
@@ -127,4 +139,5 @@ func TestLocalKeymanager_DeleteKeystores(t *testing.T) {
 		require.LogsContain(t, hook, fmt.Sprintf("%#x", bytesutil.Trunc(accountPubKey[:])))
 		require.LogsContain(t, hook, "Successfully deleted validator key(s)")
 	})
+
 }

validator/keymanager/local/import.go

@@ -3,7 +3,6 @@ package local
 import (
 	"context"
 	"encoding/hex"
-	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -13,10 +12,16 @@ import (
 	ethpbservice "github.com/prysmaticlabs/prysm/v4/proto/eth/service"
 	"github.com/prysmaticlabs/prysm/v4/validator/keymanager"
 	"github.com/schollz/progressbar/v3"
+	"github.com/sirupsen/logrus"
 	keystorev4 "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4"
 )
 
 // ImportKeystores into the local keymanager from an external source.
+// 1) Copy the in memory keystore
+// 2) Update copied keystore with new keys
+// 3) Save the copy to disk
+// 4) Reinitialize account store and updating the keymanager
+// 5) Return Statuses
 func (km *Keymanager) ImportKeystores(
 	ctx context.Context,
 	keystores []*keymanager.Keystore,
@@ -33,7 +38,13 @@ func (km *Keymanager) ImportKeystores(
 	keys := map[string]string{}
 	statuses := make([]*ethpbservice.ImportedKeystoreStatus, len(keystores))
 	var err error
-
+	// 1) Copy the in memory keystore
+	storeCopy := km.accountsStore.Copy()
+	importedKeys := make([][]byte, 0)
+	existingPubKeys := make(map[string]bool)
+	for i := 0; i < len(storeCopy.PrivateKeys); i++ {
+		existingPubKeys[string(storeCopy.PublicKeys[i])] = true
+	}
 	for i := 0; i < len(keystores); i++ {
 		var privKeyBytes []byte
 		var pubKeyBytes []byte
@@ -49,52 +60,67 @@ func (km *Keymanager) ImportKeystores(
 			log.Error(err)
 		}
 		// if key exists prior to being added then output log that duplicate key was found
-		if _, ok := keys[string(pubKeyBytes)]; ok {
+		_, isDuplicateInArray := keys[string(pubKeyBytes)]
+		_, isDuplicateInExisting := existingPubKeys[string(pubKeyBytes)]
+		if isDuplicateInArray || isDuplicateInExisting {
 			log.Warnf("Duplicate key in import will be ignored: %#x", pubKeyBytes)
 			statuses[i] = &ethpbservice.ImportedKeystoreStatus{
 				Status: ethpbservice.ImportedKeystoreStatus_DUPLICATE,
 			}
 			continue
 		}
+
 		keys[string(pubKeyBytes)] = string(privKeyBytes)
+		importedKeys = append(importedKeys, pubKeyBytes)
 		statuses[i] = &ethpbservice.ImportedKeystoreStatus{
 			Status: ethpbservice.ImportedKeystoreStatus_IMPORTED,
 		}
 	}
-	privKeys := make([][]byte, 0)
-	pubKeys := make([][]byte, 0)
+	if len(importedKeys) == 0 {
+		log.Warn("no keys were imported")
+		return statuses, nil
+	}
+	// 2) Update copied keystore with new keys,clear duplicates in existing set
+	// duplicates,errored ones are already skipped
 	for pubKey, privKey := range keys {
-		pubKeys = append(pubKeys, []byte(pubKey))
-		privKeys = append(privKeys, []byte(privKey))
+		storeCopy.PublicKeys = append(storeCopy.PublicKeys, []byte(pubKey))
+		storeCopy.PrivateKeys = append(storeCopy.PrivateKeys, []byte(privKey))
+	}
+	//3 & 4) save to disk and re-initializes keystore
+	if err := km.SaveStoreAndReInitialize(ctx, storeCopy); err != nil {
+		return nil, err
 	}
 
-	// Write the accounts to disk into a single keystore.
-	accountsKeystore, err := km.CreateAccountsKeystore(ctx, privKeys, pubKeys)
-	if err != nil {
-		return nil, err
-	}
-	encodedAccounts, err := json.MarshalIndent(accountsKeystore, "", "\t")
-	if err != nil {
-		return nil, err
-	}
-	if err := km.wallet.WriteFileAtPath(ctx, AccountsPath, AccountsKeystoreFileName, encodedAccounts); err != nil {
-		return nil, err
-	}
+	log.WithFields(logrus.Fields{
+		"publicKeys": CreatePrintoutOfKeys(importedKeys),
+	}).Info("Successfully imported validator key(s)")
+
+	// 5) Return Statuses
 	return statuses, nil
 }
 
 // ImportKeypairs directly into the keymanager.
 func (km *Keymanager) ImportKeypairs(ctx context.Context, privKeys, pubKeys [][]byte) error {
-	// Write the accounts to disk into a single keystore.
-	accountsKeystore, err := km.CreateAccountsKeystore(ctx, privKeys, pubKeys)
-	if err != nil {
-		return errors.Wrap(err, "could not import account keypairs")
+	if len(privKeys) != len(pubKeys) {
+		return fmt.Errorf(
+			"number of private keys and public keys is not equal: %d != %d", len(privKeys), len(pubKeys),
+		)
 	}
-	encodedAccounts, err := json.MarshalIndent(accountsKeystore, "", "\t")
-	if err != nil {
-		return errors.Wrap(err, "could not marshal accounts keystore into JSON")
+	// 1) Copy the in memory keystore
+	storeCopy := km.accountsStore.Copy()
+
+	// 2) Update store and remove duplicates
+	updateAccountsStoreKeys(storeCopy, privKeys, pubKeys)
+
+	// 3 & 4) save to disk and re-initializes keystore
+	if err := km.SaveStoreAndReInitialize(ctx, storeCopy); err != nil {
+		return err
 	}
-	return km.wallet.WriteFileAtPath(ctx, AccountsPath, AccountsKeystoreFileName, encodedAccounts)
+	// 5) verify if store was not updated
+	if len(km.accountsStore.PublicKeys) < len(storeCopy.PublicKeys) {
+		return fmt.Errorf("keys were not imported successfully, expected %d got %d", len(storeCopy.PublicKeys), len(km.accountsStore.PublicKeys))
+	}
+	return nil
 }
 
 // Retrieves the private key and public key from an EIP-2335 keystore file

validator/keymanager/local/import_test.go

@@ -6,13 +6,16 @@ import (
 	"strconv"
 	"testing"
 
+	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/google/uuid"
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls"
+	"github.com/prysmaticlabs/prysm/v4/encoding/bytesutil"
 	ethpbservice "github.com/prysmaticlabs/prysm/v4/proto/eth/service"
 	"github.com/prysmaticlabs/prysm/v4/testing/assert"
 	"github.com/prysmaticlabs/prysm/v4/testing/require"
 	mock "github.com/prysmaticlabs/prysm/v4/validator/accounts/testing"
 	"github.com/prysmaticlabs/prysm/v4/validator/keymanager"
+	logTest "github.com/sirupsen/logrus/hooks/test"
 	keystorev4 "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4"
 )
 
@@ -94,6 +97,7 @@ func TestLocalKeymanager_NoDuplicates(t *testing.T) {
 }
 
 func TestLocalKeymanager_ImportKeystores(t *testing.T) {
+	hook := logTest.NewGlobal()
 	ctx := context.Background()
 	// Setup the keymanager.
 	wallet := &mock.Wallet{
@@ -123,6 +127,7 @@ func TestLocalKeymanager_ImportKeystores(t *testing.T) {
 		for _, status := range statuses {
 			require.Equal(t, ethpbservice.ImportedKeystoreStatus_IMPORTED, status.Status)
 		}
+		require.LogsContain(t, hook, "Successfully imported validator key(s)")
 	})
 	t.Run("each imported keystore with a different password succeeds", func(t *testing.T) {
 		numKeystores := 5
@@ -143,6 +148,7 @@ func TestLocalKeymanager_ImportKeystores(t *testing.T) {
 		for _, status := range statuses {
 			require.Equal(t, ethpbservice.ImportedKeystoreStatus_IMPORTED, status.Status)
 		}
+		require.LogsContain(t, hook, "Successfully imported validator key(s)")
 	})
 	t.Run("some succeed, some fail to decrypt, some duplicated", func(t *testing.T) {
 		keystores := make([]*keymanager.Keystore, 0)
@@ -189,5 +195,70 @@ func TestLocalKeymanager_ImportKeystores(t *testing.T) {
 			fmt.Sprintf("incorrect password for key 0x%s", keystores[2].Pubkey),
 			statuses[2].Message,
 		)
+		b, err := hexutil.Decode("0x" + keystore1.Pubkey)
+		require.NoError(t, err)
+		require.LogsContain(t, hook, fmt.Sprintf("%#x", bytesutil.Trunc(b)))
+		require.LogsContain(t, hook, "Successfully imported validator key(s)")
+	})
+	t.Run("All fail or duplicated", func(t *testing.T) {
+		// First keystore is normal.
+		keystore1 := createRandomKeystore(t, password)
+		// First Import successfully
+		statuses, err := dr.ImportKeystores(
+			ctx,
+			[]*keymanager.Keystore{keystore1},
+			[]string{password},
+		)
+		require.NoError(t, err)
+		require.Equal(t, len(statuses), 1)
+
+		keystores := make([]*keymanager.Keystore, 0)
+		passwords := make([]string, 0)
+		// Second keystore is a duplicate of the first.
+		keystores = append(keystores, keystore1)
+		passwords = append(passwords, password)
+
+		// Third keystore has a wrong password.
+		keystore3 := createRandomKeystore(t, password)
+		keystores = append(keystores, keystore3)
+		passwords = append(passwords, "foobar")
+
+		statuses, err = dr.ImportKeystores(
+			ctx,
+			keystores,
+			passwords,
+		)
+		require.NoError(t, err)
+		require.Equal(t, len(keystores), len(statuses))
+		require.Equal(
+			t,
+			ethpbservice.ImportedKeystoreStatus_DUPLICATE,
+			statuses[0].Status,
+		)
+		require.Equal(
+			t,
+			ethpbservice.ImportedKeystoreStatus_ERROR,
+			statuses[1].Status,
+		)
+		require.Equal(
+			t,
+			fmt.Sprintf("incorrect password for key 0x%s", keystores[1].Pubkey),
+			statuses[1].Message,
+		)
+		require.LogsContain(t, hook, "no keys were imported")
+	})
+	t.Run("file write fails during import", func(t *testing.T) {
+		wallet.HasWriteFileError = true
+		copyStore := dr.accountsStore.Copy()
+		keystore1 := createRandomKeystore(t, password)
+		statuses, err := dr.ImportKeystores(
+			ctx,
+			[]*keymanager.Keystore{keystore1},
+			[]string{password},
+		)
+		require.ErrorContains(t, "could not write keystore file for accounts", err)
+		require.Equal(t, len(statuses), 0)
+		// local copy did not update due to bad file write
+		require.DeepEqual(t, dr.accountsStore, copyStore)
 	})
 }

validator/keymanager/local/keymanager.go

@@ -59,6 +59,14 @@ type accountStore struct {
 	PublicKeys  [][]byte `json:"public_keys"`
 }
 
+// Copy creates a deep copy of accountStore
+func (a *accountStore) Copy() *accountStore {
+	storeCopy := &accountStore{}
+	storeCopy.PrivateKeys = bytesutil.SafeCopy2dBytes(a.PrivateKeys)
+	storeCopy.PublicKeys = bytesutil.SafeCopy2dBytes(a.PublicKeys)
+	return storeCopy
+}
+
 // AccountsKeystoreRepresentation defines an internal Prysm representation
 // of validator accounts, encrypted according to the EIP-2334 standard.
 type AccountsKeystoreRepresentation struct {
@@ -256,55 +264,55 @@ func (km *Keymanager) initializeAccountKeystore(ctx context.Context) error {
 }
 
 // CreateAccountsKeystore creates a new keystore holding the provided keys.
-func (km *Keymanager) CreateAccountsKeystore(
+func (km *Keymanager) CreateAccountsKeystore(ctx context.Context, privateKeys [][]byte, publicKeys [][]byte) (*AccountsKeystoreRepresentation, error) {
+	if err := km.CreateOrUpdateInMemoryAccountsStore(ctx, privateKeys, publicKeys); err != nil {
+		return nil, err
+	}
+	return CreateAccountsKeystoreRepresentation(ctx, km.accountsStore, km.wallet.Password())
+}
+
+// SaveStoreAndReInitialize saves the store to disk and re-initializes the account keystore from file
+func (km *Keymanager) SaveStoreAndReInitialize(ctx context.Context, store *accountStore) error {
+	// Save the copy to disk
+	accountsKeystore, err := CreateAccountsKeystoreRepresentation(ctx, store, km.wallet.Password())
+	if err != nil {
+		return err
+
+	}
+	encodedAccounts, err := json.MarshalIndent(accountsKeystore, "", "\t")
+	if err != nil {
+		return err
+	}
+	if err := km.wallet.WriteFileAtPath(ctx, AccountsPath, AccountsKeystoreFileName, encodedAccounts); err != nil {
+		return err
+	}
+
+	// Reinitialize account store and cache
+	// This will update the in-memory information instead of reading from the file itself for safety concerns
+	km.accountsStore = store
+	err = km.initializeKeysCachesFromKeystore()
+	if err != nil {
+		return errors.Wrap(err, "failed to initialize keys caches")
+	}
+	return err
+}
+
+// CreateAccountsKeystoreRepresentation is a pure function that takes an accountStore and wallet password and returns the encrypted formatted json version for local writing.
+func CreateAccountsKeystoreRepresentation(
 	_ context.Context,
-	privateKeys, publicKeys [][]byte,
+	store *accountStore,
+	walletPW string,
 ) (*AccountsKeystoreRepresentation, error) {
 	encryptor := keystorev4.New()
 	id, err := uuid.NewRandom()
 	if err != nil {
 		return nil, err
 	}
-	if len(privateKeys) != len(publicKeys) {
-		return nil, fmt.Errorf(
-			"number of private keys and public keys is not equal: %d != %d", len(privateKeys), len(publicKeys),
-		)
-	}
-	if km.accountsStore == nil {
-		km.accountsStore = &accountStore{
-			PrivateKeys: privateKeys,
-			PublicKeys:  publicKeys,
-		}
-	} else {
-		existingPubKeys := make(map[string]bool)
-		existingPrivKeys := make(map[string]bool)
-		for i := 0; i < len(km.accountsStore.PrivateKeys); i++ {
-			existingPrivKeys[string(km.accountsStore.PrivateKeys[i])] = true
-			existingPubKeys[string(km.accountsStore.PublicKeys[i])] = true
-		}
-		// We append to the accounts store keys only
-		// if the private/secret key do not already exist, to prevent duplicates.
-		for i := 0; i < len(privateKeys); i++ {
-			sk := privateKeys[i]
-			pk := publicKeys[i]
-			_, privKeyExists := existingPrivKeys[string(sk)]
-			_, pubKeyExists := existingPubKeys[string(pk)]
-			if privKeyExists || pubKeyExists {
-				continue
-			}
-			km.accountsStore.PublicKeys = append(km.accountsStore.PublicKeys, pk)
-			km.accountsStore.PrivateKeys = append(km.accountsStore.PrivateKeys, sk)
-		}
-	}
-	err = km.initializeKeysCachesFromKeystore()
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to initialize keys caches")
-	}
-	encodedStore, err := json.MarshalIndent(km.accountsStore, "", "\t")
+	encodedStore, err := json.MarshalIndent(store, "", "\t")
 	if err != nil {
 		return nil, err
 	}
-	cryptoFields, err := encryptor.Encrypt(encodedStore, km.wallet.Password())
+	cryptoFields, err := encryptor.Encrypt(encodedStore, walletPW)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not encrypt accounts")
 	}
@@ -316,6 +324,51 @@ func (km *Keymanager) CreateAccountsKeystore(
 	}, nil
 }
 
+// CreateOrUpdateInMemoryAccountsStore will set or update the local accounts store and update the local cache.
+// This function DOES NOT save the accounts store to disk.
+func (km *Keymanager) CreateOrUpdateInMemoryAccountsStore(_ context.Context, privateKeys, publicKeys [][]byte) error {
+	if len(privateKeys) != len(publicKeys) {
+		return fmt.Errorf(
+			"number of private keys and public keys is not equal: %d != %d", len(privateKeys), len(publicKeys),
+		)
+	}
+	if km.accountsStore == nil {
+		km.accountsStore = &accountStore{
+			PrivateKeys: privateKeys,
+			PublicKeys:  publicKeys,
+		}
+	} else {
+		updateAccountsStoreKeys(km.accountsStore, privateKeys, publicKeys)
+	}
+	err := km.initializeKeysCachesFromKeystore()
+	if err != nil {
+		return errors.Wrap(err, "failed to initialize keys caches")
+	}
+	return nil
+}
+
+func updateAccountsStoreKeys(store *accountStore, privateKeys, publicKeys [][]byte) {
+	existingPubKeys := make(map[string]bool)
+	existingPrivKeys := make(map[string]bool)
+	for i := 0; i < len(store.PrivateKeys); i++ {
+		existingPrivKeys[string(store.PrivateKeys[i])] = true
+		existingPubKeys[string(store.PublicKeys[i])] = true
+	}
+	// We append to the accounts store keys only
+	// if the private/secret key do not already exist, to prevent duplicates.
+	for i := 0; i < len(privateKeys); i++ {
+		sk := privateKeys[i]
+		pk := publicKeys[i]
+		_, privKeyExists := existingPrivKeys[string(sk)]
+		_, pubKeyExists := existingPubKeys[string(pk)]
+		if privKeyExists || pubKeyExists {
+			continue
+		}
+		store.PublicKeys = append(store.PublicKeys, pk)
+		store.PrivateKeys = append(store.PrivateKeys, sk)
+	}
+}
+
 func (km *Keymanager) ListKeymanagerAccounts(ctx context.Context, cfg keymanager.ListKeymanagerAccountConfig) error {
 	au := aurora.NewAurora(true)
 	// We initialize the wallet's keymanager.
@@ -369,3 +422,17 @@ func (km *Keymanager) ListKeymanagerAccounts(ctx context.Context, cfg keymanager
 	fmt.Println("")
 	return nil
 }
+
+func CreatePrintoutOfKeys(keys [][]byte) string {
+	var keysStr string
+	for i, k := range keys {
+		if i == 0 {
+			keysStr += fmt.Sprintf("%#x", bytesutil.Trunc(k))
+		} else if i == len(keys)-1 {
+			keysStr += fmt.Sprintf("%#x", bytesutil.Trunc(k))
+		} else {
+			keysStr += fmt.Sprintf(",%#x", bytesutil.Trunc(k))
+		}
+	}
+	return keysStr
+}