comments

sei-protocol · Nov 15, 2024 · a306f8c · a306f8c
1 parent e889bae
commit a306f8c
Showing 1 changed file with 122 additions and 101 deletions.
diff --git a/x/confidentialtransfers/client/cli/query.go b/x/confidentialtransfers/client/cli/query.go
@@ -12,6 +12,8 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const decryptAvailableBalanceFlag = "decrypt-available-balance"
+
 // GetQueryCmd returns the cli query commands for the minting module.
 func GetQueryCmd() *cobra.Command {
 	confidentialTransfersQueryCmd := &cobra.Command{
@@ -36,124 +38,143 @@ func GetCmdQueryAccount() *cobra.Command {
 		Use:   "account [denom] [address] [flags]",
 		Short: "Query the account state",
 		Long: "Queries the account state associated with the address and denom." +
-			"Pass the --from flag to decrypt the account." +
-			"Pass the --decryptAvailableBalance flag to attempt to decrypt the available balance.",
+			"Pass the --from flag to decrypt the account" +
+			"Pass the --decrypt-available-balance flag to attempt to decrypt the available balance.",
 		Args: cobra.ExactArgs(2),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			clientCtx, err := client.GetClientQueryContext(cmd)
-			if err != nil {
-				return err
-			}
-			queryClient := types.NewQueryClient(clientCtx)
-
-			denom := args[0]
-			address := args[1]
-			// Validate address
-			_, err = sdk.AccAddressFromBech32(address)
-			if err != nil {
-				return fmt.Errorf("invalid address: %v", err)
-			}
-
-			from, err := cmd.Flags().GetString(flags.FlagFrom)
-			if err != nil {
-				return err
-			}
-			fromAddr, _, _, err := client.GetFromFields(clientCtx, clientCtx.Keyring, from)
-			if err != nil {
-				return err
-			}
-
-			res, err := queryClient.GetCtAccount(cmd.Context(), &types.GetCtAccountRequest{
-				Address: address,
-				Denom:   denom,
-			})
-
-			if fromAddr.String() == args[0] {
-				account, err := res.Account.FromProto()
-				if err != nil {
-					return err
-				}
-				privateKey, err := getPrivateKey(cmd)
-				if err != nil {
-					return err
-				}
-
-				aesKey, err := encryption.GetAESKey(*privateKey, denom)
-				if err != nil {
-					return err
-				}
-
-				decryptor := elgamal.NewTwistedElgamal()
-				keyPair, err := decryptor.KeyGen(*privateKey, denom)
-				if err != nil {
-					return err
-				}
-
-				decryptAvailableBalance, err := cmd.Flags().GetBool("decryptAvailableBalance")
-				if err != nil {
-					return err
-				}
-
-				decryptedAccount, err := account.Decrypt(decryptor, keyPair, aesKey, decryptAvailableBalance)
-				if err != nil {
-					return err
-				}
-
-				return clientCtx.PrintProto(decryptedAccount)
-
-			} else {
-				return clientCtx.PrintProto(res.Account)
-			}
-		},
+		RunE: queryAccount,
 	}
 
 	flags.AddQueryFlagsToCmd(cmd)
 	cmd.Flags().String(flags.FlagFrom, "", "Name or address of private key to decrypt the account")
-	cmd.Flags().Bool("decryptAvailableBalance", false, "Set this to attempt to decrypt the available balance")
+	cmd.Flags().Bool(decryptAvailableBalanceFlag, false, "Set this to attempt to decrypt the available balance")
 	return cmd
 }
 
