Skip to content

Commit

Permalink
Use public parameters in parameter store for ami id (#114)
Browse files Browse the repository at this point in the history 
* Use public parameters in parameter store for ami id
  • Loading branch information
@tiationg-kho
tiationg-kho authored Jan 2, 2025
1 parent 232ca66 commit d1d4c8b
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 42 deletions.
144 changes: 103 additions & 41 deletions pkg/ec2helper/ec2helper.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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

Expand All @@ -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{}
Expand Down
1 change: 1 addition & 0 deletions pkg/question/question_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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{
{
Expand Down
30 changes: 29 additions & 1 deletion test/e2e/e2e-ec2helper-test/e2e_ec2helper_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit d1d4c8b

Please sign in to comment.