From 0e634fa3cadce280030fb4c43df6b3d02964679f Mon Sep 17 00:00:00 2001 From: Hector Eryx Paredes Camacho Date: Mon, 20 Jan 2025 12:09:04 -0600 Subject: [PATCH] Add cross region support for EC2 images import (#1441) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary This covers an edge case where the EC2 instance exists in one region, but the AMI is from another region . There isn't much documentation available, but I noticed it more on Amazon Linux recommended AMIs for EKS. ### Related issues or links - https://winder.ai/how-to-list-all-amis-for-each-region-in-aws/ - https://docs.aws.amazon.com/eks/latest/userguide/retrieve-ami-id.html - https://docs.aws.amazon.com/eks/latest/userguide/eks-optimized-ami.html With no multi region support: ![Screenshot 2025-01-16 at 2 08 54 p m](https://github.com/user-attachments/assets/0e58e4d2-e083-4886-ad82-9f4c8add624c) With multi region support ![Screenshot 2025-01-16 at 2 08 43 p m](https://github.com/user-attachments/assets/8a30cb2c-8d3c-41f2-bb83-4c89a6e63bdf) ### Checklist Provide proof that this works (this makes reviews move faster). Please perform one or more of the following: - [x] Update/add unit or integration tests. - [ ] Include a screenshot showing what the graph looked like before and after your changes. - [ ] Include console log trace showing what happened before and after your changes. If you are changing a node or relationship: - [ ] Update the [schema](https://github.com/lyft/cartography/tree/master/docs/root/modules) and [readme](https://github.com/lyft/cartography/blob/master/docs/schema/README.md). If you are implementing a new intel module: - [ ] Use the NodeSchema [data model](https://cartography-cncf.github.io/cartography/dev/writing-intel-modules.html#defining-a-node). --------- Signed-off-by: Alex Chantavy Signed-off-by: Eryx Paredes Co-authored-by: Alex Chantavy Co-authored-by: i_virus --- cartography/intel/aws/ec2/images.py | 22 +++++- .../intel/aws/ec2/test_ec2_images.py | 72 +++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) create mode 100644 tests/unit/cartography/intel/aws/ec2/test_ec2_images.py diff --git a/cartography/intel/aws/ec2/images.py b/cartography/intel/aws/ec2/images.py index d5af19d657..ecac08db80 100644 --- a/cartography/intel/aws/ec2/images.py +++ b/cartography/intel/aws/ec2/images.py @@ -9,6 +9,7 @@ from cartography.client.core.tx import load from cartography.graph.job import GraphJob +from cartography.intel.aws.ec2 import get_ec2_regions from cartography.intel.aws.ec2.util import get_botocore_config from cartography.models.aws.ec2.images import EC2ImageSchema from cartography.util import aws_handle_regions @@ -48,7 +49,7 @@ def get_images(boto3_session: boto3.session.Session, region: str, image_ids: Lis self_images = client.describe_images(Owners=['self'])['Images'] images.extend(self_images) except ClientError as e: - logger.warning(f"Failed retrieve images for region - {region}. Error - {e}") + logger.warning(f"Failed to retrieve private images for region - {region}. Error - {e}") try: if image_ids: image_ids = [image_id for image_id in image_ids if image_id is not None] @@ -58,8 +59,25 @@ def get_images(boto3_session: boto3.session.Session, region: str, image_ids: Lis for image in images_in_use: if image["ImageId"] not in _ids: images.append(image) + _ids.append(image["ImageId"]) + # Handle cross region image ids + if len(_ids) != len(image_ids): + logger.info("Attempting to retrieve images from other regions") + pending_ids = [image_id for image_id in image_ids if image_id not in _ids] + all_regions = get_ec2_regions(boto3_session) + clients = { + other_region: boto3_session.client('ec2', region_name=other_region, config=get_botocore_config()) + for other_region in all_regions if other_region != region + } + for other_region, client in clients.items(): + for _id in pending_ids: + try: + pending_image = client.describe_images(ImageIds=[_id])['Images'] + images.extend(pending_image) + except ClientError as e: + logger.warning(f"Image {id} could not be found at region - {other_region}. Error - {e}") except ClientError as e: - logger.warning(f"Failed retrieve images for region - {region}. Error - {e}") + logger.warning(f"Failed to retrieve public images for region - {region}. Error - {e}") return images diff --git a/tests/unit/cartography/intel/aws/ec2/test_ec2_images.py b/tests/unit/cartography/intel/aws/ec2/test_ec2_images.py new file mode 100644 index 0000000000..1914fffb92 --- /dev/null +++ b/tests/unit/cartography/intel/aws/ec2/test_ec2_images.py @@ -0,0 +1,72 @@ +from unittest.mock import MagicMock +from unittest.mock import patch + +from botocore.exceptions import ClientError + +from cartography.intel.aws.ec2.images import get_images + + +@patch('cartography.intel.aws.ec2.images.get_botocore_config') +@patch('cartography.intel.aws.ec2.images.get_ec2_regions') +@patch('boto3.session.Session') +def test_get_images_all_sources(mock_boto3_session, mock_get_ec2_regions, mock_get_botocore_config): + region = 'us-east-1' + image_ids = ['ami-55555555', 'ami-12345678', 'ami-87654321'] + + mock_client = MagicMock() + mock_boto3_session.client.return_value = mock_client + mock_client.describe_images.side_effect = [ + {'Images': [{'ImageId': 'ami-55555555'}]}, + {'Images': [{'ImageId': 'ami-12345678'}]}, + {'Images': [{'ImageId': 'ami-87654321'}]}, + ] + + mock_get_ec2_regions.return_value = ['us-east-1', 'us-west-2'] + mock_get_botocore_config.return_value = {} + + result = get_images(mock_boto3_session, region, image_ids) + + assert len(result) == 3 + assert result[0]['ImageId'] == 'ami-55555555' + assert result[1]['ImageId'] == 'ami-12345678' + assert result[2]['ImageId'] == 'ami-87654321' + + +@patch('cartography.intel.aws.ec2.images.get_botocore_config') +@patch('boto3.session.Session') +def test_get_images_not_found(mock_boto3_session, mock_get_botocore_config): + region = 'us-west-2' + image_ids = ['ami-12345678', 'ami-87654321'] + + mock_client = MagicMock() + mock_boto3_session.return_value = mock_client + mock_client.describe_images.side_effect = ClientError( + error_response={'Error': {'Code': 'InvalidAMIID.NotFound', 'Message': 'The image id does not exist'}}, + operation_name='DescribeImages', + ) + + mock_get_botocore_config.return_value = {} + + result = get_images(mock_boto3_session, region, image_ids) + + assert result == [] + + +@patch('cartography.intel.aws.ec2.images.get_botocore_config') +@patch('boto3.session.Session') +def test_get_images_no_image_ids(mock_boto3_session, mock_get_botocore_config): + region = 'us-west-2' + image_ids = [] + + mock_client = MagicMock() + mock_boto3_session.client.return_value = mock_client + mock_client.describe_images.side_effect = [ + {'Images': [{'ImageId': 'ami-12345678'}]}, + ] + + mock_get_botocore_config.return_value = {} + + result = get_images(mock_boto3_session, region, image_ids) + + assert len(result) == 1 + assert result[0]['ImageId'] == 'ami-12345678'