fix: eotsd --output flag (#166)

- Respects the output flag of text or json Closes #156
babylonlabs-io · Nov 27, 2024 · 0fca54e · 0fca54e
1 parent 86626fc
commit 0fca54e
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 103 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -37,6 +37,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## Unreleased
 
+### Bug Fixes
+
+* [#166](https://github.com/babylonlabs-io/finality-provider/pull/166) fix: `eotsd keys add` `--output` flag
+
 ### Improvements
 
 * [#149](https://github.com/babylonlabs-io/finality-provider/pull/149) Remove update of config after `fpd keys add`

diff --git a/eotsmanager/cmd/eotsd/daemon/keys.go b/eotsmanager/cmd/eotsd/daemon/keys.go
@@ -2,21 +2,26 @@ package daemon
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
+	"io"
 
+	"github.com/babylonlabs-io/babylon/types"
 	"github.com/babylonlabs-io/finality-provider/eotsmanager"
 	"github.com/babylonlabs-io/finality-provider/eotsmanager/config"
 	"github.com/babylonlabs-io/finality-provider/log"
 	"github.com/babylonlabs-io/finality-provider/util"
 	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/client/keys"
+	cryptokeyring "github.com/cosmos/cosmos-sdk/crypto/keyring"
 	"github.com/spf13/cobra"
 	"sigs.k8s.io/yaml"
 )
 
-type KeyOutput struct {
-	Name      string `json:"name" yaml:"name"`
-	PubKeyHex string `json:"pub_key_hex" yaml:"pub_key_hex"`
+type KeyOutputWithPubKeyHex struct {
+	keys.KeyOutput
+	PubKeyHex string `json:"pubkey_hex" yaml:"pubkey_hex"`
 }
 
 func NewKeysCmd() *cobra.Command {
@@ -32,6 +37,8 @@ func NewKeysCmd() *cobra.Command {
 	// the sdk, but it allows empty hd path and allow to save the key
 	// in the name mapping
 	addCmd.RunE = func(cmd *cobra.Command, args []string) error {
+		oldOut := cmd.OutOrStdout()
+
 		// Create a buffer to intercept the key items
 		var buf bytes.Buffer
 		cmd.SetOut(&buf)
@@ -42,6 +49,7 @@ func NewKeysCmd() *cobra.Command {
 			return err
 		}
 
+		cmd.SetOut(oldOut)
 		return saveKeyNameMapping(cmd, args)
 	}
 
@@ -91,21 +99,93 @@ func saveKeyNameMapping(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to save key name mapping: %w", err)
 	}
 
-	return printKey(
-		&KeyOutput{
-			Name:      keyName,
-			PubKeyHex: eotsPk.MarshalHex(),
-		},
-	)
+	k, err := clientCtx.Keyring.Key(keyName)
+	if err != nil {
+		return fmt.Errorf("failed to get public get key %s: %w", keyName, err)
+	}
+
+	ctx := cmd.Context()
+	mnemonic := ctx.Value(mnemonicCtxKey).(string) //nolint: forcetypeassert
+	showMnemonic := ctx.Value(mnemonicShowCtxKey).(bool)
+	return printCreatePubKeyHex(cmd, k, eotsPk, showMnemonic, mnemonic, clientCtx.OutputFormat)
 }
 
-func printKey(keyRecord *KeyOutput) error {
-	out, err := yaml.Marshal(keyRecord)
+func printCreatePubKeyHex(cmd *cobra.Command, k *cryptokeyring.Record, eotsPk *types.BIP340PubKey, showMnemonic bool, mnemonic, outputFormat string) error {
+	out, err := keys.MkAccKeyOutput(k)
 	if err != nil {
 		return err
 	}
+	keyOutput := newKeyOutputWithPubKeyHex(out, eotsPk)
+
+	switch outputFormat {
+	case flags.OutputFormatText:
+		cmd.PrintErrln()
+		if err := printKeyringRecord(cmd.OutOrStdout(), keyOutput, outputFormat); err != nil {
+			return err
+		}
+
+		// print mnemonic unless requested not to.
+		if showMnemonic {
+			if _, err := fmt.Fprintf(cmd.ErrOrStderr(), "\n**Important** write this mnemonic phrase in a safe place.\nIt is the only way to recover your account if you ever forget your password.\n\n%s\n", mnemonic); err != nil {
+				return fmt.Errorf("failed to print mnemonic: %s", err.Error())
+			}
+		}
+	case flags.OutputFormatJSON:
+		if showMnemonic {
+			keyOutput.Mnemonic = mnemonic
+		}
+
+		jsonString, err := json.MarshalIndent(keyOutput, "", "  ")
+		if err != nil {
+			return err
+		}
+
+		cmd.Println(string(jsonString))
+
+	default:
+		return fmt.Errorf("invalid output format %s", outputFormat)
+	}
 
-	fmt.Printf("\n%s\n", out)
+	return nil
+}
+
+func newKeyOutputWithPubKeyHex(k keys.KeyOutput, eotsPk *types.BIP340PubKey) KeyOutputWithPubKeyHex {
+	return KeyOutputWithPubKeyHex{
+		KeyOutput: k,
+		PubKeyHex: eotsPk.MarshalHex(),
+	}
+}
+
+func printKeyringRecord(w io.Writer, ko KeyOutputWithPubKeyHex, output string) error {
+	switch output {
+	case flags.OutputFormatText:
+		if err := printTextRecords(w, []KeyOutputWithPubKeyHex{ko}); err != nil {
+			return err
+		}
+
+	case flags.OutputFormatJSON:
+		out, err := json.Marshal(ko)
+		if err != nil {
+			return err
+		}
+
+		if _, err := fmt.Fprintln(w, string(out)); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func printTextRecords(w io.Writer, kos []KeyOutputWithPubKeyHex) error {
+	out, err := yaml.Marshal(&kos)
+	if err != nil {
+		return err
+	}
+
+	if _, err := fmt.Fprintln(w, string(out)); err != nil {
+		return err
+	}
 
 	return nil
 }
diff --git a/eotsmanager/cmd/eotsd/daemon/keysadd.go b/eotsmanager/cmd/eotsd/daemon/keysadd.go
@@ -3,6 +3,7 @@ package daemon
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -24,11 +25,12 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/go-bip39"
 	"github.com/spf13/cobra"
-	"sigs.k8s.io/yaml"
 )
 
 const (
 	mnemonicEntropySize = 256
+	mnemonicCtxKey      = "mnemonic_ctx"
+	mnemonicShowCtxKey  = "mnemonic_show_ctx"
 )
 
 func runAddCmdPrepare(cmd *cobra.Command, args []string) error {
@@ -64,7 +66,6 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 	noBackup, _ := cmd.Flags().GetBool(flagNoBackup)
 	showMnemonic := !noBackup
 	kb := ctx.Keyring
-	outputFormat := ctx.OutputFormat
 
 	keyringAlgos, _ := kb.SupportedAlgorithms()
 	algoStr, _ := cmd.Flags().GetString(flags.FlagKeyType)
@@ -123,12 +124,12 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 			}
 
 			pk := multisig.NewLegacyAminoPubKey(multisigThreshold, pks)
-			k, err := kb.SaveMultisig(name, pk)
+			_, err := kb.SaveMultisig(name, pk)
 			if err != nil {
 				return err
 			}
 
-			return printCreate(cmd, k, false, "", outputFormat)
+			return nil
 		}
 	}
 
@@ -143,12 +144,12 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 			return err
 		}
 
-		k, err := kb.SaveOfflineKey(name, pk)
+		_, err := kb.SaveOfflineKey(name, pk)
 		if err != nil {
 			return err
 		}
 
-		return printCreate(cmd, k, false, "", outputFormat)
+		return nil
 	}
 	if pubKeyBase64 != "" {
 		b64, err := base64.StdEncoding.DecodeString(pubKeyBase64)
@@ -175,12 +176,12 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 			return err
 		}
 
-		k, err := kb.SaveOfflineKey(name, pk)
+		_, err = kb.SaveOfflineKey(name, pk)
 		if err != nil {
 			return fmt.Errorf("failed to save offline key: %w", err)
 		}
 
-		return printCreate(cmd, k, false, "", outputFormat)
+		return nil
 	}
 
 	coinType, _ := cmd.Flags().GetUint32(flagCoinType)
@@ -199,12 +200,12 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 	// If we're using ledger, only thing we need is the path and the bech32 prefix.
 	if useLedger {
 		bech32PrefixAccAddr := sdk.GetConfig().GetBech32AccountAddrPrefix()
-		k, err := kb.SaveLedgerKey(name, hd.Secp256k1, bech32PrefixAccAddr, coinType, account, index)
+		_, err := kb.SaveLedgerKey(name, hd.Secp256k1, bech32PrefixAccAddr, coinType, account, index)
 		if err != nil {
 			return err
 		}
 
-		return printCreate(cmd, k, false, "", outputFormat)
+		return nil
 	}
 
 	// Get bip39 mnemonic
@@ -281,7 +282,7 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 		}
 	}
 