+func queryAccount(cmd *cobra.Command, args []string) error {
+	clientCtx, err := client.GetClientQueryContext(cmd)
+	if err != nil {
+		return err
+	}
+	queryClient := types.NewQueryClient(clientCtx)
+
+	denom := args[0]
+
+	// Validate denom
+	err = sdk.ValidateDenom(denom)
+	if err != nil {
+		return fmt.Errorf("invalid denom: %v", err)
+	}
+
+	address := args[1]
+	// Validate address
+	_, err = sdk.AccAddressFromBech32(address)
+	if err != nil {
+		return fmt.Errorf("invalid address: %v", err)
+	}
+
+	from, err := cmd.Flags().GetString(flags.FlagFrom)
+	if err != nil {
+		return err
+	}
+	fromAddr, _, _, err := client.GetFromFields(clientCtx, clientCtx.Keyring, from)
+	if err != nil {
+		return err
+	}
+
+	res, err := queryClient.GetCtAccount(cmd.Context(), &types.GetCtAccountRequest{
+		Address: address,
+		Denom:   denom,
+	})
+	if err != nil {
+		return err
+	}
+
+	// If the --from flag passed matches the queried address, attempt to decrypt the contents
+	if fromAddr.String() == address {
+		account, err := res.Account.FromProto()
+		if err != nil {
+			return err
+		}
+		privateKey, err := getPrivateKey(cmd)
+		if err != nil {
+			return err
+		}
+
+		aesKey, err := encryption.GetAESKey(*privateKey, denom)
+		if err != nil {
+			return err
+		}
+
+		decryptor := elgamal.NewTwistedElgamal()
+		keyPair, err := decryptor.KeyGen(*privateKey, denom)
+		if err != nil {
+			return err
+		}
+
+		decryptAvailableBalance, err := cmd.Flags().GetBool(decryptAvailableBalanceFlag)
+		if err != nil {
+			return err
+		}
+
+		decryptedAccount, err := account.Decrypt(decryptor, keyPair, aesKey, decryptAvailableBalance)
+		if err != nil {
+			return err
+		}
+
+		return clientCtx.PrintProto(decryptedAccount)
+
+	} else {
+		return clientCtx.PrintProto(res.Account)
+	}
+}
+
 // GetCmdQueryAccount implements a command to return an account asssociated with the address and denom
 func GetCmdQueryAllAccount() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "accounts [address]",
 		Short: "Query all the confidential token accounts associated with the address",
 		Args:  cobra.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			clientCtx, err := client.GetClientQueryContext(cmd)
-			if err != nil {
-				return err
-			}
-			queryClient := types.NewQueryClient(clientCtx)
-
-			address := args[0]
-			// Validate address
-			_, err = sdk.AccAddressFromBech32(address)
-			if err != nil {
-				return fmt.Errorf("invalid address: %v", err)
-			}
-
-			res, err := queryClient.GetAllCtAccounts(cmd.Context(), &types.GetAllCtAccountsRequest{
-				Address: args[0],
-			})
-
-			if err != nil {
-				return err
-			}
-
-			for i := range res.Accounts {
-				clientCtx.PrintString("\n")
-				err = clientCtx.PrintProto(&res.Accounts[i])
-				if err != nil {
-					return err
-				}
-			}
-
-			return nil
-		},
+		RunE:  queryAllAccounts,
 	}
 
 	flags.AddQueryFlagsToCmd(cmd)
 
 	return cmd
 }
+
+func queryAllAccounts(cmd *cobra.Command, args []string) error {
+	clientCtx, err := client.GetClientQueryContext(cmd)
+	if err != nil {
+		return err
+	}
+	queryClient := types.NewQueryClient(clientCtx)
+
+	address := args[0]
+	// Validate address
+	_, err = sdk.AccAddressFromBech32(address)
+	if err != nil {
+		return fmt.Errorf("invalid address: %v", err)
+	}
+
+	res, err := queryClient.GetAllCtAccounts(cmd.Context(), &types.GetAllCtAccountsRequest{
+		Address: args[0],
+	})
+
+	if err != nil {
+		return err
+	}
+
+	for i := range res.Accounts {
+		err = clientCtx.PrintString("\n")
+		if err != nil {
+			return err
+		}
+
+		err = clientCtx.PrintProto(&res.Accounts[i])
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}