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

cartography/intel/aws/ec2/vpc.py

@@ -6,23 +6,29 @@ from typing import List
 import boto3
 import neo4j
 
-from .util import get_botocore_config
 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
+
 logger = logging.getLogger(__name__)
 
 
 @timeit
 @aws_handle_regions
 def get_ec2_vpcs(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    return client.describe_vpcs()['Vpcs']
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    return client.describe_vpcs()["Vpcs"]
 
 
 def _get_cidr_association_statement(block_type: str) -> str:
-    INGEST_CIDR_TEMPLATE = Template("""
+    INGEST_CIDR_TEMPLATE = Template(
+        """
     MATCH (vpc:AWSVpc{id: $VpcId})
     WITH vpc
     UNWIND $CidrBlock as block_data
@@ -36,7 +42,8 @@ def _get_cidr_association_statement(block_type: str) -> str:
         WITH vpc, new_block
         MERGE (vpc)-[r:BLOCK_ASSOCIATION]->(new_block)
         ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag""")
+        SET r.lastupdated = $update_tag""",
+    )
 
     BLOCK_CIDR = "CidrBlock"
     STATE_NAME = "CidrBlockState"
@@ -53,12 +60,19 @@ def _get_cidr_association_statement(block_type: str) -> str:
     else:
         raise ValueError(f"Unsupported block type specified - {block_type}")
 
-    return INGEST_CIDR_TEMPLATE.safe_substitute(block_label=BLOCK_TYPE, block_cidr=BLOCK_CIDR, state_name=STATE_NAME)
+    return INGEST_CIDR_TEMPLATE.safe_substitute(
+        block_label=BLOCK_TYPE,
+        block_cidr=BLOCK_CIDR,
+        state_name=STATE_NAME,
+    )
 
 
 @timeit
 def load_cidr_association_set(
-    neo4j_session: neo4j.Session, vpc_id: str, vpc_data: Dict, block_type: str,
+    neo4j_session: neo4j.Session,
+    vpc_id: str,
+    vpc_data: Dict,
+    block_type: str,
     update_tag: int,
 ) -> None:
     ingest_statement = _get_cidr_association_statement(block_type)
@@ -78,7 +92,10 @@ def load_cidr_association_set(
 
 @timeit
 def load_ec2_vpcs(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    current_aws_account_id: str,
     update_tag: int,
 ) -> None:
     # https://docs.aws.amazon.com/cli/latest/reference/ec2/describe-vpcs.html
@@ -161,16 +178,24 @@ def load_ec2_vpcs(
 
 @timeit
 def cleanup_ec2_vpcs(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_vpc_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job("aws_import_vpc_cleanup.json", neo4j_session, common_job_parameters)
 
 
 @timeit
 def sync_vpc(
-    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 EC2 VPC for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 VPC for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_ec2_vpcs(boto3_session, region)
         load_ec2_vpcs(neo4j_session, data, region, current_aws_account_id, update_tag)
     cleanup_ec2_vpcs(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/vpc_peerings.py

@@ -5,25 +5,36 @@ from typing import List
 import boto3
 import neo4j
 
-from .util import get_botocore_config
 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
+
 logger = logging.getLogger(__name__)
 
 
 @timeit
 @aws_handle_regions
-def get_vpc_peerings_data(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    return client.describe_vpc_peering_connections()['VpcPeeringConnections']
+def get_vpc_peerings_data(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    return client.describe_vpc_peering_connections()["VpcPeeringConnections"]
 
 
 @timeit
 def load_vpc_peerings(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str,
-    aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     ingest_vpc_peerings = """
     UNWIND $vpc_peerings AS vpc_peering
@@ -77,15 +88,21 @@ def load_vpc_peerings(
     """
 
     neo4j_session.run(
-        ingest_vpc_peerings, vpc_peerings=data, update_tag=update_tag,
-        region=region, aws_account_id=aws_account_id,
+        ingest_vpc_peerings,
+        vpc_peerings=data,
+        update_tag=update_tag,
+        region=region,
+        aws_account_id=aws_account_id,
     )
 
 
 @timeit
 def load_accepter_cidrs(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str,
-    aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
 
     ingest_accepter_cidr = """
@@ -110,15 +127,21 @@ def load_accepter_cidrs(
     """
 
     neo4j_session.run(
-        ingest_accepter_cidr, vpc_peerings=data, update_tag=update_tag,
-        region=region, aws_account_id=aws_account_id,
+        ingest_accepter_cidr,
+        vpc_peerings=data,
+        update_tag=update_tag,
+        region=region,
+        aws_account_id=aws_account_id,
     )
 
 
 @timeit
 def load_requester_cidrs(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str,
-    aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
 
     ingest_requester_cidr = """
@@ -143,25 +166,61 @@ def load_requester_cidrs(
     """
 
     neo4j_session.run(
-        ingest_requester_cidr, vpc_peerings=data, update_tag=update_tag,
-        region=region, aws_account_id=aws_account_id,
+        ingest_requester_cidr,
+        vpc_peerings=data,
+        update_tag=update_tag,
+        region=region,
+        aws_account_id=aws_account_id,
     )
 
 
 @timeit
-def cleanup_vpc_peerings(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_vpc_peering_cleanup.json', neo4j_session, common_job_parameters)
+def cleanup_vpc_peerings(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
+    run_cleanup_job(
+        "aws_import_vpc_peering_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
 def sync_vpc_peerings(
-    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.debug("Syncing EC2 VPC peering for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.debug(
+            "Syncing EC2 VPC peering for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_vpc_peerings_data(boto3_session, region)
-        load_vpc_peerings(neo4j_session, data, region, current_aws_account_id, update_tag)
-        load_accepter_cidrs(neo4j_session, data, region, current_aws_account_id, update_tag)
-        load_requester_cidrs(neo4j_session, data, region, current_aws_account_id, update_tag)
+        load_vpc_peerings(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
+        load_accepter_cidrs(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
+        load_requester_cidrs(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_vpc_peerings(neo4j_session, common_job_parameters)

cartography/intel/aws/ecr.py

@@ -18,33 +18,48 @@ logger = logging.getLogger(__name__)
 
 @timeit
 @aws_handle_regions
-def get_ecr_repositories(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
+def get_ecr_repositories(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
     logger.info("Getting ECR repositories for region '%s'.", region)
-    client = boto3_session.client('ecr', region_name=region)
-    paginator = client.get_paginator('describe_repositories')
+    client = boto3_session.client("ecr", region_name=region)
+    paginator = client.get_paginator("describe_repositories")
     ecr_repositories: List[Dict] = []
     for page in paginator.paginate():
-        ecr_repositories.extend(page['repositories'])
+        ecr_repositories.extend(page["repositories"])
     return ecr_repositories
 
 
 @timeit
 @aws_handle_regions
-def get_ecr_repository_images(boto3_session: boto3.session.Session, region: str, repository_name: str) -> List[Dict]:
-    logger.debug("Getting ECR images in repository '%s' for region '%s'.", repository_name, region)
-    client = boto3_session.client('ecr', region_name=region)
-    list_paginator = client.get_paginator('list_images')
+def get_ecr_repository_images(
+    boto3_session: boto3.session.Session, region: str, repository_name: str
+) -> List[Dict]:
+    logger.debug(
+        "Getting ECR images in repository '%s' for region '%s'.",
+        repository_name,
+        region,
+    )
+    client = boto3_session.client("ecr", region_name=region)
+    list_paginator = client.get_paginator("list_images")
     ecr_repository_images: List[Dict] = []
     for page in list_paginator.paginate(repositoryName=repository_name):
-        image_ids = page['imageIds']
+        image_ids = page["imageIds"]
         if not image_ids:
             continue
-        describe_paginator = client.get_paginator('describe_images')
-        describe_response = describe_paginator.paginate(repositoryName=repository_name, imageIds=image_ids)
+        describe_paginator = client.get_paginator("describe_images")
+        describe_response = describe_paginator.paginate(
+            repositoryName=repository_name, imageIds=image_ids
+        )
         for response in describe_response:
-            image_details = response['imageDetails']
+            image_details = response["imageDetails"]
             image_details = [
-                {**detail, 'imageTag': detail['imageTags'][0]} if detail.get('imageTags') else detail
+                (
+                    {**detail, "imageTag": detail["imageTags"][0]}
+                    if detail.get("imageTags")
+                    else detail
+                )
                 for detail in image_details
             ]
             ecr_repository_images.extend(image_details)
@@ -53,7 +68,10 @@ def get_ecr_repository_images(boto3_session: boto3.session.Session, region: str,
 
 @timeit
 def load_ecr_repositories(
-    neo4j_session: neo4j.Session, repos: List[Dict], region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    repos: List[Dict],
+    region: str,
+    current_aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
     query = """
@@ -73,7 +91,9 @@ def load_ecr_repositories(
         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.")
+    logger.info(
+        f"Loading {len(repos)} ECR repositories for region {region} into graph.",
+    )
     neo4j_session.run(
         query,
         Repositories=repos,
@@ -91,21 +111,23 @@ def transform_ecr_repository_images(repo_data: Dict) -> List[Dict]:
     repo_images_list = []
     for repo_uri, repo_images in repo_data.items():
         for img in repo_images:
-            if 'imageDigest' in img and img['imageDigest']:
-                img['repo_uri'] = repo_uri
+            if "imageDigest" in img and img["imageDigest"]:
+                img["repo_uri"] = repo_uri
                 repo_images_list.append(img)
             else:
                 logger.warning(
                     "Repo %s has an image that has no imageDigest. Its tag is %s. Continuing on.",
                     repo_uri,
-                    img.get('imageTag'),
+                    img.get("imageTag"),
                 )
 
     return repo_images_list
 
 
 def _load_ecr_repo_img_tx(
-    tx: neo4j.Transaction, repo_images_list: List[Dict], aws_update_tag: int,
+    tx: neo4j.Transaction,
+    repo_images_list: List[Dict],
+    aws_update_tag: int,
     region: str,
 ) -> None:
     query = """
@@ -139,23 +161,37 @@ def _load_ecr_repo_img_tx(
         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)
+    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,
+    neo4j_session: neo4j.Session,
+    repo_images_list: List[Dict],
+    region: str,
     aws_update_tag: int,
 ) -> None:
-    logger.info(f"Loading {len(repo_images_list)} ECR repository images in {region} into graph.")
+    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)
+        neo4j_session.write_transaction(
+            _load_ecr_repo_img_tx,
+            repo_image_batch,
+            aws_update_tag,
+            region,
+        )
 
 
 @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)
+    run_cleanup_job("aws_import_ecr_cleanup.json", neo4j_session, common_job_parameters)
 
 
 def _get_image_data(
@@ -163,15 +199,21 @@ def _get_image_data(
     region: str,
     repositories: List[Dict[str, Any]],
 ) -> Dict[str, Any]:
-    '''
+    """
     Given a list of repositories, get the image data for each repository,
     return as a mapping from repositoryUri to image object
-    '''
+    """
     image_data = {}
 
     async def async_get_images(repo: Dict[str, Any]) -> None:
-        repo_image_obj = await to_asynchronous(get_ecr_repository_images, boto3_session, region, repo['repositoryName'])
-        image_data[repo['repositoryUri']] = repo_image_obj
+        repo_image_obj = await to_asynchronous(
+            get_ecr_repository_images,
+            boto3_session,
+            region,
+            repo["repositoryName"],
+        )
+        image_data[repo["repositoryUri"]] = repo_image_obj
+
     to_synchronous(*[async_get_images(repo) for repo in repositories])
 
     return image_data
@@ -179,15 +221,29 @@ def _get_image_data(
 
 @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 ECR for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing ECR for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         image_data = {}
         repositories = get_ecr_repositories(boto3_session, region)
         image_data = _get_image_data(boto3_session, region, repositories)
-        load_ecr_repositories(neo4j_session, repositories, region, current_aws_account_id, update_tag)
+        load_ecr_repositories(
+            neo4j_session,
+            repositories,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
         repo_images_list = transform_ecr_repository_images(image_data)
         load_ecr_repository_images(neo4j_session, repo_images_list, region, update_tag)
     cleanup(neo4j_session, common_job_parameters)