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

cartography/intel/aws/ec2/vpc_peerings.py

@@ -1,12 +1,19 @@
 import logging
+from typing import Any
 from typing import Dict
 from typing import List
+from typing import Tuple
 
 import boto3
 import neo4j
 
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.aws.ec2.vpc import AWSVpcSchema
+from cartography.models.aws.ec2.vpc_cidr import AWSIPv4CidrBlockSchema
+from cartography.models.aws.ec2.vpc_peering import AWSAccountVPCPeeringSchema
+from cartography.models.aws.ec2.vpc_peering import AWSPeeringConnectionSchema
 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
@@ -29,148 +36,258 @@ def get_vpc_peerings_data(
 
 
 @timeit
-def load_vpc_peerings(
+def transform_vpc_peering_data(
+    vpc_peerings: List[Dict],
+) -> Tuple[
+    List[Dict[str, Any]],
+    List[Dict[str, Any]],
+    List[Dict[str, Any]],
+    List[Dict[str, Any]],
+]:
+    transformed_peerings: List[Dict[str, Any]] = []
+    accepter_cidr_blocks: List[Dict[str, Any]] = []
+    requester_cidr_blocks: List[Dict[str, Any]] = []
+    vpc_nodes: List[Dict[str, Any]] = []
+
+    for peering in vpc_peerings:
+        accepter_cidr_ids: List[str] = []
+        for c_b in peering.get("AccepterVpcInfo", {}).get("CidrBlockSet", []):
+            block_id = f"{peering.get('AccepterVpcInfo', {}).get('VpcId')}|{c_b.get('CidrBlock')}"
+            accepter_cidr_blocks.append(
+                {
+                    "Id": block_id,
+                    "VpcId": peering.get("AccepterVpcInfo", {}).get("VpcId"),
+                    "AssociationId": c_b.get("AssociationId"),
+                    "CidrBlock": c_b.get("CidrBlock"),
+                    "BlockState": c_b.get("CidrBlockState", {}).get("State"),
+                    "BlockStateMessage": c_b.get("CidrBlockState", {}).get(
+                        "StatusMessage",
+                    ),
+                },
+            )
+            accepter_cidr_ids.append(block_id)
+
+        requester_cidr_ids: List[str] = []
+        for c_b in peering.get("RequesterVpcInfo", {}).get("CidrBlockSet", []):
+            block_id = f"{peering.get('RequesterVpcInfo', {}).get('VpcId')}|{c_b.get('CidrBlock')}"
+            requester_cidr_blocks.append(
+                {
+                    "Id": block_id,
+                    "VpcId": peering.get("RequesterVpcInfo", {}).get("VpcId"),
+                    "AssociationId": c_b.get("AssociationId"),
+                    "CidrBlock": c_b.get("CidrBlock"),
+                    "BlockState": c_b.get("CidrBlockState", {}).get("State"),
+                    "BlockStateMessage": c_b.get("CidrBlockState", {}).get(
+                        "StatusMessage",
+                    ),
+                },
+            )
+            requester_cidr_ids.append(block_id)
+
+        # Create VPC nodes for accepter and requester VPCs
+        accepter_vpc_id = peering.get("AccepterVpcInfo", {}).get("VpcId")
+        accepter_owner_id = peering.get("AccepterVpcInfo", {}).get("OwnerId")
+        if accepter_vpc_id:
+            vpc_nodes.append(
+                {
+                    "VpcId": accepter_vpc_id,
+                    "PrimaryCIDRBlock": None,  # VPCs from peering data don't have complete info
+                    "InstanceTenancy": None,
+                    "State": None,
+                    "IsDefault": None,
+                    "DhcpOptionsId": None,
+                    "AccountId": accepter_owner_id,  # Account that owns this VPC
+                }
+            )
+
+        requester_vpc_id = peering.get("RequesterVpcInfo", {}).get("VpcId")
+        requester_owner_id = peering.get("RequesterVpcInfo", {}).get("OwnerId")
+        if requester_vpc_id:
+            vpc_nodes.append(
+                {
+                    "VpcId": requester_vpc_id,
+                    "PrimaryCIDRBlock": None,  # VPCs from peering data don't have complete info
+                    "InstanceTenancy": None,
+                    "State": None,
+                    "IsDefault": None,
+                    "DhcpOptionsId": None,
+                    "AccountId": requester_owner_id,  # Account that owns this VPC
+                }
+            )
+
+        transformed_peerings.append(
+            {
+                "VpcPeeringConnectionId": peering.get("VpcPeeringConnectionId"),
+                "AllowDnsResolutionFromRemoteVpc": peering.get(
+                    "RequesterVpcInfo",
+                    {},
+                )
+                .get("PeeringOptions", {})
+                .get(
+                    "AllowDnsResolutionFromRemoteVpc",
+                ),
+                "AllowEgressFromLocalClassicLinkToRemoteVpc": peering.get(
+                    "RequesterVpcInfo",
+                    {},
+                )
+                .get("PeeringOptions", {})
+                .get(
+                    "AllowEgressFromLocalClassicLinkToRemoteVpc",
+                ),
+                "AllowEgressFromLocalVpcToRemoteClassicLink": peering.get(
+                    "RequesterVpcInfo",
+                    {},
+                )
+                .get("PeeringOptions", {})
+                .get(
+                    "AllowEgressFromLocalVpcToRemoteClassicLink",
+                ),
+                "RequesterRegion": peering.get("RequesterVpcInfo", {}).get(
+                    "Region",
+                ),
+                "AccepterRegion": peering.get("AccepterVpcInfo", {}).get(
+                    "Region",
+                ),
+                "StatusCode": peering.get("Status", {}).get("Code"),
+                "StatusMessage": peering.get("Status", {}).get("Message"),
+                "AccepterVpcId": peering.get("AccepterVpcInfo", {}).get("VpcId"),
+                "RequesterVpcId": peering.get("RequesterVpcInfo", {}).get(
+                    "VpcId",
+                ),
+                "ACCEPTER_CIDR_BLOCK_IDS": accepter_cidr_ids,
+                "REQUESTER_CIDR_BLOCK_IDS": requester_cidr_ids,
+            },
+        )
+
+    return transformed_peerings, accepter_cidr_blocks, requester_cidr_blocks, vpc_nodes
+
+
+@timeit
+def transform_aws_accounts_from_vpc_peering(
+    vpc_peerings: List[Dict],
+) -> List[Dict[str, Any]]:
+    """
+    Transform VPC peering data to extract AWS account information.
+    Creates composite AWS account nodes with context from VPC peering.
+    """
+    account_data: Dict[str, Dict[str, Any]] = {}
+
+    for peering in vpc_peerings:
+        # Extract accepter account
+        accepter_owner_id = peering.get("AccepterVpcInfo", {}).get("OwnerId")
+        if accepter_owner_id:
+            account_data[accepter_owner_id] = {
+                "id": accepter_owner_id,
+            }
+
+        # Extract requester account
+        requester_owner_id = peering.get("RequesterVpcInfo", {}).get("OwnerId")
+        if requester_owner_id:
+            account_data[requester_owner_id] = {
+                "id": requester_owner_id,
+            }
+
+    return list(account_data.values())
+
+
+@timeit
+def load_accepter_cidrs(
     neo4j_session: neo4j.Session,
-    data: List[Dict],
+    accepter_cidrs: List[Dict[str, Any]],
     region: str,
     aws_account_id: str,
     update_tag: int,
 ) -> None:
-    ingest_vpc_peerings = """
-    UNWIND $vpc_peerings AS vpc_peering
-
-    MERGE (pcx:AWSPeeringConnection{id: vpc_peering.VpcPeeringConnectionId})
-    ON CREATE SET pcx.firstseen = timestamp()
-    SET pcx.lastupdated = $update_tag,
-    pcx.allow_dns_resolution_from_remote_vpc =
-        vpc_peering.RequesterVpcInfo.PeeringOptions.AllowDnsResolutionFromRemoteVpc,
-    pcx.allow_egress_from_local_classic_link_to_remote_vpc =
-        vpc_peering.RequesterVpcInfo.PeeringOptions.AllowEgressFromLocalClassicLinkToRemoteVpc,
-    pcx.allow_egress_from_local_vpc_to_remote_classic_link =
-        vpc_peering.RequesterVpcInfo.PeeringOptions.AllowEgressFromLocalVpcToRemoteClassicLink,
-    pcx.requester_region = vpc_peering.RequesterVpcInfo.Region,
-    pcx.accepter_region = vpc_peering.AccepterVpcInfo.Region,
-    pcx.status_code = vpc_peering.Status.Code,
-    pcx.status_message = vpc_peering.Status.Message
-
-    MERGE (avpc:AWSVpc{id: vpc_peering.AccepterVpcInfo.VpcId})
-    ON CREATE SET avpc.firstseen = timestamp()
-    SET avpc.lastupdated = $update_tag, avpc.vpcid = vpc_peering.AccepterVpcInfo.VpcId
-
-    MERGE (rvpc:AWSVpc{id: vpc_peering.RequesterVpcInfo.VpcId})
-    ON CREATE SET rvpc.firstseen = timestamp()
-    SET rvpc.lastupdated = $update_tag, rvpc.vpcid = vpc_peering.RequesterVpcInfo.VpcId
-
-    MERGE (aaccount:AWSAccount{id: vpc_peering.AccepterVpcInfo.OwnerId})
-    ON CREATE SET aaccount.firstseen = timestamp()
-    SET aaccount.lastupdated = $update_tag
-
-    MERGE (raccount:AWSAccount{id: vpc_peering.RequesterVpcInfo.OwnerId})
-    ON CREATE SET raccount.firstseen = timestamp()
-    SET raccount.lastupdated = $update_tag
-
-    MERGE (pcx)-[rav:ACCEPTER_VPC]->(avpc)
-    ON CREATE SET rav.firstseen = timestamp()
-    SET rav.lastupdated = $update_tag
-
-    MERGE (aaccount)-[ra:RESOURCE]->(avpc)
-    ON CREATE SET ra.firstseen = timestamp()
-    SET ra.lastupdated = $update_tag
-
-    MERGE (pcx)-[rrv:REQUESTER_VPC]->(rvpc)
-    ON CREATE SET rrv.firstseen = timestamp()
-    SET rrv.lastupdated = $update_tag
-
-    MERGE (raccount)-[rr:RESOURCE]->(rvpc)
-    ON CREATE SET rr.firstseen = timestamp()
-    SET rr.lastupdated = $update_tag
-
-    """
-
-    neo4j_session.run(
-        ingest_vpc_peerings,
-        vpc_peerings=data,
-        update_tag=update_tag,
-        region=region,
-        aws_account_id=aws_account_id,
+    load(
+        neo4j_session,
+        AWSIPv4CidrBlockSchema(),
+        accepter_cidrs,
+        lastupdated=update_tag,
     )
 
 
 @timeit
-def load_accepter_cidrs(
+def load_requester_cidrs(
     neo4j_session: neo4j.Session,
-    data: List[Dict],
+    requester_cidrs: List[Dict[str, Any]],
     region: str,
     aws_account_id: str,
     update_tag: int,
 ) -> None:
+    load(
+        neo4j_session,
+        AWSIPv4CidrBlockSchema(),
+        requester_cidrs,
+        lastupdated=update_tag,
+    )
 
-    ingest_accepter_cidr = """
-    UNWIND $vpc_peerings AS vpc_peering
-    UNWIND vpc_peering.AccepterVpcInfo.CidrBlockSet AS c_b
-
-    MERGE (ac_b:AWSCidrBlock:AWSIpv4CidrBlock{id: vpc_peering.AccepterVpcInfo.VpcId + '|' + c_b.CidrBlock})
-    ON CREATE SET ac_b.firstseen = timestamp()
-    SET ac_b.lastupdated = $update_tag, ac_b.cidr_block  =  c_b.CidrBlock
-
-    WITH vpc_peering, ac_b
-    MATCH (pcx:AWSPeeringConnection{id: vpc_peering.VpcPeeringConnectionId}), (cb:AWSCidrBlock{id: ac_b.id})
-    MERGE (pcx)-[r:ACCEPTER_CIDR]->(cb)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-
-    WITH vpc_peering, ac_b
-    MATCH (vpc:AWSVpc{id: vpc_peering.AccepterVpcInfo.VpcId}), (cb:AWSCidrBlock{id: ac_b.id})
-    MERGE (vpc)-[r:BLOCK_ASSOCIATION]->(cb)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    """
 
-    neo4j_session.run(
-        ingest_accepter_cidr,
-        vpc_peerings=data,
-        update_tag=update_tag,
-        region=region,
-        aws_account_id=aws_account_id,
+@timeit
+def load_aws_accounts_from_vpc_peering(
+    neo4j_session: neo4j.Session,
+    aws_accounts: List[Dict[str, Any]],
+    update_tag: int,
+) -> None:
+    """
+    Load AWS account nodes using the composite schema.
+    This allows VPC peering data to provide additional context about AWS accounts.
+    """
+    load(
+        neo4j_session,
+        AWSAccountVPCPeeringSchema(),
+        aws_accounts,
+        lastupdated=update_tag,
     )
 
 
 @timeit
-def load_requester_cidrs(
+def load_vpc_nodes(
     neo4j_session: neo4j.Session,
-    data: List[Dict],
+    vpc_nodes: List[Dict[str, Any]],
     region: str,
     aws_account_id: str,
     update_tag: int,
 ) -> None:
+    # Group VPC nodes by their actual account ID
+    vpc_nodes_by_account: Dict[str, List[Dict[str, Any]]] = {}
+    for vpc in vpc_nodes:
+        account_id = vpc.get(
+            "AccountId", aws_account_id
+        )  # Use VPC's own account or fallback
+        if account_id not in vpc_nodes_by_account:
+            vpc_nodes_by_account[account_id] = []
+        vpc_nodes_by_account[account_id].append(vpc)
+
+    # Load VPCs for each account separately
+    for account_id, account_vpc_nodes in vpc_nodes_by_account.items():
+        # Remove the AccountId field as it's not part of the VPC schema
+        for vpc in account_vpc_nodes:
+            vpc.pop("AccountId", None)
+
+        load(
+            neo4j_session,
+            AWSVpcSchema(),
+            account_vpc_nodes,
+            lastupdated=update_tag,
+            AWS_ID=account_id,
+            Region=region,
+        )
 
-    ingest_requester_cidr = """
-    UNWIND $vpc_peerings AS vpc_peering
-    UNWIND vpc_peering.RequesterVpcInfo.CidrBlockSet AS c_b
-
-    MERGE (rc_b:AWSCidrBlock:AWSIpv4CidrBlock{id: vpc_peering.RequesterVpcInfo.VpcId + '|' + c_b.CidrBlock})
-    ON CREATE SET rc_b.firstseen = timestamp()
-    SET rc_b.lastupdated = $update_tag, rc_b.cidr_block  =  c_b.CidrBlock
-
-    WITH vpc_peering, rc_b
-    MATCH (pcx:AWSPeeringConnection{id: vpc_peering.VpcPeeringConnectionId}), (cb:AWSCidrBlock{id: rc_b.id})
-    MERGE (pcx)-[r:REQUESTER_CIDR]->(cb)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-
-    WITH vpc_peering, rc_b
-    MATCH (vpc:AWSVpc{id: vpc_peering.RequesterVpcInfo.VpcId}), (cb:AWSCidrBlock{id: rc_b.id})
-    MERGE (vpc)-[r:BLOCK_ASSOCIATION]->(cb)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    """
 
-    neo4j_session.run(
-        ingest_requester_cidr,
-        vpc_peerings=data,
-        update_tag=update_tag,
-        region=region,
-        aws_account_id=aws_account_id,
+@timeit
+def load_vpc_peerings(
+    neo4j_session: neo4j.Session,
+    vpc_peerings: List[Dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        AWSPeeringConnectionSchema(),
+        vpc_peerings,
+        lastupdated=update_tag,
+        AWS_ID=aws_account_id,
     )
 
 
@@ -179,11 +296,10 @@ def cleanup_vpc_peerings(
     neo4j_session: neo4j.Session,
     common_job_parameters: Dict,
 ) -> None:
-    run_cleanup_job(
-        "aws_import_vpc_peering_cleanup.json",
-        neo4j_session,
+    GraphJob.from_node_schema(
+        AWSPeeringConnectionSchema(),
         common_job_parameters,
-    )
+    ).run(neo4j_session)
 
 
 @timeit
@@ -201,24 +317,45 @@ def sync_vpc_peerings(
             region,
             current_aws_account_id,
         )
-        data = get_vpc_peerings_data(boto3_session, region)
-        load_vpc_peerings(
+        raw_data = get_vpc_peerings_data(boto3_session, region)
+        vpc_peerings, accepter_cidrs, requester_cidrs, vpc_nodes = (
+            transform_vpc_peering_data(
+                raw_data,
+            )
+        )
+        aws_accounts = transform_aws_accounts_from_vpc_peering(raw_data)
+
+        # Load AWS accounts first (composite pattern)
+        load_aws_accounts_from_vpc_peering(
+            neo4j_session,
+            aws_accounts,
+            update_tag,
+        )
+
+        load_vpc_nodes(
             neo4j_session,
-            data,
+            vpc_nodes,
             region,
             current_aws_account_id,
             update_tag,
         )
         load_accepter_cidrs(
             neo4j_session,
-            data,
+            accepter_cidrs,
             region,
             current_aws_account_id,
             update_tag,
         )
         load_requester_cidrs(
             neo4j_session,
-            data,
+            requester_cidrs,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
+        load_vpc_peerings(
+            neo4j_session,
+            vpc_peerings,
             region,
             current_aws_account_id,
             update_tag,