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

cartography/intel/create_indexes.py

@@ -5,15 +5,16 @@ import neo4j
 
 from cartography.config import Config
 from cartography.util import load_resource_binary
+
 logger = logging.getLogger(__name__)
 
 
 def get_index_statements() -> List[str]:
     statements = []
-    with load_resource_binary('cartography.data', 'indexes.cypher') as f:
+    with load_resource_binary("cartography.data", "indexes.cypher") as f:
         for line in f.readlines():
             statements.append(
-                line.decode('UTF-8').rstrip('\r\n'),
+                line.decode("UTF-8").rstrip("\r\n"),
             )
     return statements
 

cartography/intel/crowdstrike/__init__.py

@@ -20,7 +20,8 @@ stat_handler = get_stats_client(__name__)
 
 @timeit
 def start_crowdstrike_ingestion(
-    neo4j_session: neo4j.Session, config: Config,
+    neo4j_session: neo4j.Session,
+    config: Config,
 ) -> None:
     """
     Perform ingestion of crowdstrike data.
@@ -31,10 +32,7 @@ def start_crowdstrike_ingestion(
     common_job_parameters = {
         "UPDATE_TAG": config.update_tag,
     }
-    if (
-        not config.crowdstrike_client_id or
-        not config.crowdstrike_client_secret
-    ):
+    if not config.crowdstrike_client_id or not config.crowdstrike_client_secret:
         logger.error("crowdstrike config not found")
         return
 
@@ -60,18 +58,22 @@ def start_crowdstrike_ingestion(
         group_id = config.crowdstrike_api_url
     merge_module_sync_metadata(
         neo4j_session,
-        group_type='crowdstrike',
+        group_type="crowdstrike",
         group_id=group_id,
-        synced_type='crowdstrike',
+        synced_type="crowdstrike",
         update_tag=config.update_tag,
         stat_handler=stat_handler,
     )
 
 
 @timeit
-def cleanup(neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]) -> None:
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
+) -> None:
     logger.info("Running Crowdstrike cleanup")
-    GraphJob.from_node_schema(CrowdstrikeHostSchema(), common_job_parameters).run(neo4j_session)
+    GraphJob.from_node_schema(CrowdstrikeHostSchema(), common_job_parameters).run(
+        neo4j_session
+    )
 
     # Cleanup other crowdstrike assets not handled by the data model
     run_cleanup_job(

cartography/intel/crowdstrike/endpoints.py

@@ -44,7 +44,11 @@ def load_host_data(
     )
 
 
-def get_host_ids(client: Hosts, crowdstrikeapi_filter: str = '', crowdstrikeapi_limit: int = 5000) -> List[List[str]]:
+def get_host_ids(
+    client: Hosts,
+    crowdstrikeapi_filter: str = "",
+    crowdstrikeapi_limit: int = 5000,
+) -> List[List[str]]:
     ids = []
     parameters = {"filter": crowdstrikeapi_filter, "limit": crowdstrikeapi_limit}
     response = client.QueryDevicesByFilter(parameters=parameters)

cartography/intel/crowdstrike/spotlight.py

@@ -25,7 +25,9 @@ def sync_vulnerabilities(
 
 
 def load_vulnerability_data(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Transform and load scan information
@@ -111,7 +113,9 @@ def _load_cves(neo4j_session: neo4j.Session, data: List[Dict], update_tag: int)
     )
 
 
-def get_spotlight_vulnerability_ids(client: Spotlight_Vulnerabilities) -> List[List[str]]:
+def get_spotlight_vulnerability_ids(
+    client: Spotlight_Vulnerabilities,
+) -> List[List[str]]:
     ids = []
     parameters = {"filter": 'status:!"closed"', "limit": 400}
     response = client.queryVulnerabilities(parameters=parameters)
@@ -135,7 +139,8 @@ def get_spotlight_vulnerability_ids(client: Spotlight_Vulnerabilities) -> List[L
 
 
 def get_spotlight_vulnerabilities(
-    client: Spotlight_Vulnerabilities, ids: List[str],
+    client: Spotlight_Vulnerabilities,
+    ids: List[str],
 ) -> List[Dict]:
     response = client.getVulnerabilities(ids=",".join(ids))
     body = response.get("body", {})

cartography/intel/cve/__init__.py

@@ -38,21 +38,33 @@ def _sync_year_archives(
 ) -> None:
     existing_years = feed.get_cve_sync_metadata(neo4j_session)
     current_year = datetime.now().year
-    logger.info(f"Syncing CVE data for year archives. Existing years: {existing_years}. Current year: {current_year}")
+    logger.info(
+        f"Syncing CVE data for year archives. Existing years: {existing_years}. Current year: {current_year}",
+    )
     for year in range(1999, current_year + 1):
         if year in existing_years:
             continue
         logger.info(f"Syncing CVE data for year {year}")
-        cves = feed.get_published_cves_per_year(http_session, config.nist_cve_url, str(year), cve_api_key)
+        cves = feed.get_published_cves_per_year(
+            http_session,
+            config.nist_cve_url,
+            str(year),
+            cve_api_key,
+        )
         feed_metadata = feed.transform_cve_feed(cves)
         feed.load_cve_feed(neo4j_session, [feed_metadata], config.update_tag)
         published_cves = feed.transform_cves(cves)
-        feed.load_cves(neo4j_session, published_cves, feed_metadata['FEED_ID'], config.update_tag)
+        feed.load_cves(
+            neo4j_session,
+            published_cves,
+            feed_metadata["FEED_ID"],
+            config.update_tag,
+        )
         merge_module_sync_metadata(
             neo4j_session,
-            group_type='CVE',
+            group_type="CVE",
             group_id=year,
-            synced_type='year',
+            synced_type="year",
             update_tag=config.update_tag,
             stat_handler=stat_handler,
         )
@@ -66,16 +78,26 @@ def _sync_modified_data(
 ) -> None:
     logger.info("Syncing CVE data for modified data")
     last_modified_date = feed.get_last_modified_cve_date(neo4j_session)
-    cves = feed.get_modified_cves(http_session, config.nist_cve_url, last_modified_date, cve_api_key)
+    cves = feed.get_modified_cves(
+        http_session,
+        config.nist_cve_url,
+        last_modified_date,
+        cve_api_key,
+    )
     feed_metadata = feed.transform_cve_feed(cves)
     feed.load_cve_feed(neo4j_session, [feed_metadata], config.update_tag)
     modified_cves = feed.transform_cves(cves)
-    feed.load_cves(neo4j_session, modified_cves, feed_metadata['FEED_ID'], config.update_tag)
+    feed.load_cves(
+        neo4j_session,
+        modified_cves,
+        feed_metadata["FEED_ID"],
+        config.update_tag,
+    )
     merge_module_sync_metadata(
         neo4j_session,
-        group_type='CVE',
-        group_id=feed_metadata['timestamp'][:4],
-        synced_type='modified',
+        group_type="CVE",
+        group_id=feed_metadata["timestamp"][:4],
+        synced_type="modified",
         update_tag=config.update_tag,
         stat_handler=stat_handler,
     )
@@ -83,7 +105,8 @@ def _sync_modified_data(
 
 @timeit
 def start_cve_ingestion(
-    neo4j_session: neo4j.Session, config: Config,
+    neo4j_session: neo4j.Session,
+    config: Config,
 ) -> None:
     """
     Perform ingestion of CVE data from NIST APIs.
