Skip to content

Commit 73bb53a

Browse files
achantavyachantavy
authored
Update AWS EC2 keypair sync to use the new data model (#1424)
### Summary > Describe your changes. Updates EC2 key pair sync to use the data model. Fixes an issue that created duplicate key pairs because the older version of the key pair sync had extra node labels but the EC2KeyPairInstance model did not: ![image](https://github.com/user-attachments/assets/b9a004bc-de4a-458d-a7d1-9c6728548ba8) ### 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: - [x] Use the NodeSchema [data model](https://cartography-cncf.github.io/cartography/dev/writing-intel-modules.html#defining-a-node). --------- Signed-off-by: Alex Chantavy <[email protected]>
1 parent aacaba9 commit 73bb53a

File tree

5 files changed

+167
-72
lines changed

5 files changed

+167
-72
lines changed

‎cartography/intel/aws/ec2/instances.py

+5-4
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from cartography.intel.aws.ec2.util import get_botocore_config
1414
from cartography.models.aws.ec2.auto_scaling_groups import EC2InstanceAutoScalingGroupSchema
1515
from cartography.models.aws.ec2.instances import EC2InstanceSchema
16-
from cartography.models.aws.ec2.keypairs import EC2KeyPairSchema
16+
from cartography.models.aws.ec2.keypair_instance import EC2KeyPairInstanceSchema
1717
from cartography.models.aws.ec2.networkinterface_instance import EC2NetworkInterfaceInstanceSchema
1818
from cartography.models.aws.ec2.reservations import EC2ReservationSchema
1919
from cartography.models.aws.ec2.securitygroup_instance import EC2SecurityGroupInstanceSchema
@@ -193,16 +193,17 @@ def load_ec2_subnets(
193193

194194

195195
@timeit
196-
def load_ec2_key_pairs(
196+
def load_ec2_keypair_instances(
197197
neo4j_session: neo4j.Session,
198198
key_pair_list: List[Dict[str, Any]],
199199
region: str,
200200
current_aws_account_id: str,
201201
update_tag: int,
202202
) -> None:
203+
# Load EC2 keypairs as known by describe-instances.
203204
load(
204205
neo4j_session,
205-
EC2KeyPairSchema(),
206+
EC2KeyPairInstanceSchema(),
206207
key_pair_list,
207208
Region=region,
208209
AWS_ID=current_aws_account_id,
@@ -299,7 +300,7 @@ def load_ec2_instance_data(
299300
load_ec2_instance_nodes(neo4j_session, instance_list, region, current_aws_account_id, update_tag)
300301
load_ec2_subnets(neo4j_session, subnet_list, region, current_aws_account_id, update_tag)
301302
load_ec2_security_groups(neo4j_session, sg_list, region, current_aws_account_id, update_tag)
302-
load_ec2_key_pairs(neo4j_session, key_pair_list, region, current_aws_account_id, update_tag)
303+
load_ec2_keypair_instances(neo4j_session, key_pair_list, region, current_aws_account_id, update_tag)
303304
load_ec2_network_interfaces(neo4j_session, nic_list, region, current_aws_account_id, update_tag)
304305
load_ec2_instance_ebs_volumes(neo4j_session, ebs_volumes_list, region, current_aws_account_id, update_tag)
305306

‎cartography/intel/aws/ec2/key_pairs.py

+44-35
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
import logging
2+
from typing import Any
23
from typing import Dict
3-
from typing import List
44

55
import boto3
66
import neo4j
77

8-
from .util import get_botocore_config
8+
from cartography.client.core.tx import load
99
from cartography.graph.job import GraphJob
10-
from cartography.models.aws.ec2.keypairs import EC2KeyPairSchema
10+
from cartography.intel.aws.ec2.util import get_botocore_config
11+
from cartography.models.aws.ec2.keypair import EC2KeyPairSchema
1112
from cartography.util import aws_handle_regions
1213
from cartography.util import timeit
1314

@@ -16,42 +17,45 @@
1617

1718
@timeit
1819
@aws_handle_regions
19-
def get_ec2_key_pairs(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
20+
def get_ec2_key_pairs(boto3_session: boto3.session.Session, region: str) -> list[dict[str, Any]]:
2021
client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
2122
return client.describe_key_pairs()['KeyPairs']
2223

2324

25+
def transform_ec2_key_pairs(
26+
key_pairs: list[dict[str, Any]],
27+
region: str,
28+
current_aws_account_id: str,
29+
) -> list[dict[str, Any]]:
30+
transformed_key_pairs = []
31+
for key_pair in key_pairs:
32+
key_name = key_pair["KeyName"]
33+
transformed_key_pairs.append({
34+
'KeyPairArn': f'arn:aws:ec2:{region}:{current_aws_account_id}:key-pair/{key_name}',
35+
'KeyName': key_name,
36+
'KeyFingerprint': key_pair.get("KeyFingerprint"),
37+
})
38+
return transformed_key_pairs
39+
40+
2441
@timeit
2542
def load_ec2_key_pairs(
26-
neo4j_session: neo4j.Session, data: List[Dict], region: str, current_aws_account_id: str,
27-
update_tag: int,
43+
neo4j_session: neo4j.Session,
44+
data: list[dict[str, Any]],
45+
region: str,
46+
current_aws_account_id: str,
47+
update_tag: int,
2848
) -> None:
29-
ingest_key_pair = """
30-
MERGE (keypair:KeyPair:EC2KeyPair{arn: $ARN, id: $ARN})
31-
ON CREATE SET keypair.firstseen = timestamp()
32-
SET keypair.keyname = $KeyName, keypair.keyfingerprint = $KeyFingerprint, keypair.region = $Region,
33-
keypair.lastupdated = $update_tag
34-
WITH keypair
35-
MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
36-
MERGE (aa)-[r:RESOURCE]->(keypair)
37-
ON CREATE SET r.firstseen = timestamp()
38-
SET r.lastupdated = $update_tag
39-
"""
40-
41-
for key_pair in data:
42-
key_name = key_pair["KeyName"]
43-
key_fingerprint = key_pair.get("KeyFingerprint")
44-
key_pair_arn = f'arn:aws:ec2:{region}:{current_aws_account_id}:key-pair/{key_name}'
45-
46-
neo4j_session.run(
47-
ingest_key_pair,
48-
ARN=key_pair_arn,
49-
KeyName=key_name,
50-
KeyFingerprint=key_fingerprint,
51-
AWS_ACCOUNT_ID=current_aws_account_id,
52-
Region=region,
53-
update_tag=update_tag,
54-
)
49+
# Load EC2 keypairs as known by describe-key-pairs
50+
logger.info(f"Loading {len(data)} EC2 keypairs for region '{region}' into graph.")
51+
load(
52+
neo4j_session,
53+
EC2KeyPairSchema(),
54+
data,
55+
Region=region,
56+
AWS_ID=current_aws_account_id,
57+
lastupdated=update_tag,
58+
)
5559

5660

5761
@timeit
@@ -61,11 +65,16 @@ def cleanup_ec2_key_pairs(neo4j_session: neo4j.Session, common_job_parameters: D
6165

6266
@timeit
6367
def sync_ec2_key_pairs(
64-
neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
65-
update_tag: int, common_job_parameters: Dict,
68+
neo4j_session: neo4j.Session,
69+
boto3_session: boto3.session.Session,
70+
regions: list[str],
71+
current_aws_account_id: str,
72+
update_tag: int,
73+
common_job_parameters: dict[str, Any],
6674
) -> None:
6775
for region in regions:
6876
logger.info("Syncing EC2 key pairs for region '%s' in account '%s'.", region, current_aws_account_id)
6977
data = get_ec2_key_pairs(boto3_session, region)
70-
load_ec2_key_pairs(neo4j_session, data, region, current_aws_account_id, update_tag)
78+
transformed_data = transform_ec2_key_pairs(data, region, current_aws_account_id)
79+
load_ec2_key_pairs(neo4j_session, transformed_data, region, current_aws_account_id, update_tag)
7180
cleanup_ec2_key_pairs(neo4j_session, common_job_parameters)

‎cartography/models/aws/ec2/keypairs.py renamed to ‎cartography/models/aws/ec2/keypair.py

+13-23
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,23 @@
33
from cartography.models.core.common import PropertyRef
44
from cartography.models.core.nodes import CartographyNodeProperties
55
from cartography.models.core.nodes import CartographyNodeSchema
6+
from cartography.models.core.nodes import ExtraNodeLabels
67
from cartography.models.core.relationships import CartographyRelProperties
78
from cartography.models.core.relationships import CartographyRelSchema
89
from cartography.models.core.relationships import LinkDirection
910
from cartography.models.core.relationships import make_target_node_matcher
10-
from cartography.models.core.relationships import OtherRelationships
1111
from cartography.models.core.relationships import TargetNodeMatcher
1212

1313

1414
@dataclass(frozen=True)
1515
class EC2KeyPairNodeProperties(CartographyNodeProperties):
16+
"""
17+
Properties for EC2 keypairs from describe-key-pairs
18+
"""
1619
id: PropertyRef = PropertyRef('KeyPairArn')
1720
arn: PropertyRef = PropertyRef('KeyPairArn', extra_index=True)
1821
keyname: PropertyRef = PropertyRef('KeyName')
22+
keyfingerprint: PropertyRef = PropertyRef('KeyFingerprint')
1923
region: PropertyRef = PropertyRef('Region', set_in_kwargs=True)
2024
lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
2125

@@ -27,38 +31,24 @@ class EC2KeyPairToAwsAccountRelProperties(CartographyRelProperties):
2731

2832
@dataclass(frozen=True)
2933
class EC2KeyPairToAWSAccount(CartographyRelSchema):
34+
"""
35+
Relationship schema for EC2 keypairs to AWS Accounts
36+
"""
3037
target_node_label: str = 'AWSAccount'
3138
target_node_matcher: TargetNodeMatcher = make_target_node_matcher(
3239
{'id': PropertyRef('AWS_ID', set_in_kwargs=True)},
3340
)
3441
direction: LinkDirection = LinkDirection.INWARD
35-
rel_label: str = "RESOURCE"
42+
rel_label: str = 'RESOURCE'
3643
properties: EC2KeyPairToAwsAccountRelProperties = EC2KeyPairToAwsAccountRelProperties()
3744

3845

39-
@dataclass(frozen=True)
40-
class EC2KeyPairToEC2InstanceRelProperties(CartographyRelProperties):
41-
lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
42-
43-
44-
@dataclass(frozen=True)
45-
class EC2KeyPairToEC2Instance(CartographyRelSchema):
46-
target_node_label: str = 'EC2Instance'
47-
target_node_matcher: TargetNodeMatcher = make_target_node_matcher(
48-
{'id': PropertyRef('InstanceId')},
49-
)
50-
direction: LinkDirection = LinkDirection.OUTWARD
51-
rel_label: str = "SSH_LOGIN_TO"
52-
properties: EC2KeyPairToEC2InstanceRelProperties = EC2KeyPairToEC2InstanceRelProperties()
53-
54-
5546
@dataclass(frozen=True)
5647
class EC2KeyPairSchema(CartographyNodeSchema):
48+
"""
49+
Schema for EC2 keypairs from describe-key-pairs
50+
"""
5751
label: str = 'EC2KeyPair'
52+
extra_node_labels: ExtraNodeLabels = ExtraNodeLabels(['KeyPair'])
5853
properties: EC2KeyPairNodeProperties = EC2KeyPairNodeProperties()
5954
sub_resource_relationship: EC2KeyPairToAWSAccount = EC2KeyPairToAWSAccount()
60-
other_relationships: OtherRelationships = OtherRelationships(
61-
[
62-
EC2KeyPairToEC2Instance(),
63-
],
64-
)

‎cartography/models/aws/ec2/keypair_instance.py

+69
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
from dataclasses import dataclass
2+
3+
from cartography.models.core.common import PropertyRef
4+
from cartography.models.core.nodes import CartographyNodeProperties
5+
from cartography.models.core.nodes import CartographyNodeSchema
6+
from cartography.models.core.nodes import ExtraNodeLabels
7+
from cartography.models.core.relationships import CartographyRelProperties
8+
from cartography.models.core.relationships import CartographyRelSchema
9+
from cartography.models.core.relationships import LinkDirection
10+
from cartography.models.core.relationships import make_target_node_matcher
11+
from cartography.models.core.relationships import OtherRelationships
12+
from cartography.models.core.relationships import TargetNodeMatcher
13+
14+
15+
@dataclass(frozen=True)
16+
class EC2KeyPairInstanceNodeProperties(CartographyNodeProperties):
17+
id: PropertyRef = PropertyRef('KeyPairArn')
18+
arn: PropertyRef = PropertyRef('KeyPairArn', extra_index=True)
19+
keyname: PropertyRef = PropertyRef('KeyName')
20+
region: PropertyRef = PropertyRef('Region', set_in_kwargs=True)
21+
lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
22+
23+
24+
@dataclass(frozen=True)
25+
class EC2KeyPairInstanceToAwsAccountRelProperties(CartographyRelProperties):
26+
lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
27+
28+
29+
@dataclass(frozen=True)
30+
class EC2KeyPairInstanceToAWSAccount(CartographyRelSchema):
31+
target_node_label: str = 'AWSAccount'
32+
target_node_matcher: TargetNodeMatcher = make_target_node_matcher(
33+
{'id': PropertyRef('AWS_ID', set_in_kwargs=True)},
34+
)
35+
direction: LinkDirection = LinkDirection.INWARD
36+
rel_label: str = "RESOURCE"
37+
properties: EC2KeyPairInstanceToAwsAccountRelProperties = EC2KeyPairInstanceToAwsAccountRelProperties()
38+
39+
40+
@dataclass(frozen=True)
41+
class EC2KeyPairInstanceToEC2InstanceRelProperties(CartographyRelProperties):
42+
lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
43+
44+
45+
@dataclass(frozen=True)
46+
class EC2KeyPairInstanceToEC2Instance(CartographyRelSchema):
47+
target_node_label: str = 'EC2Instance'
48+
target_node_matcher: TargetNodeMatcher = make_target_node_matcher(
49+
{'id': PropertyRef('InstanceId')},
50+
)
51+
direction: LinkDirection = LinkDirection.OUTWARD
52+
rel_label: str = "SSH_LOGIN_TO"
53+
properties: EC2KeyPairInstanceToEC2InstanceRelProperties = EC2KeyPairInstanceToEC2InstanceRelProperties()
54+
55+
56+
@dataclass(frozen=True)
57+
class EC2KeyPairInstanceSchema(CartographyNodeSchema):
58+
"""
59+
EC2 keypairs as known by describe-instances.
60+
"""
61+
label: str = 'EC2KeyPair'
62+
extra_node_labels: ExtraNodeLabels = ExtraNodeLabels(['KeyPair'])
63+
properties: EC2KeyPairInstanceNodeProperties = EC2KeyPairInstanceNodeProperties()
64+
sub_resource_relationship: EC2KeyPairInstanceToAWSAccount = EC2KeyPairInstanceToAWSAccount()
65+
other_relationships: OtherRelationships = OtherRelationships(
66+
[
67+
EC2KeyPairInstanceToEC2Instance(),
68+
],
69+
)

‎tests/integration/cartography/intel/aws/ec2/test_ec2_key_pairs.py

+36-10
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
1+
from unittest.mock import MagicMock
2+
from unittest.mock import patch
3+
14
import cartography.intel.aws.ec2
25
import tests.data.aws.ec2.key_pairs
6+
from cartography.intel.aws.ec2.key_pairs import sync_ec2_key_pairs
37
from cartography.util import run_analysis_job
48

59

@@ -8,15 +12,26 @@
812
TEST_UPDATE_TAG = 123456789
913

1014

11-
def test_load_ec2_key_pairs(neo4j_session, *args):
12-
data = tests.data.aws.ec2.key_pairs.DESCRIBE_KEY_PAIRS['KeyPairs']
13-
cartography.intel.aws.ec2.key_pairs.load_ec2_key_pairs(
15+
@patch.object(
16+
cartography.intel.aws.ec2.key_pairs,
17+
'get_ec2_key_pairs',
18+
return_value=tests.data.aws.ec2.key_pairs.DESCRIBE_KEY_PAIRS['KeyPairs'],
19+
)
20+
def test_sync_ec2_key_pairs(mock_key_pairs, neo4j_session):
21+
# Arrange
22+
boto3_session = MagicMock()
23+
24+
# Act
25+
sync_ec2_key_pairs(
1426
neo4j_session,
15-
data,
16-
TEST_REGION,
27+
boto3_session,
28+
[TEST_REGION],
1729
TEST_ACCOUNT_ID,
1830
TEST_UPDATE_TAG,
31+
{'UPDATE_TAG': TEST_UPDATE_TAG, 'AWS_ID': TEST_ACCOUNT_ID},
1932
)
33+
34+
# Assert
2035
expected_nodes = {
2136
(
2237
"arn:aws:ec2:us-east-1:000000000000:key-pair/sample_key_pair_1",
@@ -51,20 +66,31 @@ def test_load_ec2_key_pairs(neo4j_session, *args):
5166
assert actual_nodes == expected_nodes
5267

5368

54-
def test_ec2_key_pairs_analysis_job(neo4j_session, *args):
55-
data = tests.data.aws.ec2.key_pairs.DESCRIBE_KEY_PAIRS['KeyPairs']
56-
cartography.intel.aws.ec2.key_pairs.load_ec2_key_pairs(
69+
@patch.object(
70+
cartography.intel.aws.ec2.key_pairs,
71+
'get_ec2_key_pairs',
72+
return_value=tests.data.aws.ec2.key_pairs.DESCRIBE_KEY_PAIRS['KeyPairs'],
73+
)
74+
def test_ec2_key_pairs_analysis_job(mock_key_pairs, neo4j_session):
75+
# Arrange
76+
boto3_session = MagicMock()
77+
sync_ec2_key_pairs(
5778
neo4j_session,
58-
data,
59-
TEST_REGION,
79+
boto3_session,
80+
[TEST_REGION],
6081
TEST_ACCOUNT_ID,
6182
TEST_UPDATE_TAG,
83+
{'UPDATE_TAG': TEST_UPDATE_TAG, 'AWS_ID': TEST_ACCOUNT_ID},
6284
)
85+
86+
# Act
6387
run_analysis_job(
6488
'aws_ec2_keypair_analysis.json',
6589
neo4j_session,
6690
{'UPDATE_TAG': TEST_UPDATE_TAG},
6791
)
92+
93+
# Assert
6894
expected_nodes = {
6995
(
7096
"arn:aws:ec2:us-east-1:000000000000:key-pair/sample_key_pair_4",

0 commit comments

Comments
 (0)