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

cartography/intel/aws/kms.py

@@ -23,18 +23,23 @@ logger = logging.getLogger(__name__)
 @timeit
 @aws_handle_regions
 def get_kms_key_list(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('kms', region_name=region)
-    paginator = client.get_paginator('list_keys')
+    client = boto3_session.client("kms", region_name=region)
+    paginator = client.get_paginator("list_keys")
     key_list: List[Any] = []
     for page in paginator.paginate():
-        key_list.extend(page['Keys'])
+        key_list.extend(page["Keys"])
 
     described_key_list = []
     for key in key_list:
         try:
-            response = client.describe_key(KeyId=key["KeyId"])['KeyMetadata']
+            response = client.describe_key(KeyId=key["KeyId"])["KeyMetadata"]
         except ClientError as e:
-            logger.warning("Failed to describe key with key id - {}. Error - {}".format(key["KeyId"], e))
+            logger.warning(
+                "Failed to describe key with key id - {}. Error - {}".format(
+                    key["KeyId"],
+                    e,
+                ),
+            )
             continue
 
         described_key_list.append(response)
@@ -45,17 +50,19 @@ def get_kms_key_list(boto3_session: boto3.session.Session, region: str) -> List[
 @timeit
 @aws_handle_regions
 def get_kms_key_details(
-    boto3_session: boto3.session.Session, kms_key_data: Dict, region: str,
+    boto3_session: boto3.session.Session,
+    kms_key_data: Dict,
+    region: str,
 ) -> Generator[Any, Any, Any]:
     """
     Iterates over all KMS Keys.
     """
-    client = boto3_session.client('kms', region_name=region)
+    client = boto3_session.client("kms", region_name=region)
     for key in kms_key_data:
         policy = get_policy(key, client)
         aliases = get_aliases(key, client)
         grants = get_grants(key, client)
-        yield key['KeyId'], policy, aliases, grants
+        yield key["KeyId"], policy, aliases, grants
 
 
 @timeit
@@ -64,10 +71,10 @@ def get_policy(key: Dict, client: botocore.client.BaseClient) -> Any:
     Gets the KMS Key policy. Returns policy string or None if we are unable to retrieve it.
     """
     try:
-        policy = client.get_key_policy(KeyId=key["KeyId"], PolicyName='default')
+        policy = client.get_key_policy(KeyId=key["KeyId"], PolicyName="default")
     except ClientError as e:
         policy = None
-        if e.response['Error']['Code'] == 'AccessDeniedException':
+        if e.response["Error"]["Code"] == "AccessDeniedException":
             logger.warning(
                 f"kms:get_key_policy on key id {key['KeyId']} failed with AccessDeniedException; continuing sync.",
                 exc_info=True,
@@ -84,9 +91,9 @@ def get_aliases(key: Dict, client: botocore.client.BaseClient) -> List[Any]:
     Gets the KMS Key Aliases.
     """
     aliases: List[Any] = []
-    paginator = client.get_paginator('list_aliases')
-    for page in paginator.paginate(KeyId=key['KeyId']):
-        aliases.extend(page['Aliases'])
+    paginator = client.get_paginator("list_aliases")
+    for page in paginator.paginate(KeyId=key["KeyId"]):
+        aliases.extend(page["Aliases"])
 
     return aliases
 
@@ -97,12 +104,12 @@ def get_grants(key: Dict, client: botocore.client.BaseClient) -> List[Any]:
     Gets the KMS Key Grants.
     """
     grants: List[Any] = []
-    paginator = client.get_paginator('list_grants')
+    paginator = client.get_paginator("list_grants")
     try:
-        for page in paginator.paginate(KeyId=key['KeyId']):
-            grants.extend(page['Grants'])
+        for page in paginator.paginate(KeyId=key["KeyId"]):
+            grants.extend(page["Grants"])
     except ClientError as e:
-        if e.response['Error']['Code'] == 'AccessDeniedException':
+        if e.response["Error"]["Code"] == "AccessDeniedException":
             logger.warning(
                 f'kms:list_grants on key_id {key["KeyId"]} failed with AccessDeniedException; continuing sync.',
                 exc_info=True,
@@ -113,7 +120,11 @@ def get_grants(key: Dict, client: botocore.client.BaseClient) -> List[Any]:
 
 
 @timeit
-def _load_kms_key_aliases(neo4j_session: neo4j.Session, aliases: List[Dict], update_tag: int) -> None:
+def _load_kms_key_aliases(
+    neo4j_session: neo4j.Session,
+    aliases: List[Dict],
+    update_tag: int,
+) -> None:
     """
     Ingest KMS Aliases into neo4j.
     """
@@ -137,7 +148,11 @@ def _load_kms_key_aliases(neo4j_session: neo4j.Session, aliases: List[Dict], upd
 
 
 @timeit
-def _load_kms_key_grants(neo4j_session: neo4j.Session, grants_list: List[Dict], update_tag: int) -> None:
+def _load_kms_key_grants(
+    neo4j_session: neo4j.Session,
+    grants_list: List[Dict],
+    update_tag: int,
+) -> None:
     """
     Ingest KMS Key Grants into neo4j.
     """
@@ -157,7 +172,7 @@ def _load_kms_key_grants(neo4j_session: neo4j.Session, grants_list: List[Dict],
     # neo4j does not accept datetime objects and values. This loop is used to convert
     # these values to string.
     for grant in grants_list:
-        grant['CreationDate'] = str(grant['CreationDate'])
+        grant["CreationDate"] = str(grant["CreationDate"])
 
     neo4j_session.run(
         ingest_grants,
@@ -167,7 +182,11 @@ def _load_kms_key_grants(neo4j_session: neo4j.Session, grants_list: List[Dict],
 
 
 @timeit
-def _load_kms_key_policies(neo4j_session: neo4j.Session, policies: List[Dict], update_tag: int) -> None:
+def _load_kms_key_policies(
+    neo4j_session: neo4j.Session,
+    policies: List[Dict],
+    update_tag: int,
+) -> None:
     """
     Ingest KMS Key policy results into neo4j.
     """
@@ -201,8 +220,11 @@ def _set_default_values(neo4j_session: neo4j.Session, aws_account_id: str) -> No
 
 @timeit
 def load_kms_key_details(
-        neo4j_session: neo4j.Session, policy_alias_grants_data: List[Tuple[Any, Any, Any, Any]], region: str,
-        aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    policy_alias_grants_data: List[Tuple[Any, Any, Any, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     """
     Create dictionaries for all KMS key policies, aliases and grants so we can import them in a single query for each
@@ -221,9 +243,9 @@ def load_kms_key_details(
 
     # cleanup existing policy properties
     run_cleanup_job(
-        'aws_kms_details.json',
+        "aws_kms_details.json",
         neo4j_session,
-        {'UPDATE_TAG': update_tag, 'AWS_ID': aws_account_id},
+        {"UPDATE_TAG": update_tag, "AWS_ID": aws_account_id},
     )
 
     _load_kms_key_policies(neo4j_session, policies, update_tag)
@@ -281,7 +303,7 @@ def parse_policy(key: str, policy: Policy) -> Optional[Dict[Any, Any]]:
     # }
     if policy is not None:
         # get just the policy element and convert to JSON because boto3 returns this as string
-        policy = Policy(json.loads(policy['Policy']))
+        policy = Policy(json.loads(policy["Policy"]))
         if policy.is_internet_accessible():
             return {
                 "kms_key": key,
@@ -296,7 +318,10 @@ def parse_policy(key: str, policy: Policy) -> Optional[Dict[Any, Any]]:
 
 @timeit
 def load_kms_keys(
-    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:
     ingest_keys = """
@@ -322,9 +347,9 @@ def load_kms_keys(
     # neo4j does not accept datetime objects and values. This loop is used to convert
     # these values to string.
     for key in data:
-        key['CreationDate'] = str(key['CreationDate'])
-        key['DeletionDate'] = str(key.get('DeletionDate'))
-        key['ValidTo'] = str(key.get('ValidTo'))
+        key["CreationDate"] = str(key["CreationDate"])
+        key["DeletionDate"] = str(key.get("DeletionDate"))
+        key["ValidTo"] = str(key.get("ValidTo"))
 
     neo4j_session.run(
         ingest_keys,
@@ -337,29 +362,58 @@ def load_kms_keys(
 
 @timeit
 def cleanup_kms(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_kms_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job("aws_import_kms_cleanup.json", neo4j_session, common_job_parameters)
 
 
 @timeit
 def sync_kms_keys(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     kms_keys = get_kms_key_list(boto3_session, region)
 
-    load_kms_keys(neo4j_session, kms_keys, region, current_aws_account_id, aws_update_tag)
+    load_kms_keys(
+        neo4j_session,
+        kms_keys,
+        region,
+        current_aws_account_id,
+        aws_update_tag,
+    )
 
     policy_alias_grants_data = get_kms_key_details(boto3_session, kms_keys, region)
-    load_kms_key_details(neo4j_session, policy_alias_grants_data, region, current_aws_account_id, aws_update_tag)
+    load_kms_key_details(
+        neo4j_session,
+        policy_alias_grants_data,
+        region,
+        current_aws_account_id,
+        aws_update_tag,
+    )
 
 
 @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 KMS for region %s in account '%s'.", region, current_aws_account_id)
-        sync_kms_keys(neo4j_session, boto3_session, region, current_aws_account_id, update_tag)
+        logger.info(
+            "Syncing KMS for region %s in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
+        sync_kms_keys(
+            neo4j_session,
+            boto3_session,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
 
     cleanup_kms(neo4j_session, common_job_parameters)

cartography/intel/aws/lambda_function.py

@@ -21,18 +21,22 @@ def get_lambda_data(boto3_session: boto3.session.Session, region: str) -> List[D
     """
     Create an Lambda boto3 client and grab all the lambda functions.
     """
-    client = boto3_session.client('lambda', region_name=region)
-    paginator = client.get_paginator('list_functions')
+    client = boto3_session.client("lambda", region_name=region)
+    paginator = client.get_paginator("list_functions")
     lambda_functions = []
     for page in paginator.paginate():
-        for each_function in page['Functions']:
+        for each_function in page["Functions"]:
             lambda_functions.append(each_function)
     return lambda_functions
 
 
 @timeit
 def load_lambda_functions(
-        neo4j_session: neo4j.Session, data: List[Dict], region: str, current_aws_account_id: str, aws_update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    aws_update_tag: int,
 ) -> None:
     ingest_lambda_functions = """
     UNWIND $lambda_functions_list AS lf
@@ -86,22 +90,28 @@ def load_lambda_functions(
 
 @timeit
 @aws_handle_regions
-def get_function_aliases(lambda_function: Dict, client: botocore.client.BaseClient) -> List[Any]:
+def get_function_aliases(
+    lambda_function: Dict,
+    client: botocore.client.BaseClient,
+) -> List[Any]:
     aliases: List[Any] = []
-    paginator = client.get_paginator('list_aliases')
-    for page in paginator.paginate(FunctionName=lambda_function['FunctionName']):
-        aliases.extend(page['Aliases'])
+    paginator = client.get_paginator("list_aliases")
+    for page in paginator.paginate(FunctionName=lambda_function["FunctionName"]):
+        aliases.extend(page["Aliases"])
 
     return aliases
 
 
 @timeit
 @aws_handle_regions
-def get_event_source_mappings(lambda_function: Dict, client: botocore.client.BaseClient) -> List[Any]:
+def get_event_source_mappings(
+    lambda_function: Dict,
+    client: botocore.client.BaseClient,
+) -> List[Any]:
     event_source_mappings: List[Any] = []
-    paginator = client.get_paginator('list_event_source_mappings')
-    for page in paginator.paginate(FunctionName=lambda_function['FunctionName']):
-        event_source_mappings.extend(page['EventSourceMappings'])
+    paginator = client.get_paginator("list_event_source_mappings")
+    for page in paginator.paginate(FunctionName=lambda_function["FunctionName"]):
+        event_source_mappings.extend(page["EventSourceMappings"])
 
     return event_source_mappings
 
@@ -109,22 +119,32 @@ def get_event_source_mappings(lambda_function: Dict, client: botocore.client.Bas
 @timeit
 @aws_handle_regions
 def get_lambda_function_details(
-        boto3_session: boto3.session.Session, data: List[Dict], region: str,
+    boto3_session: boto3.session.Session,
+    data: List[Dict],
+    region: str,
 ) -> List[Tuple[str, List[Any], List[Any], List[Any]]]:
-    client = boto3_session.client('lambda', region_name=region)
+    client = boto3_session.client("lambda", region_name=region)
     details = []
     for lambda_function in data:
         function_aliases = get_function_aliases(lambda_function, client)
         event_source_mappings = get_event_source_mappings(lambda_function, client)
-        layers = lambda_function.get('Layers', [])
-        details.append((lambda_function['FunctionArn'], function_aliases, event_source_mappings, layers))
+        layers = lambda_function.get("Layers", [])
+        details.append(
+            (
+                lambda_function["FunctionArn"],
+                function_aliases,
+                event_source_mappings,
+                layers,
+            ),
+        )
     return details
 
 
 @timeit
 def load_lambda_function_details(
-        neo4j_session: neo4j.Session, lambda_function_details: List[Tuple[str, List[Dict], List[Dict], List[Dict]]],
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    lambda_function_details: List[Tuple[str, List[Dict], List[Dict], List[Dict]]],
+    update_tag: int,
 ) -> None:
     lambda_aliases: List[Dict] = []
     lambda_event_source_mappings: List[Dict] = []
@@ -132,22 +152,30 @@ def load_lambda_function_details(
     for function_arn, aliases, event_source_mappings, layers in lambda_function_details:
         if len(aliases) > 0:
             for alias in aliases:
-                alias['FunctionArn'] = function_arn
+                alias["FunctionArn"] = function_arn
             lambda_aliases.extend(aliases)
         if len(event_source_mappings) > 0:
             lambda_event_source_mappings.extend(event_source_mappings)
         if len(layers) > 0:
             for layer in layers:
-                layer['FunctionArn'] = function_arn
+                layer["FunctionArn"] = function_arn
             lambda_layers.extend(layers)
 
     _load_lambda_function_aliases(neo4j_session, lambda_aliases, update_tag)
-    _load_lambda_event_source_mappings(neo4j_session, lambda_event_source_mappings, update_tag)
+    _load_lambda_event_source_mappings(
+        neo4j_session,
+        lambda_event_source_mappings,
+        update_tag,
+    )
     _load_lambda_layers(neo4j_session, lambda_layers, update_tag)
 
 
 @timeit
-def _load_lambda_function_aliases(neo4j_session: neo4j.Session, lambda_aliases: List[Dict], update_tag: int) -> None:
+def _load_lambda_function_aliases(
+    neo4j_session: neo4j.Session,
+    lambda_aliases: List[Dict],
+    update_tag: int,
+) -> None:
     ingest_aliases = """
     UNWIND $aliases_list AS alias
         MERGE (a:AWSLambdaFunctionAlias{id: alias.AliasArn})
@@ -173,7 +201,9 @@ def _load_lambda_function_aliases(neo4j_session: neo4j.Session, lambda_aliases:
 
 @timeit
 def _load_lambda_event_source_mappings(
-        neo4j_session: neo4j.Session, lambda_event_source_mappings: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    lambda_event_source_mappings: List[Dict],
+    update_tag: int,
 ) -> None:
     ingest_esms = """
     UNWIND $esm_list AS esm
@@ -208,7 +238,11 @@ def _load_lambda_event_source_mappings(
 
 
 @timeit
-def _load_lambda_layers(neo4j_session: neo4j.Session, lambda_layers: List[Dict], update_tag: int) -> None:
+def _load_lambda_layers(
+    neo4j_session: neo4j.Session,
+    lambda_layers: List[Dict],
+    update_tag: int,
+) -> None:
     ingest_layers = """
     UNWIND $layers_list AS layer
         MERGE (l:AWSLambdaLayer{id: layer.Arn})
@@ -233,29 +267,64 @@ def _load_lambda_layers(neo4j_session: neo4j.Session, lambda_layers: List[Dict],
 
 @timeit
 def cleanup_lambda(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_lambda_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_import_lambda_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
 def sync_lambda_functions(
-        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 Lambda for region in '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing Lambda for region in '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_lambda_data(boto3_session, region)
-        load_lambda_functions(neo4j_session, data, region, current_aws_account_id, aws_update_tag)
-        lambda_function_details = get_lambda_function_details(boto3_session, data, region)
-        load_lambda_function_details(neo4j_session, lambda_function_details, aws_update_tag)
+        load_lambda_functions(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            aws_update_tag,
+        )
+        lambda_function_details = get_lambda_function_details(
+            boto3_session,
+            data,
+            region,
+        )
+        load_lambda_function_details(
+            neo4j_session,
+            lambda_function_details,
+            aws_update_tag,
+        )
 
     cleanup_lambda(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_lambda_functions(
-        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,
     )

cartography/intel/aws/organizations.py

@@ -16,19 +16,22 @@ def get_account_from_arn(arn: str) -> str:
 
 
 def get_caller_identity(boto3_session: boto3.session.Session) -> Dict:
-    client = boto3_session.client('sts')
+    client = boto3_session.client("sts")
     return client.get_caller_identity()
 
 
 def get_current_aws_account_id(boto3_session: boto3.session.Session) -> Dict:
-    return get_caller_identity(boto3_session)['Account']
+    return get_caller_identity(boto3_session)["Account"]
 
 
 def get_aws_account_default(boto3_session: boto3.session.Session) -> Dict:
     try:
         return {boto3_session.profile_name: get_current_aws_account_id(boto3_session)}
     except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError) as e:
-        logger.debug("Error occurred getting default AWS account number.", exc_info=True)
+        logger.debug(
+            "Error occurred getting default AWS account number.",
+            exc_info=True,
+        )
         logger.error(
             (
                 "Unable to get AWS account number, an error occurred: '%s'. Make sure your AWS credentials are "
@@ -43,13 +46,20 @@ def get_aws_account_default(boto3_session: boto3.session.Session) -> Dict:
 def get_aws_accounts_from_botocore_config(boto3_session: boto3.session.Session) -> Dict:
     d = {}
     for profile_name in boto3_session.available_profiles:
-        if profile_name == 'default':
+        if profile_name == "default":
             logger.debug("Skipping AWS profile 'default'.")
             continue
         try:
             profile_boto3_session = boto3.Session(profile_name=profile_name)
-        except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError) as e:
-            logger.debug("Error occurred calling boto3.Session() with profile_name '%s'.", profile_name, exc_info=True)
+        except (
+            botocore.exceptions.BotoCoreError,
+            botocore.exceptions.ClientError,
+        ) as e:
+            logger.debug(
+                "Error occurred calling boto3.Session() with profile_name '%s'.",
+                profile_name,
+                exc_info=True,
+            )
             logger.error(
                 (
                     "Unable to initialize an AWS session using profile '%s', an error occurred: '%s'. Make sure your "
@@ -62,7 +72,10 @@ def get_aws_accounts_from_botocore_config(boto3_session: boto3.session.Session)
             continue
         try:
             d[profile_name] = get_current_aws_account_id(profile_boto3_session)
-        except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError) as e:
+        except (
+            botocore.exceptions.BotoCoreError,
+            botocore.exceptions.ClientError,
+        ) as e:
             logger.debug(
                 "Error occurred getting AWS account number with profile_name '%s'.",
                 profile_name,
@@ -87,7 +100,9 @@ def get_aws_accounts_from_botocore_config(boto3_session: boto3.session.Session)
 
 
 def load_aws_accounts(
-    neo4j_session: neo4j.Session, aws_accounts: Dict, aws_update_tag: int,
+    neo4j_session: neo4j.Session,
+    aws_accounts: Dict,
+    aws_update_tag: int,
     common_job_parameters: Dict,
 ) -> None:
     query = """
@@ -105,7 +120,7 @@ def load_aws_accounts(
     SET r.lastupdated = $aws_update_tag;
     """
     for account_name, account_id in aws_accounts.items():
-        root_arn = f'arn:aws:iam::{account_id}:root'
+        root_arn = f"arn:aws:iam::{account_id}:root"
         neo4j_session.run(
             query,
             ACCOUNT_ID=account_id,
@@ -116,5 +131,10 @@ def load_aws_accounts(
 
 
 @timeit
-def sync(neo4j_session: neo4j.Session, accounts: Dict, update_tag: int, common_job_parameters: Dict) -> None:
+def sync(
+    neo4j_session: neo4j.Session,
+    accounts: Dict,
+    update_tag: int,
+    common_job_parameters: Dict,
+) -> None:
     load_aws_accounts(neo4j_session, accounts, update_tag, common_job_parameters)