@@ -95,6 +118,16 @@ def start_cve_ingestion(
         return
     cve_api_key: str | None = config.cve_api_key if config.cve_api_key else None
     with _retryable_session() as http_session:
-        _sync_year_archives(http_session, neo4j_session=neo4j_session, config=config, cve_api_key=cve_api_key)
-        _sync_modified_data(http_session, neo4j_session=neo4j_session, config=config, cve_api_key=cve_api_key)
+        _sync_year_archives(
+            http_session,
+            neo4j_session=neo4j_session,
+            config=config,
+            cve_api_key=cve_api_key,
+        )
+        _sync_modified_data(
+            http_session,
+            neo4j_session=neo4j_session,
+            config=config,
+            cve_api_key=cve_api_key,
+        )
         # CVEs are never deleted, so we don't need to run a cleanup job

cartography/intel/cve/feed.py

@@ -51,7 +51,10 @@ def get_last_modified_cve_date(neo4j_session: neo4j.Session) -> str:
     ORDER BY last_modified DESC
     LIMIT 1
     """
-    result = cast(neo4j.time.DateTime, read_single_value_tx(neo4j_session, query)).to_native()
+    result = cast(
+        neo4j.time.DateTime,
+        read_single_value_tx(neo4j_session, query),
+    ).to_native()
     return result.strftime("%Y-%m-%dT%H:%M:%S")
 
 
@@ -61,13 +64,19 @@ def _map_cve_dict(cve_dict: Dict[Any, Any], data: Dict[Any, Any]) -> None:
     cve_dict["timestamp"] = data["timestamp"]
     cve_dict["totalResults"] = data["totalResults"]
     cve_dict["vulnerabilities"] = cve_dict.get("vulnerabilities", []) + data.get(
-        "vulnerabilities", [],
+        "vulnerabilities",
+        [],
     )
     cve_dict["resultsPerPage"] = data["resultsPerPage"]
     cve_dict["startIndex"] = data["startIndex"]
 
 
-def _call_cves_api(http_session: Session, url: str, api_key: str | None, params: Dict[str, Any]) -> Dict[Any, Any]:
+def _call_cves_api(
+    http_session: Session,
+    url: str,
+    api_key: str | None,
+    params: Dict[str, Any],
+) -> Dict[Any, Any]:
     total_results = 0
     params["startIndex"] = 0
     params["resultsPerPage"] = RESULTS_PER_PAGE
@@ -84,7 +93,12 @@ def _call_cves_api(http_session: Session, url: str, api_key: str | None, params:
 
     while params["resultsPerPage"] > 0 or params["startIndex"] < total_results:
         logger.info(f"Calling NIST NVD API at {url} with params {params}")
-        res = http_session.get(url, params=params, headers=headers, timeout=CONNECT_AND_READ_TIMEOUT)
+        res = http_session.get(
+            url,
+            params=params,
+            headers=headers,
+            timeout=CONNECT_AND_READ_TIMEOUT,
+        )
         res.raise_for_status()
         data = res.json()
         _map_cve_dict(results, data)
@@ -140,7 +154,10 @@ def get_cves_in_batches(
 
 
 def get_modified_cves(
-    http_session: Session, nist_cve_url: str, last_modified_date: str, api_key: str | None,
+    http_session: Session,
+    nist_cve_url: str,
+    last_modified_date: str,
+    api_key: str | None,
 ) -> Dict[Any, Any]:
     end_date = datetime.now(tz=timezone.utc)
     start_date = datetime.strptime(last_modified_date, "%Y-%m-%dT%H:%M:%S").replace(
@@ -151,13 +168,21 @@ def get_modified_cves(
         "end": "lastModEndDate",
     }
     cves = get_cves_in_batches(
-        http_session, nist_cve_url, start_date, end_date, date_param_names, api_key,
+        http_session,
+        nist_cve_url,
+        start_date,
+        end_date,
+        date_param_names,
+        api_key,
     )
     return cves
 
 
 def get_published_cves_per_year(
-    http_session: Session, nist_cve_url: str, year: str, api_key: str | None,
+    http_session: Session,
+    nist_cve_url: str,
+    year: str,
+    api_key: str | None,
 ) -> Dict[Any, Any]:
     start_of_year = datetime.strptime(f"{year}-01-01", "%Y-%m-%d")
     next_year = int(year) + 1
@@ -167,7 +192,12 @@ def get_published_cves_per_year(
         "end": "pubEndDate",
     }
     cves = get_cves_in_batches(
-        http_session, nist_cve_url, start_of_year, end_of_next_year, date_param_names, api_key,
+        http_session,
+        nist_cve_url,
+        start_of_year,
+        end_of_next_year,
+        date_param_names,
+        api_key,
     )
     return cves
 
@@ -198,9 +228,13 @@ def transform_cves(cve_json: Dict[Any, Any]) -> List[Dict[Any, Any]]:
             ]
             cve["references_urls"] = [url["url"] for url in cve["references"]]
             if cve.get("weaknesses"):
-                weakness_descriptions = [weakness["description"] for weakness in cve["weaknesses"]]
+                weakness_descriptions = [
+                    weakness["description"] for weakness in cve["weaknesses"]
+                ]
                 weakness_descriptions = reduce(
-                    lambda x, y: x + y, weakness_descriptions, [],
+                    lambda x, y: x + y,
+                    weakness_descriptions,
+                    [],
                 )
                 cve["weaknesses"] = [
                     description["value"]
@@ -226,7 +260,7 @@ def transform_cves(cve_json: Dict[Any, Any]) -> List[Dict[Any, Any]]:
                 cve["exploitabilityScore"] = cvss31.get("exploitabilityScore")
                 cve["impactScore"] = cvss31.get("impactScore")
         except Exception:
-            logger.error("Failed to transform CVE data {data}")
+            logger.error(f"Failed to transform CVE data {data}")
             raise
         cves.append(cve)
     return cves
@@ -265,7 +299,9 @@ def load_cves(
 
 
 def load_cve_feed(
-    neo4j_session: neo4j.Session, data: List[Dict[str, Any]], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    update_tag: int,
 ) -> None:
     """
     Load CVE feed information

cartography/intel/digitalocean/__init__.py

@@ -9,7 +9,6 @@ from cartography.intel.digitalocean import management
 from cartography.intel.digitalocean import platform
 from cartography.util import timeit
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -23,7 +22,9 @@ def start_digitalocean_ingestion(neo4j_session: neo4j.Session, config: Config) -
     """
 
     if not config.digitalocean_token:
-        logger.info('DigitalOcean import is not configured - skipping this module. See docs to configure.')
+        logger.info(
+            "DigitalOcean import is not configured - skipping this module. See docs to configure.",
+        )
         return
 
     common_job_parameters = {
@@ -31,14 +32,22 @@ def start_digitalocean_ingestion(neo4j_session: neo4j.Session, config: Config) -
     }
     manager = Manager(token=config.digitalocean_token)
 
-    """
-    Get Account ID related to this credentials and pass it along in `common_job_parameters` to avoid cleaning up other
-    accounts resources
-    """
-    account = manager.get_account()
-    common_job_parameters["DO_ACCOUNT_ID"] = account.uuid
-
-    platform.sync(neo4j_session, account, config.update_tag, common_job_parameters)
-    project_resources = management.sync(neo4j_session, manager, config.update_tag, common_job_parameters)
-    compute.sync(neo4j_session, manager, project_resources, config.update_tag, common_job_parameters)
-    return
+    account_id = platform.sync(
+        neo4j_session, manager, config.update_tag, common_job_parameters
+    )
+    common_job_parameters["ACCOUNT_ID"] = str(account_id)
+    projects_resources = management.sync(
+        neo4j_session,
+        manager,
+        account_id,
+        config.update_tag,
+        common_job_parameters,
+    )
+    compute.sync(
+        neo4j_session,
+        manager,
+        account_id,
+        projects_resources,
+        config.update_tag,
+        common_job_parameters,
+    )

cartography/intel/digitalocean/compute.py

@@ -1,10 +1,15 @@
 import logging
+from typing import Any
+from typing import Dict
+from typing import List
 from typing import Optional
 
 import neo4j
 from digitalocean import Manager
 
-from cartography.util import run_cleanup_job
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.digitalocean.droplet import DODropletSchema
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
@@ -12,29 +17,20 @@ logger = logging.getLogger(__name__)
 
 @timeit
 def sync(
-        neo4j_session: neo4j.Session,
-        manager: Manager,
-        projects_resources: dict,
-        digitalocean_update_tag: int,
-        common_job_parameters: dict,
-) -> None:
-    sync_droplets(neo4j_session, manager, projects_resources, digitalocean_update_tag, common_job_parameters)
-
-
-@timeit
-def sync_droplets(
-        neo4j_session: neo4j.Session,
-        manager: Manager,
-        projects_resources: dict,
-        digitalocean_update_tag: int,
-        common_job_parameters: dict,
+    neo4j_session: neo4j.Session,
+    manager: Manager,
+    account_id: str,
+    projects_resources: dict,
+    update_tag: int,
+    common_job_parameters: dict,
 ) -> None:
     logger.info("Syncing Droplets")
-    account_id = common_job_parameters['DO_ACCOUNT_ID']
     droplets_res = get_droplets(manager)
-    droplets = transform_droplets(droplets_res, account_id, projects_resources)
-    load_droplets(neo4j_session, droplets, digitalocean_update_tag)
-    cleanup_droplets(neo4j_session, common_job_parameters)
+    droplets_by_project = transform_droplets(
+        droplets_res, account_id, projects_resources
+    )
+    load_droplets(neo4j_session, account_id, droplets_by_project, update_tag)
+    cleanup(neo4j_session, list(droplets_by_project.keys()), common_job_parameters)
 
 
 @timeit
@@ -43,35 +39,45 @@ def get_droplets(manager: Manager) -> list:
 
 
 @timeit
-def transform_droplets(droplets_res: list, account_id: str, projects_resources: dict) -> list:
-    droplets = list()
+def transform_droplets(
+    droplets_res: list,
+    account_id: str,
+    projects_resources: dict,
+) -> Dict[str, List[Dict[str, Any]]]:
+    droplets_by_project: Dict[str, List[Dict[str, Any]]] = {}
     for d in droplets_res:
+        project_id = str(_get_project_id_for_droplet(d.id, projects_resources))
+        if project_id not in droplets_by_project:
+            droplets_by_project[project_id] = []
         droplet = {
-            'id': d.id,
-            'name': d.name,
-            'locked': d.locked,
-            'status': d.status,
-            'features': d.features,
-            'region': d.region['slug'],
-            'created_at': d.created_at,
-            'image': d.image['slug'],
-            'size': d.size_slug,
-            'kernel': d.kernel,
-            'tags': d.tags,
-            'volumes': d.volume_ids,
-            'vpc_uuid': d.vpc_uuid,
-            'ip_address': d.ip_address,
-            'private_ip_address': d.private_ip_address,
-            'ip_v6_address': d.ip_v6_address,
-            'account_id': account_id,
-            'project_id': _get_project_id_for_droplet(d.id, projects_resources),
+            "id": d.id,
+            "name": d.name,
+            "locked": d.locked,
+            "status": d.status,
+            "features": d.features,
+            "region": d.region["slug"],
+            "created_at": d.created_at,
+            "image": d.image["slug"],
+            "size": d.size_slug,
+            "kernel": d.kernel,
+            "tags": d.tags,
+            "volumes": d.volume_ids,
+            "vpc_uuid": d.vpc_uuid,
+            "ip_address": d.ip_address,
+            "private_ip_address": d.private_ip_address,
+            "ip_v6_address": d.ip_v6_address,
+            "account_id": account_id,
+            "project_id": _get_project_id_for_droplet(d.id, projects_resources),
         }
-        droplets.append(droplet)
-    return droplets
+        droplets_by_project[project_id].append(droplet)
+    return droplets_by_project
 
 
 @timeit
-def _get_project_id_for_droplet(droplet_id: int, project_resources: dict) -> Optional[str]:
+def _get_project_id_for_droplet(
+    droplet_id: int,
+    project_resources: dict,
+) -> Optional[str]:
     for project_id, resource_list in project_resources.items():
         droplet_resource_name = "do:droplet:" + str(droplet_id)
         if droplet_resource_name in resource_list:
@@ -80,71 +86,32 @@ def _get_project_id_for_droplet(droplet_id: int, project_resources: dict) -> Opt
 
 
 @timeit
-def load_droplets(neo4j_session: neo4j.Session, data: list, digitalocean_update_tag: int) -> None:
-    query = """
-        MERGE (p:DOProject{id:$ProjectId})
-        ON CREATE SET p.firstseen = timestamp()
-        SET p.lastupdated = $digitalocean_update_tag
-
-        MERGE (d:DODroplet{id:$DropletId})
-        ON CREATE SET d.firstseen = timestamp()
-        SET d.account_id = $AccountId,
-        d.name = $Name,
-        d.locked = $Locked,
-        d.status = $Status,
-        d.features = $Features,
-        d.region = $RegionSlug,
-        d.created_at = $CreatedAt,
-        d.image = $ImageSlug,
-        d.size = $SizeSlug,
-        d.kernel = $Kernel,
-        d.ip_address = $IpAddress,
-        d.private_ip_address = $PrivateIpAddress,
-        d.project_id = $ProjectId,
-        d.ip_v6_address = $IpV6Address,
-        d.tags = $Tags,
-        d.volumes = $Volumes,
-        d.vpc_uuid = $VpcUuid,
-        d.lastupdated = $digitalocean_update_tag
-        WITH d, p
-
-        MERGE (p)-[r:RESOURCE]->(d)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $digitalocean_update_tag
-        """
-    for droplet in data:
-        neo4j_session.run(
-            query,
-            AccountId=droplet['account_id'],
-            DropletId=droplet['id'],
-            Name=droplet['name'],
-            Locked=droplet['locked'],
-            Status=droplet['status'],
-            Features=droplet['features'],
-            RegionSlug=droplet['region'],
-            CreatedAt=droplet['created_at'],
-            ImageSlug=droplet['image'],
-            SizeSlug=droplet['size'],
-            IpAddress=droplet['ip_address'],
-            PrivateIpAddress=droplet['private_ip_address'],
-            ProjectId=droplet['project_id'],
-            IpV6Address=droplet['ip_v6_address'],
-            Kernel=droplet['kernel'],
-            Tags=droplet['tags'],
-            Volumes=droplet['volumes'],
-            VpcUuid=droplet['vpc_uuid'],
-            digitalocean_update_tag=digitalocean_update_tag,
+def load_droplets(
+    neo4j_session: neo4j.Session,
+    account_id: str,
+    data: Dict[str, List[Dict[str, Any]]],
+    update_tag: int,
+) -> None:
+    for project_id, droplets in data.items():
+        load(
+            neo4j_session,
+            DODropletSchema(),
+            droplets,
+            lastupdated=update_tag,
+            PROJECT_ID=str(project_id),
+            ACCOUNT_ID=str(account_id),
         )
-    return
 
 
 @timeit
-def cleanup_droplets(neo4j_session: neo4j.Session, common_job_parameters: dict) -> None:
-    """
-        Delete out-of-date DigitalOcean droplets and relationships
-        :param neo4j_session: The Neo4j session
-        :param common_job_parameters: dict of other job parameters to pass to Neo4j
-        :return: Nothing
-        """
-    run_cleanup_job('digitalocean_droplet_cleanup.json', neo4j_session, common_job_parameters)
-    return
+def cleanup(
+    neo4j_session: neo4j.Session,
+    projects_ids: List[str],
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    for project_id in projects_ids:
+        parameters = common_job_parameters.copy()
+        parameters["PROJECT_ID"] = str(project_id)
+        GraphJob.from_node_schema(DODropletSchema(), parameters).run(
+            neo4j_session,
+        )