From 59238a40fd5e6c07466c89973202ae5e056e8504 Mon Sep 17 00:00:00 2001 From: Keep Focused Date: Wed, 5 Feb 2025 16:48:20 +0530 Subject: [PATCH] Added pagination to the aws_ec2_ami_shared table list operation and optimized the table to perform batch operation with given image IDs and owner IDs (#2260) --- aws/table_aws_ec2_ami_shared.go | 96 +++++++++++++++++++++----- docs/tables/aws_ec2_ami_shared.md | 111 +++++++++++++++++++++++++++++- 2 files changed, 189 insertions(+), 18 deletions(-) diff --git a/aws/table_aws_ec2_ami_shared.go b/aws/table_aws_ec2_ami_shared.go index 0b50ea7f6..338963eaf 100644 --- a/aws/table_aws_ec2_ami_shared.go +++ b/aws/table_aws_ec2_ami_shared.go @@ -2,6 +2,7 @@ package aws import ( "context" + "encoding/json" "errors" "fmt" "strings" @@ -16,6 +17,7 @@ import ( "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" "github.com/turbot/steampipe-plugin-sdk/v5/plugin" "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" + "github.com/turbot/steampipe-plugin-sdk/v5/query_cache" ) //// TABLE DEFINITION @@ -28,12 +30,14 @@ func tableAwsEc2AmiShared(_ context.Context) *plugin.Table { Hydrate: listAmisByOwner, Tags: map[string]string{"service": "ec2", "action": "DescribeImages"}, KeyColumns: []*plugin.KeyColumn{ - {Name: "owner_id", Require: plugin.Optional, CacheMatch: "exact"}, + {Name: "owner_id", Require: plugin.Optional, CacheMatch: query_cache.CacheMatchExact}, + {Name: "owner_ids", Require: plugin.Optional, CacheMatch: query_cache.CacheMatchExact, Operators: []string{"="}}, {Name: "architecture", Require: plugin.Optional}, {Name: "description", Require: plugin.Optional}, {Name: "ena_support", Require: plugin.Optional, Operators: []string{"=", "<>"}}, {Name: "hypervisor", Require: plugin.Optional}, - {Name: "image_id", Require: plugin.Optional, CacheMatch: "exact"}, + {Name: "image_id", Require: plugin.Optional, CacheMatch: query_cache.CacheMatchExact}, + {Name: "image_ids", Require: plugin.Optional, CacheMatch: query_cache.CacheMatchExact, Operators: []string{"="}}, {Name: "image_type", Require: plugin.Optional}, {Name: "public", Require: plugin.Optional, Operators: []string{"=", "<>"}}, {Name: "kernel_id", Require: plugin.Optional}, @@ -191,6 +195,18 @@ func tableAwsEc2AmiShared(_ context.Context) *plugin.Table { Description: "The type of virtualization of the AMI.", Type: proto.ColumnType_STRING, }, + { + Name: "image_ids", + Description: "The ID of the AMIs in the form of array of strings.", + Type: proto.ColumnType_JSON, + Transform: transform.FromQual("image_ids"), + }, + { + Name: "owner_ids", + Description: "The AWS account IDs of the image owners.", + Type: proto.ColumnType_JSON, + Transform: transform.FromQual("owner_ids"), + }, { Name: "block_device_mappings", Description: "Any block device mapping entries.", @@ -242,10 +258,22 @@ func tableAwsEc2AmiShared(_ context.Context) *plugin.Table { func listAmisByOwner(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) { owner_id := d.EqualsQuals["owner_id"].GetStringValue() image_id := d.EqualsQuals["image_id"].GetStringValue() + image_ids := d.EqualsQuals["image_ids"].GetJsonbValue() + owner_ids := d.EqualsQuals["owner_ids"].GetJsonbValue() + // check if owner_id and image_id is empty - if owner_id == "" && image_id == "" { - return nil, errors.New("please provide either owner_id or image_id") + if owner_id == "" && image_id == "" && image_ids == "" { + return nil, errors.New("please provide either owner_id, image_id or image_ids") + } + + // Limiting the results + maxLimit := int32(1000) + if d.QueryContext.Limit != nil { + limit := int32(*d.QueryContext.Limit) + if limit < maxLimit { + maxLimit = limit + } } // Create Session @@ -263,6 +291,22 @@ func listAmisByOwner(ctx context.Context, d *plugin.QueryData, h *plugin.Hydrate if image_id != "" { input.ImageIds = []string{image_id} } + if image_ids != "" { + var imageIds []string + err := json.Unmarshal([]byte(image_ids), &imageIds) + if err != nil { + return nil, errors.New("unable to parse the 'image_ids' query parameter the value must be in the format '[\"ami-000165ee3e0c1d6c7\", \"ami-0002ab43c99ec70ec\"]'") + } + input.ImageIds = imageIds + } + if owner_ids != "" { + var ownerIds []string + err := json.Unmarshal([]byte(owner_ids), &ownerIds) + if err != nil { + return nil, errors.New("unable to parse the 'image_ids' query parameter the value must be in the format '[\"123456789089\", \"345345678567\"]'") + } + input.Owners = ownerIds + } filters := buildSharedAmisWithOwnerFilter(d.Quals, ctx, d, h) @@ -270,23 +314,41 @@ func listAmisByOwner(ctx context.Context, d *plugin.QueryData, h *plugin.Hydrate input.Filters = filters } - // apply rate limiting - d.WaitForListRateLimit(ctx) + paginator := ec2.NewDescribeImagesPaginator(svc, input, func(o *ec2.DescribeImagesPaginatorOptions) { + // api error InvalidParameterCombination: The parameter imageIdsSet cannot be used with the parameter maxResults + if len(input.ImageIds) == 0 { + o.Limit = maxLimit + } + o.StopOnDuplicateToken = true + }) - // There is no MaxResult property in param, through which we can limit the number of results - resp, err := svc.DescribeImages(ctx, input) - if err != nil { - plugin.Logger(ctx).Error("aws_ec2_ami_shared.listAmisByOwner", "api_error", err) - return nil, err - } - for _, image := range resp.Images { - d.StreamListItem(ctx, image) + // List call + for paginator.HasMorePages() { + // apply rate limiting + d.WaitForListRateLimit(ctx) + + output, err := paginator.NextPage(ctx) + if err != nil { + plugin.Logger(ctx).Error("aws_ec2_instance.listEc2Instance", "api_error", err) + return nil, err + } + + for _, item := range output.Images { + + d.StreamListItem(ctx, item) + // Check if context has been cancelled or if the limit has been hit (if specified) + // if there is a limit, it will return the number of rows required to reach this limit + if d.RowsRemaining(ctx) == 0 { + return nil, nil + } - // Context may get cancelled due to manual cancellation or if the limit has been reached - if d.RowsRemaining(ctx) == 0 { - return nil, nil + // Context can be cancelled due to manual cancellation or the limit has been hit + if d.RowsRemaining(ctx) == 0 { + return nil, nil + } } } + return nil, err } diff --git a/docs/tables/aws_ec2_ami_shared.md b/docs/tables/aws_ec2_ami_shared.md index dcdf9b69a..8e7398633 100644 --- a/docs/tables/aws_ec2_ami_shared.md +++ b/docs/tables/aws_ec2_ami_shared.md @@ -14,6 +14,7 @@ The `aws_ec2_ami_shared` table in Steampipe provides you with information about **Important Notes** - You must specify an Owner ID or Image ID in the `where` clause (`where owner_id='`), (`where image_id='`). - The `aws_ec2_ami_shared` table can list any image but you must specify `owner_id` or `image_id`. +- To optimize query timing and API calls, use the optional query parameters `owner_ids` or `image_ids` to perform batch operations. - If you want to list all of the images in your account then you can use the `aws_ec2_ami` table. ## Examples @@ -117,4 +118,112 @@ from join aws_ec2_ami_shared as ami on i.image_id = ami.image_id where ami.owner_id = '137112412989'; -``` \ No newline at end of file +``` + +### Retrieve details of multiple shared AMIs in a single query +Fetches metadata of multiple shared AMIs, including their state, visibility, and creation details, to streamline AMI management. + +```sql+postgres +select + name, + image_id, + state, + image_location, + creation_date, + public, + root_device_name +from + aws_ec2_ami_shared +where + image_ids = '["ami-08df646e18b182346", "ami-04c5f154a6c2fec00",]'; +``` + +### Batch API operation, ensuring AMIs are from trusted sources +Any AWS customer can publish an Amazon Machine Image (AMI) for other AWS customers to launch instances from. AWS only vets a handful of images in the AWS Marketplace, there is no guarantee that other publicly shared AMIs are free of vulnerabilities or malicious code. While it's common for vendors to share their software as an AMI, it's also possible someone in your organization has launched an instance from a compromised image. + +```sql+postgres +with instances as ( + select + instance_id, + instance_type, + account_id, + tags ->> 'Name' as instance_name, + _ctx ->> 'connection_name' as account_name, + instance_state, + region, + image_id + from + aws_ec2_instance +), +all_image_ids as ( + select + json_agg(image_id)::jsonb as image_ids -- Cast to jsonb + from + instances +), +shared_ami as ( + select + s.* + from + aws_ec2_ami_shared as s, + all_image_ids + where s.image_ids = all_image_ids.image_ids +) +select distinct + shared_ami.image_id as image_id, + shared_ami.owner_id as image_owner_id, + shared_ami.image_owner_alias as image_owner_name, + instances.instance_name, + instances.account_name, + instances.region, + shared_ami.name as image_name +from + instances +left join shared_ami on shared_ami.image_id=instances.image_id +where shared_ami.image_owner_alias != 'amazon' +and shared_ami.image_owner_alias != 'self'; +``` + +```sql+sqlite +with instances as ( + select + instance_id, + instance_type, + account_id, + json_extract(tags, '$.Name') as instance_name, + json_extract(_ctx, '$.connection_name') as account_name, + instance_state, + region, + image_id + from + aws_ec2_instance +), +all_image_ids as ( + select + json_group_array(image_id) as image_ids + from + instances +), +shared_ami as ( + select + s.*, + ai.image_ids + from + aws_ec2_ami_shared as s, + all_image_ids as ai + where json_array_contains(ai.image_ids, s.image_id) +) +select distinct + shared_ami.image_id as image_id, + shared_ami.owner_id as image_owner_id, + shared_ami.image_owner_alias as image_owner_name, + instances.instance_name, + instances.account_name, + instances.region, + shared_ami.name as image_name +from + instances +left join shared_ami on shared_ami.image_id = instances.image_id +where shared_ami.image_owner_alias != 'amazon' + and shared_ami.image_owner_alias != 'self'; +```