cartography 0.102.0rc1__py3-none-any.whl → 0.103.0rc1__py3-none-any.whl

cartography/intel/aws/rds.py

@@ -20,22 +20,28 @@ stat_handler = get_stats_client(__name__)
 
 @timeit
 @aws_handle_regions
-def get_rds_cluster_data(boto3_session: boto3.session.Session, region: str) -> List[Any]:
+def get_rds_cluster_data(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Any]:
     """
     Create an RDS boto3 client and grab all the DBClusters.
     """
-    client = boto3_session.client('rds', region_name=region)
-    paginator = client.get_paginator('describe_db_clusters')
+    client = boto3_session.client("rds", region_name=region)
+    paginator = client.get_paginator("describe_db_clusters")
     instances: List[Any] = []
     for page in paginator.paginate():
-        instances.extend(page['DBClusters'])
+        instances.extend(page["DBClusters"])
 
     return instances
 
 
 @timeit
 def load_rds_clusters(
-    neo4j_session: neo4j.Session, data: Dict, region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: Dict,
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     """
@@ -94,13 +100,31 @@ def load_rds_clusters(
         # TODO: track security groups
         # TODO: track subnet groups
 
-        cluster['EarliestRestorableTime'] = dict_value_to_str(cluster, 'EarliestRestorableTime')
-        cluster['LatestRestorableTime'] = dict_value_to_str(cluster, 'LatestRestorableTime')
-        cluster['ClusterCreateTime'] = dict_value_to_str(cluster, 'ClusterCreateTime')
-        cluster['EarliestBacktrackTime'] = dict_value_to_str(cluster, 'EarliestBacktrackTime')
-        cluster['ScalingConfigurationInfoMinCapacity'] = cluster.get('ScalingConfigurationInfo', {}).get('MinCapacity')
-        cluster['ScalingConfigurationInfoMaxCapacity'] = cluster.get('ScalingConfigurationInfo', {}).get('MaxCapacity')
-        cluster['ScalingConfigurationInfoAutoPause'] = cluster.get('ScalingConfigurationInfo', {}).get('AutoPause')
+        cluster["EarliestRestorableTime"] = dict_value_to_str(
+            cluster,
+            "EarliestRestorableTime",
+        )
+        cluster["LatestRestorableTime"] = dict_value_to_str(
+            cluster,
+            "LatestRestorableTime",
+        )
+        cluster["ClusterCreateTime"] = dict_value_to_str(cluster, "ClusterCreateTime")
+        cluster["EarliestBacktrackTime"] = dict_value_to_str(
+            cluster,
+            "EarliestBacktrackTime",
+        )
+        cluster["ScalingConfigurationInfoMinCapacity"] = cluster.get(
+            "ScalingConfigurationInfo",
+            {},
+        ).get("MinCapacity")
+        cluster["ScalingConfigurationInfoMaxCapacity"] = cluster.get(
+            "ScalingConfigurationInfo",
+            {},
+        ).get("MaxCapacity")
+        cluster["ScalingConfigurationInfoAutoPause"] = cluster.get(
+            "ScalingConfigurationInfo",
+            {},
+        ).get("AutoPause")
 
     neo4j_session.run(
         ingest_rds_cluster,
@@ -113,22 +137,28 @@ def load_rds_clusters(
 
 @timeit
 @aws_handle_regions
-def get_rds_instance_data(boto3_session: boto3.session.Session, region: str) -> List[Any]:
+def get_rds_instance_data(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Any]:
     """
     Create an RDS boto3 client and grab all the DBInstances.
     """
-    client = boto3_session.client('rds', region_name=region)
-    paginator = client.get_paginator('describe_db_instances')
+    client = boto3_session.client("rds", region_name=region)
+    paginator = client.get_paginator("describe_db_instances")
     instances: List[Any] = []
     for page in paginator.paginate():
-        instances.extend(page['DBInstances'])
+        instances.extend(page["DBInstances"])
 
     return instances
 
 
 @timeit
 def load_rds_instances(
-    neo4j_session: neo4j.Session, data: Dict, region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: Dict,
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     """
@@ -192,17 +222,17 @@ def load_rds_instances(
         if rds.get("DBClusterIdentifier"):
             clusters.append(rds)
 
-        if rds.get('VpcSecurityGroups'):
+        if rds.get("VpcSecurityGroups"):
             secgroups.append(rds)
 
-        if rds.get('DBSubnetGroup'):
+        if rds.get("DBSubnetGroup"):
             subnets.append(rds)
 
-        rds['InstanceCreateTime'] = dict_value_to_str(rds, 'InstanceCreateTime')
-        rds['LatestRestorableTime'] = dict_value_to_str(rds, 'LatestRestorableTime')
-        rds['EndpointAddress'] = ep.get('Address')
-        rds['EndpointHostedZoneId'] = ep.get('HostedZoneId')
-        rds['EndpointPort'] = ep.get('Port')
+        rds["InstanceCreateTime"] = dict_value_to_str(rds, "InstanceCreateTime")
+        rds["LatestRestorableTime"] = dict_value_to_str(rds, "LatestRestorableTime")
+        rds["EndpointAddress"] = ep.get("Address")
+        rds["EndpointHostedZoneId"] = ep.get("HostedZoneId")
+        rds["EndpointPort"] = ep.get("Port")
 
     neo4j_session.run(
         ingest_rds_instance,
@@ -212,24 +242,36 @@ def load_rds_instances(
         aws_update_tag=aws_update_tag,
     )
     _attach_ec2_security_groups(neo4j_session, secgroups, aws_update_tag)
-    _attach_ec2_subnet_groups(neo4j_session, subnets, region, current_aws_account_id, aws_update_tag)
+    _attach_ec2_subnet_groups(
+        neo4j_session,
+        subnets,
+        region,
+        current_aws_account_id,
+        aws_update_tag,
+    )
     _attach_read_replicas(neo4j_session, read_replicas, aws_update_tag)
     _attach_clusters(neo4j_session, clusters, aws_update_tag)
 
 
 @timeit
 @aws_handle_regions
-def get_rds_snapshot_data(boto3_session: boto3.session.Session, region: str) -> List[Any]:
+def get_rds_snapshot_data(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Any]:
     """
     Create an RDS boto3 client and grab all the DBSnapshots.
     """
-    client = boto3_session.client('rds', region_name=region)
-    return aws_paginate(client, 'describe_db_snapshots', 'DBSnapshots')
+    client = boto3_session.client("rds", region_name=region)
+    return aws_paginate(client, "describe_db_snapshots", "DBSnapshots")
 
 
 @timeit
 def load_rds_snapshots(
-    neo4j_session: neo4j.Session, data: Dict, region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: Dict,
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     """
@@ -293,7 +335,11 @@ def load_rds_snapshots(
 
 
 @timeit
-def _attach_snapshots(neo4j_session: neo4j.Session, snapshots: List[Dict], aws_update_tag: int) -> None:
+def _attach_snapshots(
+    neo4j_session: neo4j.Session,
+    snapshots: List[Dict],
+    aws_update_tag: int,
+) -> None:
     """
     Attach snapshots to their source instance
     """
@@ -314,7 +360,10 @@ def _attach_snapshots(neo4j_session: neo4j.Session, snapshots: List[Dict], aws_u
 
 @timeit
 def _attach_ec2_subnet_groups(
-    neo4j_session: neo4j.Session, instances: List[Dict], region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    instances: List[Dict],
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     """
@@ -337,22 +386,35 @@ def _attach_ec2_subnet_groups(
     """
     db_sngs = []
     for instance in instances:
-        db_sng = instance['DBSubnetGroup']
-        db_sng['arn'] = _get_db_subnet_group_arn(region, current_aws_account_id, db_sng['DBSubnetGroupName'])
-        db_sng['instance_arn'] = instance['DBInstanceArn']
+        db_sng = instance["DBSubnetGroup"]
+        db_sng["arn"] = _get_db_subnet_group_arn(
+            region,
+            current_aws_account_id,
+            db_sng["DBSubnetGroupName"],
+        )
+        db_sng["instance_arn"] = instance["DBInstanceArn"]
         db_sngs.append(db_sng)
     neo4j_session.run(
         attach_rds_to_subnet_group,
         SubnetGroups=db_sngs,
         aws_update_tag=aws_update_tag,
     )
-    _attach_ec2_subnets_to_subnetgroup(neo4j_session, db_sngs, region, current_aws_account_id, aws_update_tag)
+    _attach_ec2_subnets_to_subnetgroup(
+        neo4j_session,
+        db_sngs,
+        region,
+        current_aws_account_id,
+        aws_update_tag,
+    )
 
 
 @timeit
 def _attach_ec2_subnets_to_subnetgroup(
-    neo4j_session: neo4j.Session, db_subnet_groups: List[Dict], region: str,
-    current_aws_account_id: str, aws_update_tag: int,
+    neo4j_session: neo4j.Session,
+    db_subnet_groups: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    aws_update_tag: int,
 ) -> None:
     """
     Attach EC2Subnets to their DB Subnet Group.
@@ -375,15 +437,21 @@ def _attach_ec2_subnets_to_subnetgroup(
     """
     subnets = []
     for subnet_group in db_subnet_groups:
-        for subnet in subnet_group.get('Subnets', []):
-            sn_id = subnet.get('SubnetIdentifier')
-            sng_arn = _get_db_subnet_group_arn(region, current_aws_account_id, subnet_group['DBSubnetGroupName'])
-            az = subnet.get('SubnetAvailabilityZone', {}).get('Name')
-            subnets.append({
-                'sn_id': sn_id,
-                'sng_arn': sng_arn,
-                'az': az,
-            })
+        for subnet in subnet_group.get("Subnets", []):
+            sn_id = subnet.get("SubnetIdentifier")
+            sng_arn = _get_db_subnet_group_arn(
+                region,
+                current_aws_account_id,
+                subnet_group["DBSubnetGroupName"],
+            )
+            az = subnet.get("SubnetAvailabilityZone", {}).get("Name")
+            subnets.append(
+                {
+                    "sn_id": sn_id,
+                    "sng_arn": sng_arn,
+                    "az": az,
+                },
+            )
     neo4j_session.run(
         attach_subnets_to_sng,
         Subnets=subnets,
@@ -392,7 +460,11 @@ def _attach_ec2_subnets_to_subnetgroup(
 
 
 @timeit
-def _attach_ec2_security_groups(neo4j_session: neo4j.Session, instances: List[Dict], aws_update_tag: int) -> None:
+def _attach_ec2_security_groups(
+    neo4j_session: neo4j.Session,
+    instances: List[Dict],
+    aws_update_tag: int,
+) -> None:
     """
     Attach an RDS instance to its EC2SecurityGroups
     """
@@ -406,11 +478,13 @@ def _attach_ec2_security_groups(neo4j_session: neo4j.Session, instances: List[Di
     """
     groups = []
     for instance in instances:
-        for group in instance['VpcSecurityGroups']:
-            groups.append({
-                'arn': instance['DBInstanceArn'],
-                'group_id': group['VpcSecurityGroupId'],
-            })
+        for group in instance["VpcSecurityGroups"]:
+            groups.append(
+                {
+                    "arn": instance["DBInstanceArn"],
+                    "group_id": group["VpcSecurityGroupId"],
+                },
+            )
     neo4j_session.run(
         attach_rds_to_group,
         Groups=groups,
@@ -419,7 +493,11 @@ def _attach_ec2_security_groups(neo4j_session: neo4j.Session, instances: List[Di
 
 
 @timeit
-def _attach_read_replicas(neo4j_session: neo4j.Session, read_replicas: List[Dict], aws_update_tag: int) -> None:
+def _attach_read_replicas(
+    neo4j_session: neo4j.Session,
+    read_replicas: List[Dict],
+    aws_update_tag: int,
+) -> None:
     """
     Attach read replicas to their source instances
     """
@@ -439,7 +517,11 @@ def _attach_read_replicas(neo4j_session: neo4j.Session, read_replicas: List[Dict
 
 
 @timeit
-def _attach_clusters(neo4j_session: neo4j.Session, cluster_members: List[Dict], aws_update_tag: int) -> None:
+def _attach_clusters(
+    neo4j_session: neo4j.Session,
+    cluster_members: List[Dict],
+    aws_update_tag: int,
+) -> None:
     """
     Attach cluster members to their source clusters
     """
@@ -462,13 +544,20 @@ def _validate_rds_endpoint(rds: Dict) -> Dict:
     """
     Get Endpoint from RDS data structure.  Log to debug if an Endpoint field does not exist.
     """
-    ep = rds.get('Endpoint', {})
+    ep = rds.get("Endpoint", {})
     if not ep:
-        logger.debug("RDS instance does not have an Endpoint field.  Here is the object: %r", rds)
+        logger.debug(
+            "RDS instance does not have an Endpoint field.  Here is the object: %r",
+            rds,
+        )
     return ep
 
 
-def _get_db_subnet_group_arn(region: str, current_aws_account_id: str, db_subnet_group_name: str) -> str:
+def _get_db_subnet_group_arn(
+    region: str,
+    current_aws_account_id: str,
+    db_subnet_group_name: str,
+) -> str:
     """
     Return an ARN for the DB subnet group name by concatenating the account name and region.
     This is done to avoid another AWS API call since the describe_db_instances boto call does not return the DB subnet
@@ -476,7 +565,9 @@ def _get_db_subnet_group_arn(region: str, current_aws_account_id: str, db_subnet
     Form is arn:aws:rds:{region}:{account-id}:subgrp:{subnet-group-name}
     as per https://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html
     """
-    return f"arn:aws:rds:{region}:{current_aws_account_id}:subgrp:{db_subnet_group_name}"
+    return (
+        f"arn:aws:rds:{region}:{current_aws_account_id}:subgrp:{db_subnet_group_name}"
+    )
 
 
 @timeit
@@ -486,49 +577,90 @@ def transform_rds_snapshots(data: Dict) -> List[Dict]:
     for snapshot in data:
         snapshots.append(snapshot)
 
-        snapshot['SnapshotCreateTime'] = dict_value_to_str(snapshot, 'EarliestRestorableTime')
-        snapshot['InstanceCreateTime'] = dict_value_to_str(snapshot, 'InstanceCreateTime')
-        snapshot['ProcessorFeatures'] = dict_value_to_str(snapshot, 'ProcessorFeatures')
-        snapshot['OriginalSnapshotCreateTime'] = dict_value_to_str(snapshot, 'OriginalSnapshotCreateTime')
-        snapshot['SnapshotDatabaseTime'] = dict_value_to_str(snapshot, 'SnapshotDatabaseTime')
+        snapshot["SnapshotCreateTime"] = dict_value_to_str(
+            snapshot,
+            "EarliestRestorableTime",
+        )
+        snapshot["InstanceCreateTime"] = dict_value_to_str(
+            snapshot,
+            "InstanceCreateTime",
+        )
+        snapshot["ProcessorFeatures"] = dict_value_to_str(snapshot, "ProcessorFeatures")
+        snapshot["OriginalSnapshotCreateTime"] = dict_value_to_str(
+            snapshot,
+            "OriginalSnapshotCreateTime",
+        )
+        snapshot["SnapshotDatabaseTime"] = dict_value_to_str(
+            snapshot,
+            "SnapshotDatabaseTime",
+        )
 
     return snapshots
 
 
 @timeit
-def cleanup_rds_instances_and_db_subnet_groups(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_rds_instances_and_db_subnet_groups(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     """
     Remove RDS graph nodes and DBSubnetGroups that were created from other ingestion runs
     """
-    run_cleanup_job('aws_import_rds_instances_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_import_rds_instances_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
-def cleanup_rds_clusters(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_rds_clusters(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     """
     Remove RDS cluster graph nodes
     """
-    run_cleanup_job('aws_import_rds_clusters_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_import_rds_clusters_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
-def cleanup_rds_snapshots(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_rds_snapshots(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     """
     Remove RDS snapshots graph nodes
     """
-    run_cleanup_job('aws_import_rds_snapshots_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_import_rds_snapshots_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
 def sync_rds_clusters(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
-    update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: List[str],
+    current_aws_account_id: str,
+    update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     """
     Grab RDS instance data from AWS, ingest to neo4j, and run the cleanup job.
     """
     for region in regions:
-        logger.info("Syncing RDS for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing RDS for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_rds_cluster_data(boto3_session, region)
         load_rds_clusters(neo4j_session, data, region, current_aws_account_id, update_tag)  # type: ignore
     cleanup_rds_clusters(neo4j_session, common_job_parameters)
@@ -536,14 +668,22 @@ def sync_rds_clusters(
 
 @timeit
 def sync_rds_instances(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
-    update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: List[str],
+    current_aws_account_id: str,
+    update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     """
     Grab RDS instance data from AWS, ingest to neo4j, and run the cleanup job.
     """
     for region in regions:
-        logger.info("Syncing RDS for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing RDS for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_rds_instance_data(boto3_session, region)
         load_rds_instances(neo4j_session, data, region, current_aws_account_id, update_tag)  # type: ignore
     cleanup_rds_instances_and_db_subnet_groups(neo4j_session, common_job_parameters)
@@ -551,14 +691,22 @@ def sync_rds_instances(
 
 @timeit
 def sync_rds_snapshots(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
-    update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: List[str],
+    current_aws_account_id: str,
+    update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     """
     Grab RDS snapshot data from AWS, ingest to neo4j, and run the cleanup job.
     """
     for region in regions:
-        logger.info("Syncing RDS for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing RDS for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_rds_snapshot_data(boto3_session, region)
         load_rds_snapshots(neo4j_session, data, region, current_aws_account_id, update_tag)  # type: ignore
     cleanup_rds_snapshots(neo4j_session, common_job_parameters)
@@ -566,26 +714,42 @@ def sync_rds_snapshots(
 
 @timeit
 def sync(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
-    update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: List[str],
+    current_aws_account_id: str,
+    update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     sync_rds_clusters(
-        neo4j_session, boto3_session, regions, current_aws_account_id, update_tag,
+        neo4j_session,
+        boto3_session,
+        regions,
+        current_aws_account_id,
+        update_tag,
         common_job_parameters,
     )
     sync_rds_instances(
-        neo4j_session, boto3_session, regions, current_aws_account_id, update_tag,
+        neo4j_session,
+        boto3_session,
+        regions,
+        current_aws_account_id,
+        update_tag,
         common_job_parameters,
     )
     sync_rds_snapshots(
-        neo4j_session, boto3_session, regions, current_aws_account_id, update_tag,
+        neo4j_session,
+        boto3_session,
+        regions,
+        current_aws_account_id,
+        update_tag,
         common_job_parameters,
     )
     merge_module_sync_metadata(
         neo4j_session,
-        group_type='AWSAccount',
+        group_type="AWSAccount",
         group_id=current_aws_account_id,
-        synced_type='RDSCluster',
+        synced_type="RDSCluster",
         update_tag=update_tag,
         stat_handler=stat_handler,
     )