Use public parameters in parameter store for ami id (#114)

* Use public parameters in parameter store for ami id

Author: tiationg-kho

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{}

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{
 			{

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")