diff --git a/pkg/ec2helper/ec2helper.go b/pkg/ec2helper/ec2helper.go index 46dfb47..1c00582 100644 --- a/pkg/ec2helper/ec2helper.go +++ b/pkg/ec2helper/ec2helper.go @@ -35,6 +35,7 @@ import ( "github.com/aws/aws-sdk-go/aws/awserr" "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/ec2" + "github.com/aws/aws-sdk-go/service/ssm" "github.com/google/uuid" ) @@ -352,76 +353,137 @@ func (h *EC2Helper) getInstanceTypes(input *ec2.DescribeInstanceTypesInput) ([]* // Define all OS and corresponding AMI name formats var osDescs = map[string]map[string]string{ "Amazon Linux": { - "ebs": "amzn-ami-hvm-????.??.?.????????.?-*-gp2", - "instance-store": "amzn-ami-hvm-????.??.?.????????.?-*-s3", + "ebs": "amzn-ami-hvm-*", + "instance-store": "amzn-ami-hvm-*", }, "Amazon Linux 2": { - "ebs": "amzn2-ami-hvm-2.?.????????.?-*-gp2", + "ebs": "amzn2-ami-hvm-*", }, "Red Hat": { - "ebs": "RHEL-?.?.?_HVM-????????-*-?-Hourly2-GP2", + "ebs": "RHEL-9*", }, "SUSE Linux": { - "ebs": "suse-sles-??-sp?-v????????-hvm-ssd-*", + "ebs": "suse-sles-*", }, - // Ubuntu 18.04 LTS "Ubuntu": { - "ebs": "ubuntu/images/hvm-ssd/ubuntu-bionic-18.04-*-server-????????", - "instance-store": "ubuntu/images/hvm-instance/ubuntu-bionic-18.04-*-server-????????", + "ebs": "ubuntu/images/*", + "instance-store": "ubuntu/images/*", }, - // 64 bit Microsoft Windows Server with Desktop Experience Locale English AMI "Windows": { - "ebs": "Windows_Server-????-English-Full-Base-????.??.??", + "ebs": "Windows_Server-*-English-Full-Base*", + }, +} + +// Define all OS and corresponding AMI public parameters path in Parameter Store +var osSsmPath = map[string]map[string]string{ + "Amazon Linux": { + "ebs": "/aws/service/ami-amazon-linux-latest", + "instance-store": "/aws/service/ami-amazon-linux-latest", + }, + "Amazon Linux 2": { + "ebs": "/aws/service/ami-amazon-linux-latest", + }, + "Red Hat": { + "ebs": "", + }, + "SUSE Linux": { + "ebs": "/aws/service/suse/sles", + }, + "Ubuntu": { + "ebs": "/aws/service/canonical/ubuntu/server/24.04/stable/current", + "instance-store": "/aws/service/canonical/ubuntu/server/24.04/stable/current", + }, + "Windows": { + "ebs": "/aws/service/ami-windows-latest", }, } // Get the appropriate input for describing images -func getDescribeImagesInputs(rootDeviceType string, architectures []*string) *map[string]ec2.DescribeImagesInput { +func (h *EC2Helper) GetDescribeImagesInputs(rootDeviceType string, architectures []*string) *map[string]ec2.DescribeImagesInput { + ssmClient := ssm.New(h.Sess) + // Construct all the inputs imageInputs := map[string]ec2.DescribeImagesInput{} for osName, rootDeviceTypes := range osDescs { // Only add inputs if the corresponding root device type is applicable for the specified os desc, found := rootDeviceTypes[rootDeviceType] - if found { - imageInputs[osName] = ec2.DescribeImagesInput{ - Filters: []*ec2.Filter{ - { - Name: aws.String("name"), - Values: []*string{ - aws.String(desc), - }, - }, - { - Name: aws.String("state"), - Values: []*string{ - aws.String("available"), - }, + if !found { + continue + } + imageInputs[osName] = ec2.DescribeImagesInput{ + Filters: []*ec2.Filter{ + { + Name: aws.String("name"), + Values: []*string{ + aws.String(desc), }, - { - Name: aws.String("root-device-type"), - Values: []*string{ - aws.String(rootDeviceType), - }, + }, + { + Name: aws.String("state"), + Values: []*string{ + aws.String("available"), }, - { - Name: aws.String("architecture"), - Values: architectures, + }, + { + Name: aws.String("root-device-type"), + Values: []*string{ + aws.String(rootDeviceType), }, - { - Name: aws.String("owner-alias"), - Values: []*string{ - aws.String("amazon"), - }, + }, + { + Name: aws.String("architecture"), + Values: architectures, + }, + { + Name: aws.String("owner-alias"), + Values: []*string{ + aws.String("amazon"), }, }, - } + }, + } + ssmPath, found := osSsmPath[osName][rootDeviceType] + if !found || ssmPath == "" { + continue + } + if imageIds, err := h.GetImageIdsFromSSM(ssmClient, ssmPath); err == nil { + input := imageInputs[osName] + input.ImageIds = imageIds + imageInputs[osName] = input } } return &imageInputs } +func (h *EC2Helper) GetImageIdsFromSSM(ssmClient *ssm.SSM, ssmPath string) ([]*string, error) { + var imageIds []*string + + input := &ssm.GetParametersByPathInput{ + Path: aws.String(ssmPath), + Recursive: aws.Bool(true), + WithDecryption: aws.Bool(false), + } + + err := ssmClient.GetParametersByPathPages(input, + func(page *ssm.GetParametersByPathOutput, lastPage bool) bool { + for _, parameter := range page.Parameters { + if !strings.HasPrefix(*parameter.Value, "ami") { + continue + } + imageIds = append(imageIds, parameter.Value) + } + return !lastPage + }) + + if err != nil { + return nil, fmt.Errorf("error getting parameters from SSM: %v", err) + } + + return imageIds, nil +} + // Sort interface for images type byCreationDate []*ec2.Image @@ -436,9 +498,9 @@ Empty result is allowed. func (h *EC2Helper) GetLatestImages(rootDeviceType *string, architectures []*string) (*map[string]*ec2.Image, error) { var inputs *map[string]ec2.DescribeImagesInput if rootDeviceType == nil { - inputs = getDescribeImagesInputs("ebs", architectures) + inputs = h.GetDescribeImagesInputs("ebs", architectures) } else { - inputs = getDescribeImagesInputs(*rootDeviceType, architectures) + inputs = h.GetDescribeImagesInputs(*rootDeviceType, architectures) } images := map[string]*ec2.Image{} diff --git a/pkg/question/question_test.go b/pkg/question/question_test.go index 6958457..e5ba17e 100644 --- a/pkg/question/question_test.go +++ b/pkg/question/question_test.go @@ -343,6 +343,7 @@ func TestAskImage_Success(t *testing.T) { const expectedImage = "ami-12345" const testInstanceType = ec2.InstanceTypeT2Micro + testEC2 = ec2helper.New(session.Must(session.NewSession())) testEC2.Svc = &th.MockedEC2Svc{ InstanceTypes: []*ec2.InstanceTypeInfo{ { diff --git a/test/e2e/e2e-ec2helper-test/e2e_ec2helper_test.go b/test/e2e/e2e-ec2helper-test/e2e_ec2helper_test.go index fbe0efd..8e7d753 100644 --- a/test/e2e/e2e-ec2helper-test/e2e_ec2helper_test.go +++ b/test/e2e/e2e-ec2helper-test/e2e_ec2helper_test.go @@ -29,6 +29,7 @@ import ( "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/cloudformation" "github.com/aws/aws-sdk-go/service/ec2" + "github.com/aws/aws-sdk-go/service/ssm" ) const testStackName = "simple-ec2-e2e-ec2helper-test" @@ -155,8 +156,21 @@ func TestGetInstanceTypesFromInstanceSelector(t *testing.T) { func TestGetLatestImages(t *testing.T) { th.Assert(t, h != nil, "EC2Helper was not initialized successfully") - _, err := h.GetLatestImages(nil, aws.StringSlice([]string{"x86_64"})) + osNames := []string{ + "Amazon Linux", + "Amazon Linux 2", + "Red Hat", + "SUSE Linux", + "Ubuntu", + "Windows", + } + + imagesMap, err := h.GetLatestImages(nil, aws.StringSlice([]string{"x86_64"})) th.Ok(t, err) + + for _, os := range osNames { + th.Assert(t, (*imagesMap)[os] != nil, fmt.Sprintf("GetLatestImages should fetch image for %s", os)) + } } func TestGetDefaultImageForAmd(t *testing.T) { @@ -403,6 +417,20 @@ func ValidateInstanceTags(t *testing.T, actualInstanceTags []*ec2.Tag, launchReq th.Assert(t, countOfExpectedTags == countOfActualTagsMatched, "Didn't find all of the expected tags on the actual instance") } +func TestGetImageIdsFromSSM(t *testing.T) { + ssmClient := ssm.New(h.Sess) + validSsmPath := "/aws/service/ami-amazon-linux-latest" + imageIds, err := h.GetImageIdsFromSSM(ssmClient, validSsmPath) + th.Ok(t, err) + th.Assert(t, imageIds != nil, "imageIds should not be nil") + th.Assert(t, len(imageIds) > 0, "imageIds should not be empty") + + invalidSsmPath := "/aws/service/ami-amazon-linux" + _, err = h.GetImageIdsFromSSM(ssmClient, invalidSsmPath) + expectedErrorMsg := fmt.Sprintf("%s is not a valid namespace", invalidSsmPath[1:]) + th.Assert(t, strings.Contains(err.Error(), expectedErrorMsg), "Error message should contain invalid SSM path (namespace) information") +} + func TestTerminateInstances(t *testing.T) { th.Assert(t, h != nil, "EC2Helper was not initialized successfully") th.Assert(t, instanceId != nil, "No test instance ID found")