From f86f73b503106e79ba4fe8a95419febaae0902d2 Mon Sep 17 00:00:00 2001
From: RafilxTenfen <rafaeltenfen.rt@gmail.com>
Date: Mon, 16 Dec 2024 23:32:07 -0300
Subject: [PATCH] feat: add save key on mapping after import of keys (#228)

Closes: #169
---
 CHANGELOG.md                         |  3 +-
 eotsmanager/cmd/eotsd/daemon/keys.go | 51 ++++++++++++++++++++++------
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2637c866..dcd83069 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -46,6 +46,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 * [#216](https://github.com/babylonlabs-io/finality-provider/pull/216) Add multiple fpd connecting to one eotsd in e2e tests
 * [#218](https://github.com/babylonlabs-io/finality-provider/pull/218) Prune used merkle proof
 * [#221](https://github.com/babylonlabs-io/finality-provider/pull/221) Cleanup TODOs
+* [#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
 
 ## v0.13.1
@@ -119,7 +120,7 @@ finality vote submission
 
 * [#117](https://github.com/babylonlabs-io/finality-provider/pull/117) Spec of
 commit public randomness
-* [#130](https://github.com/babylonlabs-io/finality-provider/pull/130) Finality 
+* [#130](https://github.com/babylonlabs-io/finality-provider/pull/130) Finality
 Provider operation documentation
 
 ### Bug Fixes
diff --git a/eotsmanager/cmd/eotsd/daemon/keys.go b/eotsmanager/cmd/eotsd/daemon/keys.go
index b7248a46..ceead6f6 100644
--- a/eotsmanager/cmd/eotsd/daemon/keys.go
+++ b/eotsmanager/cmd/eotsd/daemon/keys.go
@@ -50,53 +50,84 @@ func NewKeysCmd() *cobra.Command {
 		}
 
 		cmd.SetOut(oldOut)
-		return saveKeyNameMapping(cmd, args)
+		keyName := args[0]
+		eotsPk, err := saveKeyNameMapping(cmd, keyName)
+		if err != nil {
+			return err
+		}
+
+		return printFromKey(cmd, keyName, eotsPk)
 	}
 
+	saveKeyOnPostRun(keysCmd, "import")
+	saveKeyOnPostRun(keysCmd, "import-hex")
+
 	return keysCmd
 }
 
-func saveKeyNameMapping(cmd *cobra.Command, args []string) error {
+func saveKeyOnPostRun(cmd *cobra.Command, commandName string) {
+	subCmd := util.GetSubCommand(cmd, commandName)
+	if subCmd == nil {
+		panic(fmt.Sprintf("failed to find keys %s command", commandName))
+	}
+
+	subCmd.PostRunE = func(cmd *cobra.Command, args []string) error {
+		keyName := args[0]
+		_, err := saveKeyNameMapping(cmd, keyName)
+		return err
+	}
+}
+
+func saveKeyNameMapping(cmd *cobra.Command, keyName string) (*types.BIP340PubKey, error) {
 	clientCtx, err := client.GetClientQueryContext(cmd)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	keyName := args[0]
 
 	// Load configuration
 	cfg, err := config.LoadConfig(clientCtx.HomeDir)
 	if err != nil {
-		return fmt.Errorf("failed to load config: %w", err)
+		return nil, fmt.Errorf("failed to load config: %w", err)
 	}
 
 	// Setup logger
 	logger, err := log.NewRootLoggerWithFile(config.LogFile(clientCtx.HomeDir), cfg.LogLevel)
 	if err != nil {
-		return fmt.Errorf("failed to load the logger: %w", err)
+		return nil, fmt.Errorf("failed to load the logger: %w", err)
 	}
 
 	// Get database backend
 	dbBackend, err := cfg.DatabaseConfig.GetDBBackend()
 	if err != nil {
-		return fmt.Errorf("failed to create db backend: %w", err)
+		return nil, fmt.Errorf("failed to create db backend: %w", err)
 	}
 	defer dbBackend.Close()
 
 	// Create EOTS manager
 	eotsManager, err := eotsmanager.NewLocalEOTSManager(clientCtx.HomeDir, clientCtx.Keyring.Backend(), dbBackend, logger)
 	if err != nil {
-		return fmt.Errorf("failed to create EOTS manager: %w", err)
+		return nil, fmt.Errorf("failed to create EOTS manager: %w", err)
 	}
 
 	// Get the public key for the newly added key
 	eotsPk, err := eotsManager.LoadBIP340PubKeyFromKeyName(keyName)
 	if err != nil {
-		return fmt.Errorf("failed to get public key for key %s: %w", keyName, err)
+		return nil, fmt.Errorf("failed to get public key for key %s: %w", keyName, err)
 	}
 
 	// Save the public key to key name mapping
 	if err := eotsManager.SaveEOTSKeyName(eotsPk.MustToBTCPK(), keyName); err != nil {
-		return fmt.Errorf("failed to save key name mapping: %w", err)
+		return nil, fmt.Errorf("failed to save key name mapping: %w", err)
+	}
+
+	cmd.Printf("Successfully wrote key name %s to mapping", keyName)
+	return eotsPk, nil
+}
+
+func printFromKey(cmd *cobra.Command, keyName string, eotsPk *types.BIP340PubKey) error {
+	clientCtx, err := client.GetClientQueryContext(cmd)
+	if err != nil {
+		return err
 	}
 
 	k, err := clientCtx.Keyring.Key(keyName)