Merge branch 'master' into george-elo-412-cli-gate-rewards-v2-command…

…s-to-holesky
Layr-Labs · Dec 13, 2024 · 7084ab7 · 7084ab7
2 parents 752049f + d41fc29
commit 7084ab7
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 16 deletions.
diff --git a/pkg/keys/import.go b/pkg/keys/import.go
@@ -44,18 +44,35 @@ This command will import keys in $HOME/.eigenlayer/operator_keys/ location
 		After: telemetry.AfterRunAction(),
 		Action: func(ctx *cli.Context) error {
 			args := ctx.Args()
-			if args.Len() != 2 {
-				return fmt.Errorf("%w: accepts 2 arg, received %d", ErrInvalidNumberOfArgs, args.Len())
+			if args.Len() != 1 && args.Len() != 2 {
+				return fmt.Errorf("%w: accepts 1 or 2 arg, received %d", ErrInvalidNumberOfArgs, args.Len())
 			}
 
+			var err error
 			keyName := args.Get(0)
-			if err := validateKeyName(keyName); err != nil {
+			if err = validateKeyName(keyName); err != nil {
 				return err
 			}
 
 			privateKey := args.Get(1)
+
+			// In case user doesn't provide private key, prompt for it
+			// This is to ensure backward compatibility
 			if privateKey == "" {
-				return ErrEmptyPrivateKey
+				privateKey, err = p.InputString(
+					"Enter private key or Mnemonic (mnemonic only work for ecdsa right now): ",
+					"",
+					"Please provide private key or Mnemonic to import key",
+					func(s string) error {
+						if s == "" {
+							return ErrEmptyPrivateKey
+						}
+						return nil
+					},
+				)
+				if err != nil {
+					return err
+				}
 			}
 
 			pkSlice := strings.Split(privateKey, " ")

diff --git a/pkg/keys/import_test.go b/pkg/keys/import_test.go
@@ -37,16 +37,10 @@ func TestImportCmd(t *testing.T) {
 			args: []string{},
 			err:  errors.New("Required flag \"key-type\" not set"),
 		},
-		{
-			name: "one argument",
-			args: []string{"--key-type", "ecdsa", "arg1"},
-			err:  fmt.Errorf("%w: accepts 2 arg, received 1", ErrInvalidNumberOfArgs),
-		},
-
 		{
 			name: "more than two argument",
 			args: []string{"--key-type", "ecdsa", "arg1", "arg2", "arg3"},
-			err:  fmt.Errorf("%w: accepts 2 arg, received 3", ErrInvalidNumberOfArgs),
+			err:  fmt.Errorf("%w: accepts 1 or 2 arg, received 3", ErrInvalidNumberOfArgs),
 		},
 		{
 			name: "empty key name argument",
@@ -58,11 +52,6 @@ func TestImportCmd(t *testing.T) {
 			args: []string{"--key-type", "ecdsa", "hello world", ""},
 			err:  ErrKeyContainsWhitespaces,
 		},
-		{
-			name: "empty private key argument",
-			args: []string{"--key-type", "ecdsa", "hello", ""},
-			err:  ErrEmptyPrivateKey,
-		},
 		{
 			name: "keyname with whitespaces",
 			args: []string{"--key-type", "ecdsa", "hello", "hello world"},
@@ -110,6 +99,24 @@ func TestImportCmd(t *testing.T) {
 			expectedPrivKey: "6842fb8f5fa574d0482818b8a825a15c4d68f542693197f2c2497e3562f335f6",
 			keyPath:         filepath.Join(homePath, OperatorKeystoreSubFolder, "/test.ecdsa.key.json"),
 		},
+		{
+			name: "valid ecdsa key import with prompt",
+			args: []string{
+				"--key-type",
+				"ecdsa",
+				"test",
+			},
+			err: nil,
+			promptMock: func(p *prompterMock.MockPrompter) {
+				p.EXPECT().
+					InputString(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+					Return("6842fb8f5fa574d0482818b8a825a15c4d68f542693197f2c2497e3562f335f6", nil)
+				p.EXPECT().InputHiddenString(gomock.Any(), gomock.Any(), gomock.Any()).Return("", nil)
+				p.EXPECT().InputHiddenString(gomock.Any(), gomock.Any(), gomock.Any()).Return("", nil)
+			},
+			expectedPrivKey: "6842fb8f5fa574d0482818b8a825a15c4d68f542693197f2c2497e3562f335f6",
+			keyPath:         filepath.Join(homePath, OperatorKeystoreSubFolder, "/test.ecdsa.key.json"),
+		},
 		{
 			name: "valid ecdsa key import with 0x prefix",
 			args: []string{
@@ -158,6 +165,24 @@ func TestImportCmd(t *testing.T) {
 			expectedPrivKey: "20030410000080487431431153104351076122223465926814327806350179952713280726583",
 			keyPath:         filepath.Join(homePath, OperatorKeystoreSubFolder, "/test.bls.key.json"),
 		},
+		{
+			name: "valid bls key import with prompt",
+			args: []string{
+				"--key-type",
+				"bls",
+				"test",
+			},
+			err: nil,
+			promptMock: func(p *prompterMock.MockPrompter) {
+				p.EXPECT().
+					InputString(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+					Return("20030410000080487431431153104351076122223465926814327806350179952713280726583", nil)
+				p.EXPECT().InputHiddenString(gomock.Any(), gomock.Any(), gomock.Any()).Return("", nil)
+				p.EXPECT().InputHiddenString(gomock.Any(), gomock.Any(), gomock.Any()).Return("", nil)
+			},
+			expectedPrivKey: "20030410000080487431431153104351076122223465926814327806350179952713280726583",
+			keyPath:         filepath.Join(homePath, OperatorKeystoreSubFolder, "/test.bls.key.json"),
+		},
 		{
 			name: "valid bls key import for hex key",
 			args: []string{