-	k, err := kb.NewAccount(name, mnemonic, bip39Passphrase, hdPath, algo)
+	_, err = kb.NewAccount(name, mnemonic, bip39Passphrase, hdPath, algo)
 	if err != nil {
 		return err
 	}
@@ -293,7 +294,10 @@ func runAddCmd(ctx client.Context, cmd *cobra.Command, args []string, inBuf *buf
 		mnemonic = ""
 	}
 
-	return printCreate(cmd, k, showMnemonic, mnemonic, outputFormat)
+	// used later for printing the values if needed
+	ctxWithValues := context.WithValue(cmd.Context(), mnemonicCtxKey, mnemonic)        //nolint: revive,staticcheck
+	cmd.SetContext(context.WithValue(ctxWithValues, mnemonicShowCtxKey, showMnemonic)) //nolint: revive,staticcheck
+	return nil
 }
 
 func validateMultisigThreshold(k, nKeys int) error {
@@ -307,72 +311,6 @@ func validateMultisigThreshold(k, nKeys int) error {
 	return nil
 }
 
-func printCreate(cmd *cobra.Command, k *cryptokeyring.Record, showMnemonic bool, mnemonic, outputFormat string) error {
-	switch outputFormat {
-	case flags.OutputFormatText:
-		cmd.PrintErrln()
-		if err := printKeyringRecord(cmd.OutOrStdout(), k, keys.MkAccKeyOutput, outputFormat); err != nil {
-			return err
-		}
-
-		// print mnemonic unless requested not to.
-		if showMnemonic {
-			if _, err := fmt.Fprintf(cmd.ErrOrStderr(), "\n**Important** write this mnemonic phrase in a safe place.\nIt is the only way to recover your account if you ever forget your password.\n\n%s\n", mnemonic); err != nil {
-				return fmt.Errorf("failed to print mnemonic: %s", err.Error())
-			}
-		}
-	case flags.OutputFormatJSON:
-		out, err := keys.MkAccKeyOutput(k)
-		if err != nil {
-			return err
-		}
-
-		if showMnemonic {
-			out.Mnemonic = mnemonic
-		}
-
-		jsonString, err := json.Marshal(out)
-		if err != nil {
-			return err
-		}
-
-		cmd.Println(string(jsonString))
-
-	default:
-		return fmt.Errorf("invalid output format %s", outputFormat)
-	}
-
-	return nil
-}
-
-type bechKeyOutFn func(k *cryptokeyring.Record) (keys.KeyOutput, error)
-
-func printKeyringRecord(w io.Writer, k *cryptokeyring.Record, bechKeyOut bechKeyOutFn, output string) error {
-	ko, err := bechKeyOut(k)
-	if err != nil {
-		return err
-	}
-
-	switch output {
-	case flags.OutputFormatText:
-		if err := printTextRecords(w, []keys.KeyOutput{ko}); err != nil {
-			return err
-		}
-
-	case flags.OutputFormatJSON:
-		out, err := json.Marshal(ko)
-		if err != nil {
-			return err
-		}
-
-		if _, err := fmt.Fprintln(w, string(out)); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
 func readMnemonicFromFile(filePath string) (string, error) {
 	file, err := os.Open(filepath.Clean(filePath))
 	if err != nil {
@@ -386,16 +324,3 @@ func readMnemonicFromFile(filePath string) (string, error) {
 	}
 	return string(bz), nil
 }
-
-func printTextRecords(w io.Writer, kos []keys.KeyOutput) error {
-	out, err := yaml.Marshal(&kos)
-	if err != nil {
-		return err
-	}
-
-	if _, err := fmt.Fprintln(w, string(out)); err != nil {
-		return err
-	}
-
-	return nil
-}