VAULT-27953: Bootstrap cmd for AWS Key Pair Setup

go.mod

@@ -37,17 +37,19 @@ require (
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect
 	github.com/armon/go-radix v1.0.0 // indirect
-	github.com/aws/aws-sdk-go-v2 v1.21.0 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.36.3 // indirect
 	github.com/aws/aws-sdk-go-v2/credentials v1.13.37 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.13.11 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.41 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.35 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.3.42 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.35 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ec2 v1.224.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.13.6 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.15.6 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.21.5 // indirect
-	github.com/aws/smithy-go v1.14.2 // indirect
+	github.com/aws/smithy-go v1.22.2 // indirect
 	github.com/bgentry/speakeasy v0.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fatih/color v1.17.0 // indirect

go.sum

@@ -22,6 +22,8 @@ github.com/armon/go-radix v1.0.0 h1:F4z6KzEeeQIMeLFa97iZU6vupzoecKdU5TX24SNppXI=
 github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
 github.com/aws/aws-sdk-go-v2 v1.21.0 h1:gMT0IW+03wtYJhRqTVYn0wLzwdnK9sRMcxmtfGzRdJc=
 github.com/aws/aws-sdk-go-v2 v1.21.0/go.mod h1:/RfNgGmRxI+iFOB1OeJUyxiU+9s88k3pfHvDagGEp0M=
+github.com/aws/aws-sdk-go-v2 v1.36.3 h1:mJoei2CxPutQVxaATCzDUjcZEjVRdpsiiXi2o38yqWM=
+github.com/aws/aws-sdk-go-v2 v1.36.3/go.mod h1:LLXuLpgzEbD766Z5ECcRmi8AzSwfZItDtmABVkRLGzg=
 github.com/aws/aws-sdk-go-v2/config v1.18.39 h1:oPVyh6fuu/u4OiW4qcuQyEtk7U7uuNBmHmJSLg1AJsQ=
 github.com/aws/aws-sdk-go-v2/config v1.18.39/go.mod h1:+NH/ZigdPckFpgB1TRcRuWCB/Kbbvkxc/iNAKTq5RhE=
 github.com/aws/aws-sdk-go-v2/credentials v1.13.37 h1:BvEdm09+ZEh2XtN+PVHPcYwKY3wIeB6pw7vPRM4M9/U=
@@ -30,12 +32,22 @@ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.13.11 h1:uDZJF1hu0EVT/4bogChk8D
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.13.11/go.mod h1:TEPP4tENqBGO99KwVpV9MlOX4NSrSLP8u3KRy2CDwA8=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.41 h1:22dGT7PneFMx4+b3pz7lMTRyN8ZKH7M2cW4GP9yUS2g=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.41/go.mod h1:CrObHAuPneJBlfEJ5T3szXOUkLEThaGfvnhTf33buas=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34 h1:ZK5jHhnrioRkUNOc+hOgQKlUL5JeC3S6JgLxtQ+Rm0Q=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34/go.mod h1:p4VfIceZokChbA9FzMbRGz5OV+lekcVtHlPKEO0gSZY=
 github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.35 h1:SijA0mgjV8E+8G45ltVHs0fvKpTj8xmZJ3VwhGKtUSI=
 github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.35/go.mod h1:SJC1nEVVva1g3pHAIdCp7QsRIkMmLAgoDquQ9Rr8kYw=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34 h1:SZwFm17ZUNNg5Np0ioo/gq8Mn6u9w19Mri8DnJ15Jf0=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34/go.mod h1:dFZsC0BLo346mvKQLWmoJxT+Sjp+qcVR1tRVHQGOH9Q=
 github.com/aws/aws-sdk-go-v2/internal/ini v1.3.42 h1:GPUcE/Yq7Ur8YSUk6lVkoIMWnJNO0HT18GUzCWCgCI0=
 github.com/aws/aws-sdk-go-v2/internal/ini v1.3.42/go.mod h1:rzfdUlfA+jdgLDmPKjd3Chq9V7LVLYo1Nz++Wb91aRo=
+github.com/aws/aws-sdk-go-v2/service/ec2 v1.224.1 h1:O5i9p1FonEMcU2D5jNwv6W1IqHtZN3XiVzejefobTBI=
+github.com/aws/aws-sdk-go-v2/service/ec2 v1.224.1/go.mod h1:ouvGEfHbLaIlWwpDpOVWPWR+YwO0HDv3vm5tYLq8ImY=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3 h1:eAh2A4b5IzM/lum78bZ590jy36+d/aFLgKF/4Vd1xPE=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3/go.mod h1:0yKJC/kb8sAnmlYa6Zs3QVYqaC8ug2AbnNChv5Ox3uA=
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.35 h1:CdzPW9kKitgIiLV1+MHobfR5Xg25iYnyzWZhyQuSlDI=
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.35/go.mod h1:QGF2Rs33W5MaN9gYdEQOBBFPLwTZkEhRwI33f7KIG0o=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 h1:dM9/92u2F1JbDaGooxTq18wmmFzbJRfXfVfy96/1CXM=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15/go.mod h1:SwFBy2vjtA0vZbjjaFtfN045boopadnoVPhu4Fv66vY=
 github.com/aws/aws-sdk-go-v2/service/sso v1.13.6 h1:2PylFCfKCEDv6PeSN09pC/VUiRd10wi1VfHG5FrW0/g=
 github.com/aws/aws-sdk-go-v2/service/sso v1.13.6/go.mod h1:fIAwKQKBFu90pBxx07BFOMJLpRUGu8VOzLJakeY+0K4=
 github.com/aws/aws-sdk-go-v2/service/ssooidc v1.15.6 h1:pSB560BbVj9ZlJZF4WYj5zsytWHWKxg+NgyGV4B2L58=
@@ -44,6 +56,8 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.21.5 h1:CQBFElb0LS8RojMJlxRSo/HXipvT
 github.com/aws/aws-sdk-go-v2/service/sts v1.21.5/go.mod h1:VC7JDqsqiwXukYEDjoHh9U0fOJtNWh04FPQz4ct4GGU=
 github.com/aws/smithy-go v1.14.2 h1:MJU9hqBGbvWZdApzpvoF2WAIJDbtjK2NDJSiJP7HblQ=
 github.com/aws/smithy-go v1.14.2/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA=
+github.com/aws/smithy-go v1.22.2 h1:6D9hW43xKFrRx/tXXfAlIZc4JI+yQe6snnWcQyxSyLQ=
+github.com/aws/smithy-go v1.22.2/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
 github.com/bgentry/speakeasy v0.1.0 h1:ByYyxL9InA1OWqxJqqp2A5pYHUrCiAL6K3J+LKSsQkY=
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
 github.com/bgentry/speakeasy v0.2.0 h1:tgObeVOf8WAvtuAX6DhJ4xks4CFNwPDZiqzGqIHE51E=

internal/bootstrap/bootstrap.go

@@ -0,0 +1,90 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package bootstrap
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/smithy-go"
+)
+
+type describeKeyPairsAPI interface {
+	DescribeKeyPairs(ctx context.Context, params *ec2.DescribeKeyPairsInput, optFns ...func(*ec2.Options)) (*ec2.DescribeKeyPairsOutput, error)
+	DeleteKeyPair(ctx context.Context, params *ec2.DeleteKeyPairInput, optFns ...func(*ec2.Options)) (*ec2.DeleteKeyPairOutput, error)
+	CreateKeyPair(ctx context.Context, params *ec2.CreateKeyPairInput, optFns ...func(*ec2.Options)) (*ec2.CreateKeyPairOutput, error)
+}
+
+func Run(client describeKeyPairsAPI, keypairName, sshDir string, force bool) error {
+	exists, err := keyPairExists(client, keypairName)
+	if err != nil {
+		return err
+	}
+
+	if exists && !force {
+		fmt.Printf("Key pair %q already exists. Use --force to recreate.\n", keypairName)
+		return nil
+	}
+
+	if exists && force {
+		fmt.Printf("Deleting existing key pair...\n")
+		_, _ = client.DeleteKeyPair(context.TODO(), &ec2.DeleteKeyPairInput{
+			KeyName: aws.String(keypairName),
+		})
+	}
+
+	fmt.Printf("Creating new key pair %q...\n", keypairName)
+	out, err := client.CreateKeyPair(context.TODO(), &ec2.CreateKeyPairInput{
+		KeyName: aws.String(keypairName),
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create key pair: %w", err)
+	}
+
+	err = os.MkdirAll(sshDir, 0o700)
+	if err != nil {
+		return fmt.Errorf("failed to create ssh directory: %w", err)
+	}
+
+	expandedPath := filepath.Join(sshDir, keypairName+".pem")
+	err = os.WriteFile(expandedPath, []byte(*out.KeyMaterial), 0o600)
+	if err != nil {
+		return fmt.Errorf("failed to write private key: %w", err)
+	}
+
+	fmt.Printf(`
+Key pair %q created and saved to:
+
+    %s
+
+Please update your enos-local.vars.hcl with the following:
+
+    aws_ssh_keypair_name       = "%s"
+    aws_ssh_private_key_path   = "%s"
+
+`, keypairName, expandedPath, keypairName, expandedPath)
+
+	return nil
+}
+
+func keyPairExists(client describeKeyPairsAPI, name string) (bool, error) {
+	_, err := client.DescribeKeyPairs(context.TODO(), &ec2.DescribeKeyPairsInput{
+		KeyNames: []string{name},
+	})
+	if err != nil {
+		var apiErr smithy.APIError
+		if errors.As(err, &apiErr) && apiErr.ErrorCode() == "InvalidKeyPair.NotFound" {
+			return false, nil
+		}
+
+		return false, err
+	}
+
+	return true, nil
+}

internal/bootstrap/bootstrap_test.go

@@ -0,0 +1,133 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package bootstrap
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
+	"github.com/aws/smithy-go"
+)
+
+type mockAPIError struct {
+	code string
+}
+
+func (e *mockAPIError) Error() string                 { return e.code }
+func (e *mockAPIError) ErrorCode() string             { return e.code }
+func (e *mockAPIError) ErrorMessage() string          { return e.code }
+func (e *mockAPIError) ErrorFault() smithy.ErrorFault { return smithy.FaultClient }
+
+type mockKeyManager struct {
+	keyExists    bool
+	deleteCalled bool
+	createCalled bool
+}
+
+func (m *mockKeyManager) DescribeKeyPairs(ctx context.Context, input *ec2.DescribeKeyPairsInput, optFns ...func(*ec2.Options)) (*ec2.DescribeKeyPairsOutput, error) {
+	if m.keyExists {
+		return &ec2.DescribeKeyPairsOutput{
+			KeyPairs: []types.KeyPairInfo{
+				{KeyName: &input.KeyNames[0]},
+			},
+		}, nil
+	}
+
+	return nil, &mockAPIError{code: "InvalidKeyPair.NotFound"}
+}
+
+func (m *mockKeyManager) DeleteKeyPair(ctx context.Context, input *ec2.DeleteKeyPairInput, optFns ...func(*ec2.Options)) (*ec2.DeleteKeyPairOutput, error) {
+	m.deleteCalled = true
+	m.keyExists = false
+
+	return &ec2.DeleteKeyPairOutput{}, nil
+}
+
+func (m *mockKeyManager) CreateKeyPair(ctx context.Context, input *ec2.CreateKeyPairInput, optFns ...func(*ec2.Options)) (*ec2.CreateKeyPairOutput, error) {
+	m.createCalled = true
+	key := "mock-private-key"
+
+	return &ec2.CreateKeyPairOutput{
+		KeyName:     input.KeyName,
+		KeyMaterial: &key,
+	}, nil
+}
+
+func TestKeyPairExists(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name      string
+		client    describeKeyPairsAPI
+		expect    bool
+		expectErr bool
+	}{
+		{
+			name:      "key exists",
+			client:    &mockKeyManager{keyExists: true},
+			expect:    true,
+			expectErr: false,
+		},
+		{
+			name:      "key does not exist",
+			client:    &mockKeyManager{keyExists: false},
+			expect:    false,
+			expectErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			got, err := keyPairExists(tt.client, "enos-ec2-key")
+			if (err != nil) != tt.expectErr {
+				t.Fatalf("expected error: %v, got: %v", tt.expectErr, err)
+			}
+			if got != tt.expect {
+				t.Errorf("expected: %v, got: %v", tt.expect, got)
+			}
+		})
+	}
+}
+
+func TestRunScenarios(t *testing.T) {
+	tempHome := t.TempDir()
+	sshPath := filepath.Join(tempHome, ".ssh")
+	err := os.MkdirAll(sshPath, 0o700)
+	if err != nil {
+		t.Fatalf("failed to create temp ssh dir: %v", err)
+	}
+	t.Setenv("HOME", tempHome)
+
+	tests := []struct {
+		name        string
+		force       bool
+		keyExists   bool
+		expectWrite bool
+	}{
+		{"Key exists, no force", false, true, false},
+		{"Key exists, force", true, true, true},
+		{"Key does not exist", false, false, true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			client := &mockKeyManager{keyExists: tt.keyExists}
+			err := Run(client, "enos-ec2-key", sshPath, tt.force)
+			if err != nil && tt.expectWrite {
+				t.Fatalf("expected no error, got: %v", err)
+			}
+
+			pemPath := filepath.Join(sshPath, "enos-ec2-key.pem")
+			_, err = os.Stat(pemPath)
+			if tt.expectWrite && os.IsNotExist(err) {
+				t.Errorf("expected key file to be written, but it wasn't")
+			}
+		})
+	}
+}

