From 70df36ff3c104f80953e05d9823733467a432e9f Mon Sep 17 00:00:00 2001
From: Lazar <12626340+Lazar955@users.noreply.github.com>
Date: Thu, 19 Dec 2024 12:25:45 +0100
Subject: [PATCH] feat: list all keys in store (#234)

Adds `eotsd ls` command which prints all keys from DB

[closes](https://github.com/babylonlabs-io/finality-provider/issues/229)
---
 CHANGELOG.md                         |  1 +
 eotsmanager/cmd/eotsd/daemon/keys.go | 74 ++++++++++++++++++++++++++++
 eotsmanager/cmd/eotsd/daemon/root.go |  1 +
 eotsmanager/localmanager.go          |  4 ++
 eotsmanager/store/eotsstore.go       | 28 +++++++++++
 eotsmanager/store/eotsstore_test.go  | 51 +++++++++++++++++++
 itest/e2e_test.go                    | 49 ++++++++++++++++++
 itest/test_manager.go                |  2 +
 8 files changed, 210 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index df00a257..616cf01a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -49,6 +49,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 * [#228](https://github.com/babylonlabs-io/finality-provider/pull/228) Save key name mapping in eotsd import commands
 * [#227](https://github.com/babylonlabs-io/finality-provider/pull/227) Fix FP submission loop
 * [#226](https://github.com/babylonlabs-io/finality-provider/pull/226) Update local fp before register
+* [#234](https://github.com/babylonlabs-io/finality-provider/pull/234) eotsd ls command
 
 ## v0.13.1
 
diff --git a/eotsmanager/cmd/eotsd/daemon/keys.go b/eotsmanager/cmd/eotsd/daemon/keys.go
index ceead6f6..150147c4 100644
--- a/eotsmanager/cmd/eotsd/daemon/keys.go
+++ b/eotsmanager/cmd/eotsd/daemon/keys.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"io"
 
 	"github.com/babylonlabs-io/babylon/types"
@@ -124,6 +125,79 @@ func saveKeyNameMapping(cmd *cobra.Command, keyName string) (*types.BIP340PubKey
 	return eotsPk, nil
 }
 
+// CommandPrintAllKeys prints all EOTS keys
+func CommandPrintAllKeys() *cobra.Command {
+	var cmd = &cobra.Command{
+		Use:     "list",
+		Aliases: []string{"ls"},
+		Short:   "Print all EOTS key names and public keys mapping from database.",
+		Example: `eotsd list --home=/path/to/cfg`,
+		Args:    cobra.NoArgs,
+		RunE:    runCommandPrintAllKeys,
+	}
+
+	cmd.Flags().String(flags.FlagHome, config.DefaultEOTSDir, "The path to the eotsd home directory")
+
+	return cmd
+}
+
+func runCommandPrintAllKeys(cmd *cobra.Command, _ []string) error {
+	eotsKeys, err := getAllEOTSKeys(cmd)
+	if err != nil {
+		return err
+	}
+
+	for keyName, key := range eotsKeys {
+		pk, err := schnorr.ParsePubKey(key)
+		if err != nil {
+			return err
+		}
+		eotsPk := types.NewBIP340PubKeyFromBTCPK(pk)
+		cmd.Printf("Key Name: %s, EOTS PK: %s\n", keyName, eotsPk.MarshalHex())
+	}
+
+	return nil
+}
+
+func getAllEOTSKeys(cmd *cobra.Command) (map[string][]byte, error) {
+	homePath, err := getHomePath(cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	// Load configuration
+	cfg, err := config.LoadConfig(homePath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load config: %w", err)
+	}
+
+	// Setup logger
+	logger, err := log.NewRootLoggerWithFile(config.LogFile(homePath), cfg.LogLevel)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load the logger: %w", err)
+	}
+
+	// Get database backend
+	dbBackend, err := cfg.DatabaseConfig.GetDBBackend()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create db backend: %w", err)
+	}
+	defer dbBackend.Close()
+
+	// Create EOTS manager
+	eotsManager, err := eotsmanager.NewLocalEOTSManager(homePath, cfg.KeyringBackend, dbBackend, logger)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create EOTS manager: %w", err)
+	}
+
+	res, err := eotsManager.ListEOTSKeys()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get keys from db: %w", err)
+	}
+
+	return res, nil
+}
+
 func printFromKey(cmd *cobra.Command, keyName string, eotsPk *types.BIP340PubKey) error {
 	clientCtx, err := client.GetClientQueryContext(cmd)
 	if err != nil {
diff --git a/eotsmanager/cmd/eotsd/daemon/root.go b/eotsmanager/cmd/eotsd/daemon/root.go
index 590b83ee..de200793 100644
--- a/eotsmanager/cmd/eotsd/daemon/root.go
+++ b/eotsmanager/cmd/eotsd/daemon/root.go
@@ -25,6 +25,7 @@ func NewRootCmd() *cobra.Command {
 		NewKeysCmd(),
 		NewStartCmd(),
 		version.CommandVersion("eotsd"),
+		CommandPrintAllKeys(),
 	)
 
 	return rootCmd
diff --git a/eotsmanager/localmanager.go b/eotsmanager/localmanager.go
index 3c5a0feb..28ba6b36 100644
--- a/eotsmanager/localmanager.go
+++ b/eotsmanager/localmanager.go
@@ -369,3 +369,7 @@ func (lm *LocalEOTSManager) keyExists(name string) bool {
 	_, err := lm.kr.Key(name)
 	return err == nil
 }
+
+func (lm *LocalEOTSManager) ListEOTSKeys() (map[string][]byte, error) {
+	return lm.es.GetAllEOTSKeyNames()
+}
diff --git a/eotsmanager/store/eotsstore.go b/eotsmanager/store/eotsstore.go
index c197db67..1ced1a27 100644
--- a/eotsmanager/store/eotsstore.go
+++ b/eotsmanager/store/eotsstore.go
@@ -106,6 +106,34 @@ func (s *EOTSStore) GetEOTSKeyName(pk []byte) (string, error) {
 	return keyName, nil
 }
 
+// GetAllEOTSKeyNames retrieves all keys and values.
+// Returns keyName -> btcPK
+func (s *EOTSStore) GetAllEOTSKeyNames() (map[string][]byte, error) {
+	result := make(map[string][]byte)
+
+	err := s.db.View(func(tx kvdb.RTx) error {
+		eotsBucket := tx.ReadBucket(eotsBucketName)
+		if eotsBucket == nil {
+			return ErrCorruptedEOTSDb
+		}
+
+		return eotsBucket.ForEach(func(k, v []byte) error {
+			if k == nil || v == nil {
+				return fmt.Errorf("encountered invalid key or value in bucket")
+			}
+			result[string(v)] = k
+
+			return nil
+		})
+	}, func() {})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return result, nil
+}
+
 func (s *EOTSStore) SaveSignRecord(
 	height uint64,
 	chainID []byte,
diff --git a/eotsmanager/store/eotsstore_test.go b/eotsmanager/store/eotsstore_test.go
index bf160a18..244533d5 100644
--- a/eotsmanager/store/eotsstore_test.go
+++ b/eotsmanager/store/eotsstore_test.go
@@ -125,3 +125,54 @@ func FuzzSignStore(f *testing.F) {
 		require.False(t, found)
 	})
 }
+
+// FuzzListKeysEOTSStore tests getting all keys from store
+func FuzzListKeysEOTSStore(f *testing.F) {
+	testutil.AddRandomSeedsToFuzzer(f, 10)
+	f.Fuzz(func(t *testing.T, seed int64) {
+		t.Parallel()
+		r := rand.New(rand.NewSource(seed))
+
+		homePath := t.TempDir()
+		cfg := config.DefaultDBConfigWithHomePath(homePath)
+
+		dbBackend, err := cfg.GetDBBackend()
+		require.NoError(t, err)
+
+		vs, err := store.NewEOTSStore(dbBackend)
+		require.NoError(t, err)
+
+		defer func() {
+			dbBackend.Close()
+		}()
+
+		expected := make(map[string][]byte)
+		for i := 0; i < r.Intn(10); i++ {
+			expectedKeyName := testutil.GenRandomHexStr(r, 10)
+			_, btcPk, err := datagen.GenRandomBTCKeyPair(r)
+			require.NoError(t, err)
+			expected[expectedKeyName] = schnorr.SerializePubKey(btcPk)
+
+			err = vs.AddEOTSKeyName(
+				btcPk,
+				expectedKeyName,
+			)
+			require.NoError(t, err)
+		}
+
+		keys, err := vs.GetAllEOTSKeyNames()
+		require.NoError(t, err)
+
+		for keyName, btcPk := range expected {
+			gotBtcPk, ok := keys[keyName]
+			require.True(t, ok)
+
+			parsedGot, err := schnorr.ParsePubKey(gotBtcPk)
+			require.NoError(t, err)
+			parsedExpected, err := schnorr.ParsePubKey(btcPk)
+			require.NoError(t, err)
+
+			require.Equal(t, parsedExpected, parsedGot)
+		}
+	})
+}
diff --git a/itest/e2e_test.go b/itest/e2e_test.go
index 01b8bf0f..01ce10eb 100644
--- a/itest/e2e_test.go
+++ b/itest/e2e_test.go
@@ -1,10 +1,15 @@
 package e2etest
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
+	eotscmd "github.com/babylonlabs-io/finality-provider/eotsmanager/cmd/eotsd/daemon"
+	eotscfg "github.com/babylonlabs-io/finality-provider/eotsmanager/config"
 	"github.com/babylonlabs-io/finality-provider/finality-provider/store"
+	"github.com/btcsuite/btcd/btcec/v2/schnorr"
+	"github.com/jessevdk/go-flags"
 	"log"
 	"math/rand"
 	"os"
@@ -352,3 +357,47 @@ func TestRemoveMerkleProofsCmd(t *testing.T) {
 		return errors.Is(err, store.ErrPubRandProofNotFound)
 	}, eventuallyWaitTimeOut, eventuallyPollTime)
 }
+
+func TestPrintEotsCmd(t *testing.T) {
+	tm := StartManager(t)
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+	defer tm.Stop(t)
+
+	expected := make(map[string]string)
+	for i := 0; i < r.Intn(10); i++ {
+		eotsKeyName := fmt.Sprintf("eots-key-%s", datagen.GenRandomHexStr(r, 4))
+		ekey, err := tm.EOTSClient.CreateKey(eotsKeyName, passphrase, hdPath)
+		require.NoError(t, err)
+		pk, err := schnorr.ParsePubKey(ekey)
+		require.NoError(t, err)
+		expected[eotsKeyName] = bbntypes.NewBIP340PubKeyFromBTCPK(pk).MarshalHex()
+	}
+
+	tm.EOTSServerHandler.Stop()
+
+	cmd := eotscmd.CommandPrintAllKeys()
+
+	defaultConfig := eotscfg.DefaultConfigWithHomePath(tm.EOTSHomeDir)
+	fileParser := flags.NewParser(defaultConfig, flags.Default)
+	err := flags.NewIniParser(fileParser).WriteFile(eotscfg.CfgFile(tm.EOTSHomeDir), flags.IniIncludeDefaults)
+	require.NoError(t, err)
+
+	cmd.SetArgs([]string{
+		"--home=" + tm.EOTSHomeDir,
+	})
+
+	var outputBuffer bytes.Buffer
+	cmd.SetOut(&outputBuffer)
+	cmd.SetErr(&outputBuffer)
+
+	err = cmd.Execute()
+	require.NoError(t, err)
+
+	output := outputBuffer.String()
+	t.Logf("Captured output: %s", output)
+
+	for keyName, eotsPK := range expected {
+		require.Contains(t, output, keyName)
+		require.Contains(t, output, eotsPK)
+	}
+}
diff --git a/itest/test_manager.go b/itest/test_manager.go
index 20319795..fc428a54 100644
--- a/itest/test_manager.go
+++ b/itest/test_manager.go
@@ -59,6 +59,7 @@ var (
 type TestManager struct {
 	Wg                sync.WaitGroup
 	EOTSServerHandler *EOTSServerHandler
+	EOTSHomeDir       string
 	FpConfig          *fpcfg.Config
 	Fps               []*service.FinalityProviderApp
 	EOTSClient        *client.EOTSManagerGRpcClient
@@ -138,6 +139,7 @@ func StartManager(t *testing.T) *TestManager {
 
 	tm := &TestManager{
 		EOTSServerHandler: eh,
+		EOTSHomeDir:       eotsHomeDir,
 		FpConfig:          cfg,
 		EOTSClient:        eotsCli,
 		BBNClient:         bc,