cartography 0.104.0rc2__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/aws/ec2/load_balancer_v2s.py

@@ -6,6 +6,7 @@ import boto3
 import botocore
 import neo4j
 
+from cartography.client.core.tx import run_write_query
 from cartography.util import aws_handle_regions
 from cartography.util import run_cleanup_job
 from cartography.util import timeit
@@ -104,7 +105,8 @@ def load_load_balancer_v2s(
             logger.warning("Skipping load balancer entry with missing DNSName: %r", lb)
             continue
 
-        neo4j_session.run(
+        run_write_query(
+            neo4j_session,
             ingest_load_balancer_v2,
             ID=load_balancer_id,
             CREATED_TIME=str(lb["CreatedTime"]),
@@ -138,7 +140,8 @@ def load_load_balancer_v2s(
             SET r.lastupdated = $update_tag
             """
             for group in lb["SecurityGroups"]:
-                neo4j_session.run(
+                run_write_query(
+                    neo4j_session,
                     ingest_load_balancer_v2_security_group,
                     ID=load_balancer_id,
                     GROUP_ID=str(group),
@@ -182,7 +185,8 @@ def load_load_balancer_v2_subnets(
     SET r.lastupdated = $update_tag
     """
     for az in az_data:
-        neo4j_session.run(
+        run_write_query(
+            neo4j_session,
             ingest_load_balancer_subnet,
             ID=load_balancer_id,
             SubnetId=az["SubnetId"],
@@ -219,7 +223,8 @@ def load_load_balancer_v2_target_groups(
             continue
 
         for instance in target_group["Targets"]:
-            neo4j_session.run(
+            run_write_query(
+                neo4j_session,
                 ingest_instances,
                 ID=load_balancer_id,
                 INSTANCE_ID=instance,
@@ -253,7 +258,8 @@ def load_load_balancer_v2_listeners(
         ON CREATE SET r.firstseen = timestamp()
         SET r.lastupdated = $update_tag
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingest_listener,
         LoadBalancerId=load_balancer_id,
         Listeners=listener_data,

cartography/intel/aws/ec2/network_interfaces.py

@@ -91,13 +91,17 @@ def transform_network_interface_data(
                 "InterfaceType": network_interface["InterfaceType"],
                 "MacAddress": network_interface["MacAddress"],
                 "PrivateDnsName": network_interface.get("PrivateDnsName"),
-                "PrivateIpAddress": network_interface["PrivateIpAddress"],
+                "PrivateIpAddress": network_interface.get("PrivateIpAddress"),
                 "PublicIp": network_interface.get("Association", {}).get("PublicIp"),
                 "RequesterId": network_interface.get("RequesterId"),
                 "RequesterManaged": network_interface["RequesterManaged"],
                 "SourceDestCheck": network_interface["SourceDestCheck"],
                 "Status": network_interface["Status"],
                 "SubnetId": network_interface["SubnetId"],
+                "AttachTime": network_interface.get("Attachment", {}).get("AttachTime"),
+                "DeviceIndex": network_interface.get("Attachment", {}).get(
+                    "DeviceIndex"
+                ),
                 "ElbV1Id": elb_v1_id,
                 "ElbV2Id": elb_v2_id,
             },

cartography/intel/aws/ec2/reserved_instances.py

@@ -6,6 +6,7 @@ import boto3
 import neo4j
 from botocore.exceptions import ClientError
 
+from cartography.client.core.tx import run_write_query
 from cartography.util import aws_handle_regions
 from cartography.util import run_cleanup_job
 from cartography.util import timeit
@@ -64,7 +65,8 @@ def load_reserved_instances(
         r_instance["Start"] = str(r_instance["Start"])
         r_instance["End"] = str(r_instance["End"])
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingest_reserved_instances,
         reserved_instances_list=data,
         AWS_ACCOUNT_ID=current_aws_account_id,

cartography/intel/aws/ec2/security_groups.py

@@ -1,17 +1,22 @@
 import logging
-from string import Template
+from collections import namedtuple
+from typing import Any
 from typing import Dict
 from typing import List
 
 import boto3
 import neo4j
 
+from cartography.client.core.tx import load
 from cartography.graph.job import GraphJob
+from cartography.models.aws.ec2.security_group_rules import IpPermissionInboundSchema
+from cartography.models.aws.ec2.security_group_rules import IpRangeSchema
+from cartography.models.aws.ec2.security_group_rules import IpRuleSchema
+from cartography.models.aws.ec2.security_groups import EC2SecurityGroupSchema
 from cartography.models.aws.ec2.securitygroup_instance import (
     EC2SecurityGroupInstanceSchema,
 )
 from cartography.util import aws_handle_regions
-from cartography.util import run_cleanup_job
 from cartography.util import timeit
 
 from .util import get_botocore_config
@@ -37,138 +42,146 @@ def get_ec2_security_group_data(
     return security_groups
 
 
+Ec2SecurityGroupData = namedtuple(
+    "Ec2SecurityGroupData",
+    ["groups", "inbound_rules", "egress_rules", "ranges"],
+)
+
+
+def transform_ec2_security_group_data(
+    data: List[Dict[str, Any]],
+) -> Ec2SecurityGroupData:
+    groups: List[Dict[str, Any]] = []
+    inbound_rules: List[Dict[str, Any]] = []
+    egress_rules: List[Dict[str, Any]] = []
+    ranges: List[Dict[str, Any]] = []
+
+    for group in data:
+        group_record = {
+            "GroupId": group["GroupId"],
+            "GroupName": group.get("GroupName"),
+            "Description": group.get("Description"),
+            "VpcId": group.get("VpcId"),
+        }
+        # Collect referenced security groups for relationship loading
+        source_group_ids: set[str] = set()
+
+        for rule_type, target in (
+            ("IpPermissions", inbound_rules),
+            ("IpPermissionsEgress", egress_rules),
+        ):
+            for rule in group.get(rule_type, []):
+                protocol = rule.get("IpProtocol", "all")
+                from_port = rule.get("FromPort")
+                to_port = rule.get("ToPort")
+                rule_id = (
+                    f"{group['GroupId']}/{rule_type}/{from_port}{to_port}{protocol}"
+                )
+                target.append(
+                    {
+                        "RuleId": rule_id,
+                        "GroupId": group["GroupId"],
+                        "Protocol": protocol,
+                        "FromPort": from_port,
+                        "ToPort": to_port,
+                    },
+                )
+                for ip_range in rule.get("IpRanges", []):
+                    ranges.append({"RangeId": ip_range["CidrIp"], "RuleId": rule_id})
+                for pair in rule.get("UserIdGroupPairs", []):
+                    sg_id = pair.get("GroupId")
+                    if sg_id:
+                        source_group_ids.add(sg_id)
+
+        group_record["SOURCE_GROUP_IDS"] = list(source_group_ids)
+        groups.append(group_record)
+
+    return Ec2SecurityGroupData(
+        groups=groups,
+        inbound_rules=inbound_rules,
+        egress_rules=egress_rules,
+        ranges=ranges,
+    )
+
+
 @timeit
-def load_ec2_security_group_rule(
+def load_ip_rules(
     neo4j_session: neo4j.Session,
-    group: Dict,
-    rule_type: str,
+    data: List[Dict[str, Any]],
+    inbound: bool,
+    region: str,
+    aws_account_id: str,
     update_tag: int,
 ) -> None:
-    INGEST_RULE_TEMPLATE = Template(
-        """
-    MERGE (rule:$rule_label{ruleid: $RuleId})
-    ON CREATE SET rule :IpRule, rule.firstseen = timestamp(), rule.fromport = $FromPort, rule.toport = $ToPort,
-    rule.protocol = $Protocol
-    SET rule.lastupdated = $update_tag
-    WITH rule
-    MATCH (group:EC2SecurityGroup{groupid: $GroupId})
-    MERGE (group)<-[r:MEMBER_OF_EC2_SECURITY_GROUP]-(rule)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag;
-    """,
+    schema = IpPermissionInboundSchema() if inbound else IpRuleSchema()
+    load(
+        neo4j_session,
+        schema,
+        data,
+        Region=region,
+        AWS_ID=aws_account_id,
+        lastupdated=update_tag,
     )
 
-    ingest_rule_group_pair = """
-    MERGE (group:EC2SecurityGroup{id: $GroupId})
-    ON CREATE SET group.firstseen = timestamp(), group.groupid = $GroupId
-    SET group.lastupdated = $update_tag
-    WITH group
-    MATCH (inbound:IpRule{ruleid: $RuleId})
-    MERGE (inbound)-[r:MEMBER_OF_EC2_SECURITY_GROUP]->(group)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    """
-
-    ingest_range = """
-    MERGE (range:IpRange{id: $RangeId})
-    ON CREATE SET range.firstseen = timestamp(), range.range = $RangeId
-    SET range.lastupdated = $update_tag
-    WITH range
-    MATCH (rule:IpRule{ruleid: $RuleId})
-    MERGE (rule)<-[r:MEMBER_OF_IP_RULE]-(range)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    """
-
-    group_id = group["GroupId"]
-    rule_type_map = {
-        "IpPermissions": "IpPermissionInbound",
-        "IpPermissionsEgress": "IpPermissionEgress",
-    }
-
-    if group.get(rule_type):
-        for rule in group[rule_type]:
-            protocol = rule.get("IpProtocol", "all")
-            from_port = rule.get("FromPort")
-            to_port = rule.get("ToPort")
-
-            ruleid = f"{group_id}/{rule_type}/{from_port}{to_port}{protocol}"
-            # NOTE Cypher query syntax is incompatible with Python string formatting, so we have to do this awkward
-            # NOTE manual formatting instead.
-            neo4j_session.run(
-                INGEST_RULE_TEMPLATE.safe_substitute(
-                    rule_label=rule_type_map[rule_type],
-                ),
-                RuleId=ruleid,
-                FromPort=from_port,
-                ToPort=to_port,
-                Protocol=protocol,
-                GroupId=group_id,
-                update_tag=update_tag,
-            )
-
-            neo4j_session.run(
-                ingest_rule_group_pair,
-                GroupId=group_id,
-                RuleId=ruleid,
-                update_tag=update_tag,
-            )
-
-            for ip_range in rule["IpRanges"]:
-                range_id = ip_range["CidrIp"]
-                neo4j_session.run(
-                    ingest_range,
-                    RangeId=range_id,
-                    RuleId=ruleid,
-                    update_tag=update_tag,
-                )
+
+@timeit
+def load_ip_ranges(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        IpRangeSchema(),
+        data,
+        Region=region,
+        AWS_ID=aws_account_id,
+        lastupdated=update_tag,
+    )
 
 
 @timeit
 def load_ec2_security_groupinfo(
     neo4j_session: neo4j.Session,
-    data: List[Dict],
+    data: Ec2SecurityGroupData,
     region: str,
     current_aws_account_id: str,
     update_tag: int,
 ) -> None:
-    ingest_security_group = """
-    MERGE (group:EC2SecurityGroup{id: $GroupId})
-    ON CREATE SET group.firstseen = timestamp(), group.groupid = $GroupId
-    SET group.name = $GroupName, group.description = $Description, group.region = $Region,
-    group.lastupdated = $update_tag
-    WITH group
-    MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
-    MERGE (aa)-[r:RESOURCE]->(group)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    WITH group
-    MATCH (vpc:AWSVpc{id: $VpcId})
-    MERGE (vpc)-[rg:MEMBER_OF_EC2_SECURITY_GROUP]->(group)
-    ON CREATE SET rg.firstseen = timestamp()
-    """
-
-    for group in data:
-        group_id = group["GroupId"]
-
-        neo4j_session.run(
-            ingest_security_group,
-            GroupId=group_id,
-            GroupName=group.get("GroupName"),
-            Description=group.get("Description"),
-            VpcId=group.get("VpcId", None),
-            Region=region,
-            AWS_ACCOUNT_ID=current_aws_account_id,
-            update_tag=update_tag,
-        )
+    load(
+        neo4j_session,
+        EC2SecurityGroupSchema(),
+        data.groups,
+        Region=region,
+        AWS_ID=current_aws_account_id,
+        lastupdated=update_tag,
+    )
 
-        load_ec2_security_group_rule(neo4j_session, group, "IpPermissions", update_tag)
-        load_ec2_security_group_rule(
-            neo4j_session,
-            group,
-            "IpPermissionsEgress",
-            update_tag,
-        )
+    load_ip_rules(
+        neo4j_session,
+        data.inbound_rules,
+        inbound=True,
+        region=region,
+        aws_account_id=current_aws_account_id,
+        update_tag=update_tag,
+    )
+    load_ip_rules(
+        neo4j_session,
+        data.egress_rules,
+        inbound=False,
+        region=region,
+        aws_account_id=current_aws_account_id,
+        update_tag=update_tag,
+    )
+    load_ip_ranges(
+        neo4j_session,
+        data.ranges,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
 
 
 @timeit
@@ -176,11 +189,15 @@ def cleanup_ec2_security_groupinfo(
     neo4j_session: neo4j.Session,
     common_job_parameters: Dict,
 ) -> None:
-    run_cleanup_job(
-        "aws_import_ec2_security_groupinfo_cleanup.json",
-        neo4j_session,
+    GraphJob.from_node_schema(
+        EC2SecurityGroupSchema(),
         common_job_parameters,
+    ).run(neo4j_session)
+    GraphJob.from_node_schema(IpPermissionInboundSchema(), common_job_parameters).run(
+        neo4j_session,
     )
+    GraphJob.from_node_schema(IpRuleSchema(), common_job_parameters).run(neo4j_session)
+    GraphJob.from_node_schema(IpRangeSchema(), common_job_parameters).run(neo4j_session)
     GraphJob.from_node_schema(
         EC2SecurityGroupInstanceSchema(),
         common_job_parameters,
@@ -203,9 +220,10 @@ def sync_ec2_security_groupinfo(
             current_aws_account_id,
         )
         data = get_ec2_security_group_data(boto3_session, region)
+        transformed = transform_ec2_security_group_data(data)
         load_ec2_security_groupinfo(
             neo4j_session,
-            data,
+            transformed,
             region,
             current_aws_account_id,
             update_tag,

cartography/intel/aws/ec2/snapshots.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Any
 from typing import Dict
 from typing import List
 
@@ -6,8 +7,11 @@ import boto3
 import neo4j
 from botocore.exceptions import ClientError
 
+from cartography.client.core.tx import load
+from cartography.client.core.tx import read_list_of_values_tx
+from cartography.graph.job import GraphJob
+from cartography.models.aws.ec2.snapshots import EBSSnapshotSchema
 from cartography.util import aws_handle_regions
-from cartography.util import run_cleanup_job
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
@@ -24,12 +28,13 @@ def get_snapshots_in_use(
     WHERE v.region = $Region
     RETURN v.snapshotid as snapshot
     """
-    results = neo4j_session.run(
+    results = read_list_of_values_tx(
+        neo4j_session,
         query,
         AWS_ACCOUNT_ID=current_aws_account_id,
         Region=region,
     )
-    return [r["snapshot"] for r in results if r["snapshot"]]
+    return [str(snapshot) for snapshot in results if snapshot]
 
 
 @timeit
@@ -45,7 +50,6 @@ def get_snapshots(
     for page in paginator.paginate(OwnerIds=["self"]):
         snapshots.extend(page["Snapshots"])
 
-    # fetch in-use snapshots not in self_owned snapshots
     self_owned_snapshot_ids = {s["SnapshotId"] for s in snapshots}
     other_snapshot_ids = set(in_use_snapshot_ids) - self_owned_snapshot_ids
     if other_snapshot_ids:
@@ -55,8 +59,7 @@ def get_snapshots(
         except ClientError as e:
             if e.response["Error"]["Code"] == "InvalidSnapshot.NotFound":
                 logger.warning(
-                    f"Failed to retrieve page of in-use, \
-                    not owned snapshots. Continuing anyway. Error - {e}",
+                    f"Failed to retrieve page of in-use, not owned snapshots. Continuing anyway. Error - {e}"
                 )
             else:
                 raise
@@ -64,93 +67,53 @@ def get_snapshots(
     return snapshots
 
 
+def transform_snapshots(snapshots: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    transformed: List[Dict[str, Any]] = []
+    for snap in snapshots:
+        transformed.append(
+            {
+                "SnapshotId": snap["SnapshotId"],
+                "Description": snap.get("Description"),
+                "Encrypted": snap.get("Encrypted"),
+                "Progress": snap.get("Progress"),
+                "StartTime": snap.get("StartTime"),
+                "State": snap.get("State"),
+                "StateMessage": snap.get("StateMessage"),
+                "VolumeId": snap.get("VolumeId"),
+                "VolumeSize": snap.get("VolumeSize"),
+                "OutpostArn": snap.get("OutpostArn"),
+                "DataEncryptionKeyId": snap.get("DataEncryptionKeyId"),
+                "KmsKeyId": snap.get("KmsKeyId"),
+            }
+        )
+    return transformed
+
+
 @timeit
 def load_snapshots(
     neo4j_session: neo4j.Session,
-    data: List[Dict],
+    data: List[Dict[str, Any]],
     region: str,
     current_aws_account_id: str,
     update_tag: int,
 ) -> None:
-    ingest_snapshots = """
-    UNWIND $snapshots_list as snapshot
-        MERGE (s:EBSSnapshot{id: snapshot.SnapshotId})
-        ON CREATE SET s.firstseen = timestamp()
-        SET s.lastupdated = $update_tag, s.description = snapshot.Description, s.encrypted = snapshot.Encrypted,
-        s.progress = snapshot.Progress, s.starttime = snapshot.StartTime, s.state = snapshot.State,
-        s.statemessage = snapshot.StateMessage, s.volumeid = snapshot.VolumeId, s.volumesize = snapshot.VolumeSize,
-        s.outpostarn = snapshot.OutpostArn, s.dataencryptionkeyid = snapshot.DataEncryptionKeyId,
-        s.kmskeyid = snapshot.KmsKeyId, s.region=$Region
-        WITH s
-        MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
-        MERGE (aa)-[r:RESOURCE]->(s)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    for snapshot in data:
-        snapshot["StartTime"] = str(snapshot["StartTime"])
-
-    neo4j_session.run(
-        ingest_snapshots,
-        snapshots_list=data,
-        AWS_ACCOUNT_ID=current_aws_account_id,
+    load(
+        neo4j_session,
+        EBSSnapshotSchema(),
+        data,
+        lastupdated=update_tag,
         Region=region,
-        update_tag=update_tag,
-    )
-
-
-@timeit
-def get_snapshot_volumes(snapshots: List[Dict]) -> List[Dict]:
-    snapshot_volumes: List[Dict] = []
-    for snapshot in snapshots:
-        if snapshot.get("VolumeId"):
-            snapshot_volumes.append(snapshot)
-
-    return snapshot_volumes
-
-
-@timeit
-def load_snapshot_volume_relations(
-    neo4j_session: neo4j.Session,
-    data: List[Dict],
-    current_aws_account_id: str,
-    update_tag: int,
-) -> None:
-    ingest_volumes = """
-    UNWIND $snapshot_volumes_list as volume
-        MERGE (v:EBSVolume{id: volume.VolumeId})
-        ON CREATE SET v.firstseen = timestamp()
-        SET v.lastupdated = $update_tag, v.snapshotid = volume.SnapshotId
-        WITH v, volume
-        MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
-        MERGE (aa)-[r:RESOURCE]->(v)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-        WITH v, volume
-        MATCH (s:EBSSnapshot{id: volume.SnapshotId})
-        MERGE (s)-[r:CREATED_FROM]->(v)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    neo4j_session.run(
-        ingest_volumes,
-        snapshot_volumes_list=data,
-        AWS_ACCOUNT_ID=current_aws_account_id,
-        update_tag=update_tag,
+        AWS_ID=current_aws_account_id,
     )
 
 
 @timeit
 def cleanup_snapshots(
     neo4j_session: neo4j.Session,
-    common_job_parameters: Dict,
+    common_job_parameters: Dict[str, Any],
 ) -> None:
-    run_cleanup_job(
-        "aws_import_snapshots_cleanup.json",
-        neo4j_session,
-        common_job_parameters,
+    GraphJob.from_node_schema(EBSSnapshotSchema(), common_job_parameters).run(
+        neo4j_session
     )
 
 
@@ -161,7 +124,7 @@ def sync_ebs_snapshots(
     regions: List[str],
     current_aws_account_id: str,
     update_tag: int,
-    common_job_parameters: Dict,
+    common_job_parameters: Dict[str, Any],
 ) -> None:
     for region in regions:
         logger.debug(
@@ -174,12 +137,12 @@ def sync_ebs_snapshots(
             region,
             current_aws_account_id,
         )
-        data = get_snapshots(boto3_session, region, snapshots_in_use)
-        load_snapshots(neo4j_session, data, region, current_aws_account_id, update_tag)
-        snapshot_volumes = get_snapshot_volumes(data)
-        load_snapshot_volume_relations(
+        raw_data = get_snapshots(boto3_session, region, snapshots_in_use)
+        transformed_data = transform_snapshots(raw_data)
+        load_snapshots(
             neo4j_session,
-            snapshot_volumes,
+            transformed_data,
+            region,
             current_aws_account_id,
             update_tag,
         )