cartography 0.102.0rc2__py3-none-any.whl → 0.103.0__py3-none-any.whl

cartography/intel/aws/config.py

@@ -14,22 +14,28 @@ logger = logging.getLogger(__name__)
 
 @timeit
 @aws_handle_regions
-def get_configuration_recorders(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('config', region_name=region)
+def get_configuration_recorders(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client("config", region_name=region)
     recorders: List[Dict] = []
     response = client.describe_configuration_recorders()
-    for recorder in response.get('ConfigurationRecorders'):
+    for recorder in response.get("ConfigurationRecorders"):
         recorders.append(recorder)
     return recorders
 
 
 @timeit
 @aws_handle_regions
-def get_delivery_channels(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('config', region_name=region)
+def get_delivery_channels(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client("config", region_name=region)
     channels: List[Dict] = []
     response = client.describe_delivery_channels()
-    for channel in response.get('DeliveryChannels'):
+    for channel in response.get("DeliveryChannels"):
         channels.append(channel)
     return channels
 
@@ -37,11 +43,11 @@ def get_delivery_channels(boto3_session: boto3.session.Session, region: str) ->
 @timeit
 @aws_handle_regions
 def get_config_rules(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('config', region_name=region)
-    paginator = client.get_paginator('describe_config_rules')
+    client = boto3_session.client("config", region_name=region)
+    paginator = client.get_paginator("describe_config_rules")
     rules: List[Dict] = []
     for page in paginator.paginate():
-        rules.extend(page['ConfigRules'])
+        rules.extend(page["ConfigRules"])
     return rules
 
 
@@ -72,7 +78,7 @@ def load_configuration_recorders(
     # generate a unique id using a combo of the name, account id and region, since the name
     # itself is unique per region.
     for recorder in data:
-        recorder['_id'] = f'{recorder["name"]}:{current_aws_account_id}:{region}'
+        recorder["_id"] = f'{recorder["name"]}:{current_aws_account_id}:{region}'
 
     neo4j_session.run(
         ingest_configuration_recorders,
@@ -112,7 +118,7 @@ def load_delivery_channels(
     # generate a unique id using a combo of the name, account id and region, since the name
     # itself is unique per region.
     for channel in data:
-        channel['_id'] = f'{channel["name"]}:{current_aws_account_id}:{region}'
+        channel["_id"] = f'{channel["name"]}:{current_aws_account_id}:{region}'
 
     neo4j_session.run(
         ingest_delivery_channels,
@@ -157,9 +163,9 @@ def load_config_rules(
     """
     for rule in data:
         details = []
-        if rule.get('Source', {}).get('SourceDetails'):
+        if rule.get("Source", {}).get("SourceDetails"):
             for detail in rule["Source"]["SourceDetails"]:
-                details.append(f'{detail}')
+                details.append(f"{detail}")
         rule["_source_details"] = details
     neo4j_session.run(
         ingest_config_rules,
@@ -172,20 +178,50 @@ def load_config_rules(
 
 @timeit
 def cleanup_config(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_config_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_import_config_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @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:
     for region in regions:
-        logger.info("Syncing AWS Config for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing AWS Config for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         recorders = get_configuration_recorders(boto3_session, region)
-        load_configuration_recorders(neo4j_session, recorders, region, current_aws_account_id, update_tag)
+        load_configuration_recorders(
+            neo4j_session,
+            recorders,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
         channels = get_delivery_channels(boto3_session, region)
-        load_delivery_channels(neo4j_session, channels, region, current_aws_account_id, update_tag)
+        load_delivery_channels(
+            neo4j_session,
+            channels,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
         rules = get_config_rules(boto3_session, region)
-        load_config_rules(neo4j_session, rules, region, current_aws_account_id, update_tag)
+        load_config_rules(
+            neo4j_session,
+            rules,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_config(neo4j_session, common_job_parameters)

cartography/intel/aws/dynamodb.py

@@ -22,12 +22,19 @@ stat_handler = get_stats_client(__name__)
 
 @timeit
 @aws_handle_regions
-def get_dynamodb_tables(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('dynamodb', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('list_tables')
+def get_dynamodb_tables(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "dynamodb",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("list_tables")
     dynamodb_tables = []
     for page in paginator.paginate():
-        for table_name in page['TableNames']:
+        for table_name in page["TableNames"]:
             dynamodb_tables.append(client.describe_table(TableName=table_name))
     return dynamodb_tables
 
@@ -38,33 +45,50 @@ def transform_dynamodb_tables(dynamodb_tables: List, region: str) -> Any:
     ddb_gsi_data: List[Dict[str, Any]] = []
 
     for table in dynamodb_tables:
-        ddb_table_data.append({
-            'Arn': table['Table']['TableArn'],
-            'TableName': table['Table']['TableName'],
-            'Region': region,
-            'Rows': table['Table']['ItemCount'],
-            'Size': table['Table']['TableSizeBytes'],
-            'ProvisionedThroughputReadCapacityUnits': table['Table']['ProvisionedThroughput']['ReadCapacityUnits'],
-            'ProvisionedThroughputWriteCapacityUnits': table['Table']['ProvisionedThroughput']['WriteCapacityUnits'],
-        })
-        for gsi in table['Table'].get('GlobalSecondaryIndexes', []):
-            ddb_gsi_data.append({
-                'Arn': gsi['IndexArn'],
-                'TableArn': table['Table']['TableArn'],
-                'Region': region,
-                'ProvisionedThroughputReadCapacityUnits': gsi['ProvisionedThroughput']['ReadCapacityUnits'],
-                'ProvisionedThroughputWriteCapacityUnits': gsi['ProvisionedThroughput']['WriteCapacityUnits'],
-                'GSIName': gsi['IndexName'],
-            })
+        ddb_table_data.append(
+            {
+                "Arn": table["Table"]["TableArn"],
+                "TableName": table["Table"]["TableName"],
+                "Region": region,
+                "Rows": table["Table"]["ItemCount"],
+                "Size": table["Table"]["TableSizeBytes"],
+                "ProvisionedThroughputReadCapacityUnits": table["Table"][
+                    "ProvisionedThroughput"
+                ]["ReadCapacityUnits"],
+                "ProvisionedThroughputWriteCapacityUnits": table["Table"][
+                    "ProvisionedThroughput"
+                ]["WriteCapacityUnits"],
+            },
+        )
+        for gsi in table["Table"].get("GlobalSecondaryIndexes", []):
+            ddb_gsi_data.append(
+                {
+                    "Arn": gsi["IndexArn"],
+                    "TableArn": table["Table"]["TableArn"],
+                    "Region": region,
+                    "ProvisionedThroughputReadCapacityUnits": gsi[
+                        "ProvisionedThroughput"
+                    ]["ReadCapacityUnits"],
+                    "ProvisionedThroughputWriteCapacityUnits": gsi[
+                        "ProvisionedThroughput"
+                    ]["WriteCapacityUnits"],
+                    "GSIName": gsi["IndexName"],
+                },
+            )
     return ddb_table_data, ddb_gsi_data
 
 
 @timeit
 def load_dynamodb_tables(
-    neo4j_session: neo4j.Session, tables_data: List[Dict[str, Any]], region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    tables_data: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
-    logger.info(f"Loading Dynamo DB tables {len(tables_data)} for region '{region}' into graph.")
+    logger.info(
+        f"Loading Dynamo DB tables {len(tables_data)} for region '{region}' into graph.",
+    )
     load(
         neo4j_session,
         DynamoDBTableSchema(),
@@ -77,10 +101,15 @@ def load_dynamodb_tables(
 
 @timeit
 def load_dynamodb_gsi(
-    neo4j_session: neo4j.Session, gsi_data: List[Dict[str, Any]], region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    gsi_data: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
-    logger.info(f"Loading Dynamo DB GSI {len(gsi_data)} for region '{region}' into graph.")
+    logger.info(
+        f"Loading Dynamo DB GSI {len(gsi_data)} for region '{region}' into graph.",
+    )
     load(
         neo4j_session,
         DynamoDBGSISchema(),
@@ -92,38 +121,77 @@ def load_dynamodb_gsi(
 
 
 @timeit
-def cleanup_dynamodb_tables(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    GraphJob.from_node_schema(DynamoDBTableSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(DynamoDBGSISchema(), common_job_parameters).run(neo4j_session)
+def cleanup_dynamodb_tables(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
+    GraphJob.from_node_schema(DynamoDBTableSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(DynamoDBGSISchema(), common_job_parameters).run(
+        neo4j_session,
+    )
 
 
 @timeit
 def sync_dynamodb_tables(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
-    aws_update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: List[str],
+    current_aws_account_id: str,
+    aws_update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     for region in regions:
-        logger.info("Syncing DynamoDB for region in '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing DynamoDB for region in '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         dynamodb_tables = get_dynamodb_tables(boto3_session, region)
-        ddb_table_data, ddb_gsi_data = transform_dynamodb_tables(dynamodb_tables, region)
-        load_dynamodb_tables(neo4j_session, ddb_table_data, region, current_aws_account_id, aws_update_tag)
-        load_dynamodb_gsi(neo4j_session, ddb_gsi_data, region, current_aws_account_id, aws_update_tag)
+        ddb_table_data, ddb_gsi_data = transform_dynamodb_tables(
+            dynamodb_tables,
+            region,
+        )
+        load_dynamodb_tables(
+            neo4j_session,
+            ddb_table_data,
+            region,
+            current_aws_account_id,
+            aws_update_tag,
+        )
+        load_dynamodb_gsi(
+            neo4j_session,
+            ddb_gsi_data,
+            region,
+            current_aws_account_id,
+            aws_update_tag,
+        )
     cleanup_dynamodb_tables(neo4j_session, common_job_parameters)
 
 
 @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_dynamodb_tables(
-        neo4j_session, boto3_session, regions, current_aws_account_id, update_tag, common_job_parameters,
+        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='DynamoDBTable',
+        synced_type="DynamoDBTable",
         update_tag=update_tag,
         stat_handler=stat_handler,
     )

cartography/intel/aws/ec2/__init__.py

@@ -10,6 +10,6 @@ logger = logging.getLogger(__name__)
 
 @timeit
 def get_ec2_regions(boto3_session: boto3.session.Session) -> List[str]:
-    client = boto3_session.client('ec2')
+    client = boto3_session.client("ec2")
     result = client.describe_regions()
-    return [r['RegionName'] for r in result['Regions']]
+    return [r["RegionName"] for r in result["Regions"]]