internal/command/enos/cmd/bootstrap.go

@@ -0,0 +1,42 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package cmd
+
+import (
+	"os"
+	"path/filepath"
+
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/spf13/cobra"
+
+	"github.com/hashicorp/enos/internal/bootstrap"
+)
+
+var (
+	force       bool
+	keypairName string
+	sshDir      string
+)
+
+var bootstrapCmd = &cobra.Command{
+	Use:   "bootstrap",
+	Short: "Initialize and configure required dependencies (e.g., AWS SSH key)",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		cfg, err := config.LoadDefaultConfig(cmd.Context())
+		if err != nil {
+			return err
+		}
+		client := ec2.NewFromConfig(cfg)
+
+		return bootstrap.Run(client, keypairName, sshDir, force)
+	},
+}
+
+func init() {
+	bootstrapCmd.Flags().BoolVar(&force, "force", false, "Force creation even if key already exists")
+	bootstrapCmd.Flags().StringVar(&keypairName, "keypair-name", "enos-ec2-key", "Name of the SSH key pair")
+	bootstrapCmd.Flags().StringVar(&sshDir, "ssh-dir", filepath.Join(os.Getenv("HOME"), ".ssh"), "Directory to save the private key")
+	rootCmd.AddCommand(bootstrapCmd)
+}