Skip to content

Commit

Permalink
Added pagination to the aws_ec2_ami_shared table list operation and o…
Browse files Browse the repository at this point in the history 
…ptimized the table to perform batch operation with given image IDs and owner IDs (#2260)
  • Loading branch information
@ParthaI
ParthaI authored Feb 5, 2025
1 parent e0e8012 commit 59238a4
Show file tree
Hide file tree
Showing 2 changed files with 189 additions and 18 deletions.
96 changes: 79 additions & 17 deletions aws/table_aws_ec2_ami_shared.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package aws

import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
Expand All @@ -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
Expand All @@ -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},
Expand Down Expand Up @@ -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.",
Expand Down Expand Up @@ -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
Expand All @@ -263,30 +291,64 @@ 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)

if len(filters) != 0 {
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
}

Expand Down
111 changes: 110 additions & 1 deletion docs/tables/aws_ec2_ami_shared.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -117,4 +118,112 @@ from
join aws_ec2_ami_shared as ami on i.image_id = ami.image_id
where
ami.owner_id = '137112412989';
```
```

### 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';
```

0 comments on commit 59238a4

Please sign in to comment.