cartography 0.107.0rc3__py3-none-any.whl → 0.108.0__py3-none-any.whl

cartography/intel/aws/ec2/subnets.py

@@ -1,17 +1,17 @@
 import logging
-from typing import Dict
-from typing import List
+from typing import Any
 
 import boto3
 import neo4j
 
+from cartography.client.core.tx import load
 from cartography.graph.job import GraphJob
 from cartography.models.aws.ec2.auto_scaling_groups import (
     EC2SubnetAutoScalingGroupSchema,
 )
 from cartography.models.aws.ec2.subnet_instance import EC2SubnetInstanceSchema
+from cartography.models.aws.ec2.subnets import EC2SubnetSchema
 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
@@ -21,86 +21,53 @@ logger = logging.getLogger(__name__)
 
 @timeit
 @aws_handle_regions
-def get_subnet_data(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
+def get_subnet_data(
+    boto3_session: boto3.session.Session, region: str
+) -> list[dict[str, Any]]:
     client = boto3_session.client(
         "ec2",
         region_name=region,
         config=get_botocore_config(),
     )
     paginator = client.get_paginator("describe_subnets")
-    subnets: List[Dict] = []
+    subnets: list[dict[str, Any]] = []
     for page in paginator.paginate():
         subnets.extend(page["Subnets"])
     return subnets
 
 
+def transform_subnet_data(subnets: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Transform subnet data into a loadable format."""
+    transformed: list[dict[str, Any]] = []
+    for subnet in subnets:
+        transformed.append(subnet.copy())
+    return transformed
+
+
 @timeit
 def load_subnets(
     neo4j_session: neo4j.Session,
-    data: List[Dict],
+    data: list[dict[str, Any]],
     region: str,
     aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
-
-    ingest_subnets = """
-    UNWIND $subnets as subnet
-    MERGE (snet:EC2Subnet{subnetid: subnet.SubnetId})
-    ON CREATE SET snet.firstseen = timestamp()
-    SET snet.lastupdated = $aws_update_tag, snet.name = subnet.CidrBlock, snet.cidr_block = subnet.CidrBlock,
-    snet.available_ip_address_count = subnet.AvailableIpAddressCount, snet.default_for_az = subnet.DefaultForAz,
-    snet.map_customer_owned_ip_on_launch = subnet.MapCustomerOwnedIpOnLaunch,
-    snet.state = subnet.State, snet.assignipv6addressoncreation = subnet.AssignIpv6AddressOnCreation,
-    snet.map_public_ip_on_launch = subnet.MapPublicIpOnLaunch, snet.subnet_arn = subnet.SubnetArn,
-    snet.availability_zone = subnet.AvailabilityZone, snet.availability_zone_id = subnet.AvailabilityZoneId,
-    snet.subnet_id = subnet.SubnetId
-    """
-
-    ingest_subnet_vpc_relations = """
-    UNWIND $subnets as subnet
-    MATCH (snet:EC2Subnet{subnetid: subnet.SubnetId}), (vpc:AWSVpc{id: subnet.VpcId})
-    MERGE (snet)-[r:MEMBER_OF_AWS_VPC]->(vpc)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $aws_update_tag
-    """
-
-    ingest_subnet_aws_account_relations = """
-    UNWIND $subnets as subnet
-    MATCH (snet:EC2Subnet{subnetid: subnet.SubnetId}), (aws:AWSAccount{id: $aws_account_id})
-    MERGE (aws)-[r:RESOURCE]->(snet)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $aws_update_tag
-    """
-
-    neo4j_session.run(
-        ingest_subnets,
-        subnets=data,
-        aws_update_tag=aws_update_tag,
-        region=region,
-        aws_account_id=aws_account_id,
-    )
-    neo4j_session.run(
-        ingest_subnet_vpc_relations,
-        subnets=data,
-        aws_update_tag=aws_update_tag,
-        region=region,
-        aws_account_id=aws_account_id,
-    )
-    neo4j_session.run(
-        ingest_subnet_aws_account_relations,
-        subnets=data,
-        aws_update_tag=aws_update_tag,
-        region=region,
-        aws_account_id=aws_account_id,
+    load(
+        neo4j_session,
+        EC2SubnetSchema(),
+        data,
+        Region=region,
+        AWS_ID=aws_account_id,
+        lastupdated=aws_update_tag,
     )
 
 
 @timeit
-def cleanup_subnets(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job(
-        "aws_ingest_subnets_cleanup.json",
+def cleanup_subnets(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(EC2SubnetSchema(), common_job_parameters).run(
         neo4j_session,
-        common_job_parameters,
     )
     GraphJob.from_node_schema(EC2SubnetInstanceSchema(), common_job_parameters).run(
         neo4j_session,
@@ -115,10 +82,10 @@ def cleanup_subnets(neo4j_session: neo4j.Session, common_job_parameters: Dict) -
 def sync_subnets(
     neo4j_session: neo4j.Session,
     boto3_session: boto3.session.Session,
-    regions: List[str],
+    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.info(
@@ -127,5 +94,12 @@ def sync_subnets(
             current_aws_account_id,
         )
         data = get_subnet_data(boto3_session, region)
-        load_subnets(neo4j_session, data, region, current_aws_account_id, update_tag)
+        transformed = transform_subnet_data(data)
+        load_subnets(
+            neo4j_session,
+            transformed,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_subnets(neo4j_session, common_job_parameters)

cartography/intel/aws/ecr.py

@@ -6,9 +6,12 @@ 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.ecr.image import ECRImageSchema
+from cartography.models.aws.ecr.repository import ECRRepositorySchema
+from cartography.models.aws.ecr.repository_image import ECRRepositoryImageSchema
 from cartography.util import aws_handle_regions
-from cartography.util import batch
-from cartography.util import run_cleanup_job
 from cartography.util import timeit
 from cartography.util import to_asynchronous
 from cartography.util import to_synchronous
@@ -74,33 +77,17 @@ def load_ecr_repositories(
     current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
-    query = """
-    UNWIND $Repositories as ecr_repo
-        MERGE (repo:ECRRepository{id: ecr_repo.repositoryArn})
-        ON CREATE SET repo.firstseen = timestamp(),
-            repo.arn = ecr_repo.repositoryArn,
-            repo.name = ecr_repo.repositoryName,
-            repo.region = $Region,
-            repo.created_at = ecr_repo.createdAt
-        SET repo.lastupdated = $aws_update_tag,
-            repo.uri = ecr_repo.repositoryUri
-        WITH repo
-
-        MATCH (owner:AWSAccount{id: $AWS_ACCOUNT_ID})
-        MERGE (owner)-[r:RESOURCE]->(repo)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $aws_update_tag
-    """
     logger.info(
         f"Loading {len(repos)} ECR repositories for region {region} into graph.",
     )
-    neo4j_session.run(
-        query,
-        Repositories=repos,
+    load(
+        neo4j_session,
+        ECRRepositorySchema(),
+        repos,
+        lastupdated=aws_update_tag,
         Region=region,
-        aws_update_tag=aws_update_tag,
-        AWS_ACCOUNT_ID=current_aws_account_id,
-    ).consume()  # See issue #440
+        AWS_ID=current_aws_account_id,
+    )
 
 
 @timeit
@@ -114,8 +101,13 @@ def transform_ecr_repository_images(repo_data: Dict) -> List[Dict]:
     for repo_uri in sorted(repo_data.keys()):
         repo_images = repo_data[repo_uri]
         for img in repo_images:
-            if "imageDigest" in img and img["imageDigest"]:
+            digest = img.get("imageDigest")
+            if digest:
+                tag = img.get("imageTag")
+                uri = repo_uri + (f":{tag}" if tag else "")
                 img["repo_uri"] = repo_uri
+                img["uri"] = uri
+                img["id"] = uri
                 repo_images_list.append(img)
             else:
                 logger.warning(
@@ -127,74 +119,51 @@ def transform_ecr_repository_images(repo_data: Dict) -> List[Dict]:
     return repo_images_list
 
 
-def _load_ecr_repo_img_tx(
-    tx: neo4j.Transaction,
-    repo_images_list: List[Dict],
-    aws_update_tag: int,
-    region: str,
-) -> None:
-    query = """
-    UNWIND $RepoList as repo_img
-        MERGE (ri:ECRRepositoryImage{id: repo_img.repo_uri + COALESCE(":" + repo_img.imageTag, '')})
-        ON CREATE SET ri.firstseen = timestamp()
-        SET ri.lastupdated = $aws_update_tag,
-            ri.tag = repo_img.imageTag,
-            ri.uri = repo_img.repo_uri + COALESCE(":" + repo_img.imageTag, ''),
-            ri.image_size_bytes = repo_img.imageSizeInBytes,
-            ri.image_pushed_at = repo_img.imagePushedAt,
-            ri.image_manifest_media_type = repo_img.imageManifestMediaType,
-            ri.artifact_media_type = repo_img.artifactMediaType,
-            ri.last_recorded_pull_time = repo_img.lastRecordedPullTime
-        WITH ri, repo_img
-
-        MERGE (img:ECRImage{id: repo_img.imageDigest})
-        ON CREATE SET img.firstseen = timestamp(),
-            img.digest = repo_img.imageDigest
-        SET img.lastupdated = $aws_update_tag,
-            img.region = $Region
-        WITH ri, img, repo_img
-
-        MERGE (ri)-[r1:IMAGE]->(img)
-        ON CREATE SET r1.firstseen = timestamp()
-        SET r1.lastupdated = $aws_update_tag
-        WITH ri, repo_img
-
-        MATCH (repo:ECRRepository{uri: repo_img.repo_uri})
-        MERGE (repo)-[r2:REPO_IMAGE]->(ri)
-        ON CREATE SET r2.firstseen = timestamp()
-        SET r2.lastupdated = $aws_update_tag
-    """
-    tx.run(
-        query,
-        RepoList=repo_images_list,
-        Region=region,
-        aws_update_tag=aws_update_tag,
-    )
-
-
 @timeit
 def load_ecr_repository_images(
     neo4j_session: neo4j.Session,
     repo_images_list: List[Dict],
     region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     logger.info(
         f"Loading {len(repo_images_list)} ECR repository images in {region} into graph.",
     )
-    for repo_image_batch in batch(repo_images_list, size=10000):
-        neo4j_session.write_transaction(
-            _load_ecr_repo_img_tx,
-            repo_image_batch,
-            aws_update_tag,
-            region,
-        )
+    image_digests = {img["imageDigest"] for img in repo_images_list}
+    ecr_images = [{"imageDigest": d} for d in image_digests]
+
+    load(
+        neo4j_session,
+        ECRImageSchema(),
+        ecr_images,
+        lastupdated=aws_update_tag,
+        Region=region,
+        AWS_ID=current_aws_account_id,
+    )
+
+    load(
+        neo4j_session,
+        ECRRepositoryImageSchema(),
+        repo_images_list,
+        lastupdated=aws_update_tag,
+        Region=region,
+        AWS_ID=current_aws_account_id,
+    )
 
 
 @timeit
 def cleanup(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
     logger.debug("Running ECR cleanup job.")
-    run_cleanup_job("aws_import_ecr_cleanup.json", neo4j_session, common_job_parameters)
+    GraphJob.from_node_schema(ECRRepositorySchema(), common_job_parameters).run(
+        neo4j_session
+    )
+    GraphJob.from_node_schema(ECRRepositoryImageSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+    GraphJob.from_node_schema(ECRImageSchema(), common_job_parameters).run(
+        neo4j_session
+    )
 
 
 def _get_image_data(
@@ -251,5 +220,11 @@ def sync(
             update_tag,
         )
         repo_images_list = transform_ecr_repository_images(image_data)
-        load_ecr_repository_images(neo4j_session, repo_images_list, region, update_tag)
+        load_ecr_repository_images(
+            neo4j_session,
+            repo_images_list,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup(neo4j_session, common_job_parameters)

cartography/intel/aws/elasticache.py

@@ -1,118 +1,132 @@
 import logging
-from typing import Dict
-from typing import List
-from typing import Set
+from typing import Any
 
 import boto3
 import neo4j
 
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.aws.elasticache.cluster import ElasticacheClusterSchema
+from cartography.models.aws.elasticache.topic import ElasticacheTopicSchema
 from cartography.stats import get_stats_client
 from cartography.util import aws_handle_regions
 from cartography.util import merge_module_sync_metadata
-from cartography.util import run_cleanup_job
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
 stat_handler = get_stats_client(__name__)
 
 
-def _get_topic(cluster: Dict) -> Dict:
-    return cluster["NotificationConfiguration"]
-
-
-def transform_elasticache_topics(cluster_data: List[Dict]) -> List[Dict]:
-    """
-    Collect unique TopicArns from the cluster data
-    """
-    seen: Set[str] = set()
-    topics: List[Dict] = []
-    for cluster in cluster_data:
-        topic = _get_topic(cluster)
-        topic_arn = topic["TopicArn"]
-        if topic_arn not in seen:
-            seen.add(topic_arn)
-            topics.append(topic)
-    return topics
-
-
 @timeit
 @aws_handle_regions
 def get_elasticache_clusters(
     boto3_session: boto3.session.Session,
     region: str,
-) -> List[Dict]:
-    logger.debug(f"Getting ElastiCache Clusters in region '{region}'.")
+) -> list[dict[str, Any]]:
     client = boto3_session.client("elasticache", region_name=region)
     paginator = client.get_paginator("describe_cache_clusters")
-    clusters: List[Dict] = []
+    clusters: list[dict[str, Any]] = []
     for page in paginator.paginate():
-        clusters.extend(page["CacheClusters"])
+        clusters.extend(page.get("CacheClusters", []))
     return clusters
 
 
+def transform_elasticache_clusters(
+    clusters: list[dict[str, Any]], region: str
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    cluster_data: list[dict[str, Any]] = []
+    topics: dict[str, dict[str, Any]] = {}
+
+    for cluster in clusters:
+        notification = cluster.get("NotificationConfiguration", {})
+        topic_arn = notification.get("TopicArn")
+        cluster_record = {
+            "ARN": cluster["ARN"],
+            "CacheClusterId": cluster["CacheClusterId"],
+            "CacheNodeType": cluster.get("CacheNodeType"),
+            "Engine": cluster.get("Engine"),
+            "EngineVersion": cluster.get("EngineVersion"),
+            "CacheClusterStatus": cluster.get("CacheClusterStatus"),
+            "NumCacheNodes": cluster.get("NumCacheNodes"),
+            "PreferredAvailabilityZone": cluster.get("PreferredAvailabilityZone"),
+            "PreferredMaintenanceWindow": cluster.get("PreferredMaintenanceWindow"),
+            "CacheClusterCreateTime": cluster.get("CacheClusterCreateTime"),
+            "CacheSubnetGroupName": cluster.get("CacheSubnetGroupName"),
+            "AutoMinorVersionUpgrade": cluster.get("AutoMinorVersionUpgrade"),
+            "ReplicationGroupId": cluster.get("ReplicationGroupId"),
+            "SnapshotRetentionLimit": cluster.get("SnapshotRetentionLimit"),
+            "SnapshotWindow": cluster.get("SnapshotWindow"),
+            "AuthTokenEnabled": cluster.get("AuthTokenEnabled"),
+            "TransitEncryptionEnabled": cluster.get("TransitEncryptionEnabled"),
+            "AtRestEncryptionEnabled": cluster.get("AtRestEncryptionEnabled"),
+            "TopicArn": topic_arn,
+            "Region": region,
+        }
+        cluster_data.append(cluster_record)
+
+        if topic_arn:
+            topics.setdefault(
+                topic_arn,
+                {
+                    "TopicArn": topic_arn,
+                    "TopicStatus": notification.get("TopicStatus"),
+                    "cluster_arns": [],
+                },
+            )["cluster_arns"].append(cluster["ARN"])
+
+    return cluster_data, list(topics.values())
+
+
 @timeit
 def load_elasticache_clusters(
     neo4j_session: neo4j.Session,
-    clusters: List[Dict],
+    clusters: list[dict[str, Any]],
     region: str,
     aws_account_id: str,
     update_tag: int,
 ) -> None:
-    query = """
-    UNWIND $clusters as elasticache_cluster
-        MERGE (cluster:ElasticacheCluster{id:elasticache_cluster.ARN})
-        ON CREATE SET cluster.firstseen = timestamp(),
-            cluster.arn = elasticache_cluster.ARN,
-            cluster.topic_arn = elasticache_cluster.NotificationConfiguration.TopicArn,
-            cluster.id = elasticache_cluster.CacheClusterId,
-            cluster.region = $region
-        SET cluster.lastupdated = $aws_update_tag
-
-        WITH cluster, elasticache_cluster
-        MATCH (owner:AWSAccount{id: $aws_account_id})
-        MERGE (owner)-[r3:RESOURCE]->(cluster)
-        ON CREATE SET r3.firstseen = timestamp()
-        SET r3.lastupdated = $aws_update_tag
-
-        WITH elasticache_cluster, owner
-        WHERE NOT elasticache_cluster.NotificationConfiguration IS NULL
-        MERGE (topic:ElasticacheTopic{id: elasticache_cluster.NotificationConfiguration.TopicArn})
-        ON CREATE SET topic.firstseen = timestamp(),
-            topic.arn = elasticache_cluster.NotificationConfiguration.TopicArn
-        SET topic.lastupdated = $aws_update_tag,
-            topic.status = elasticache_cluster.NotificationConfiguration.Status
-
-        MERGE (topic)-[r:CACHE_CLUSTER]->(cluster)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $aws_update_tag
-        WITH cluster, topic
-
-        MERGE (owner)-[r2:RESOURCE]->(topic)
-        ON CREATE SET r2.firstseen = timestamp()
-        SET r2.lastupdated = $aws_update_tag
-    """
     logger.info(
-        f"Loading f{len(clusters)} ElastiCache clusters for region '{region}' into graph.",
+        f"Loading {len(clusters)} ElastiCache clusters for region '{region}' into graph."
     )
-    neo4j_session.run(
-        query,
-        clusters=clusters,
-        region=region,
-        aws_update_tag=update_tag,
-        aws_account_id=aws_account_id,
+    load(
+        neo4j_session,
+        ElasticacheClusterSchema(),
+        clusters,
+        lastupdated=update_tag,
+        Region=region,
+        AWS_ID=aws_account_id,
     )
 
 
 @timeit
-def cleanup(
+def load_elasticache_topics(
     neo4j_session: neo4j.Session,
-    current_aws_account_id: str,
+    topics: list[dict[str, Any]],
+    aws_account_id: str,
     update_tag: int,
 ) -> None:
-    run_cleanup_job(
-        "aws_import_elasticache_cleanup.json",
+    if not topics:
+        return
+    logger.info(f"Loading {len(topics)} ElastiCache topics into graph.")
+    load(
         neo4j_session,
-        {"UPDATE_TAG": update_tag, "AWS_ID": current_aws_account_id},
+        ElasticacheTopicSchema(),
+        topics,
+        lastupdated=update_tag,
+        AWS_ID=aws_account_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    GraphJob.from_node_schema(ElasticacheClusterSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+    GraphJob.from_node_schema(ElasticacheTopicSchema(), common_job_parameters).run(
+        neo4j_session
     )
 
 
@@ -120,24 +134,33 @@ def cleanup(
 def sync(
     neo4j_session: neo4j.Session,
     boto3_session: boto3.session.Session,
-    regions: List[str],
+    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.info(
-            f"Syncing ElastiCache clusters for region '{region}' in account {current_aws_account_id}",
+            "Syncing ElastiCache clusters for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
         )
-        clusters = get_elasticache_clusters(boto3_session, region)
+        raw_clusters = get_elasticache_clusters(boto3_session, region)
+        cluster_data, topic_data = transform_elasticache_clusters(raw_clusters, region)
         load_elasticache_clusters(
             neo4j_session,
-            clusters,
+            cluster_data,
             region,
             current_aws_account_id,
             update_tag,
         )
-    cleanup(neo4j_session, current_aws_account_id, update_tag)
+        load_elasticache_topics(
+            neo4j_session,
+            topic_data,
+            current_aws_account_id,
+            update_tag,
+        )
+    cleanup(neo4j_session, common_job_parameters)
     merge_module_sync_metadata(
         neo4j_session,
         group_type="AWSAccount",