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

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,
+        )

cartography/intel/digitalocean/management.py

@@ -1,9 +1,14 @@
 import logging
+from typing import Any
+from typing import Dict
+from typing import List
 
 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.project import DOProjectSchema
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
@@ -11,31 +16,19 @@ logger = logging.getLogger(__name__)
 
 @timeit
 def sync(
-        neo4j_session: neo4j.Session,
-        manager: Manager,
-        digitalocean_update_tag: int,
-        common_job_parameters: dict,
-) -> dict:
-    projects_resources = sync_projects(neo4j_session, manager, digitalocean_update_tag, common_job_parameters)
-    return projects_resources
-
-
-@timeit
-def sync_projects(
-        neo4j_session: neo4j.Session,
-        manager: Manager,
-        digitalocean_update_tag: int,
-        common_job_parameters: dict,
+    neo4j_session: neo4j.Session,
+    manager: Manager,
+    account_id: str,
+    update_tag: int,
+    common_job_parameters: dict,
 ) -> dict:
     logger.info("Syncing Projects")
-    account_id = common_job_parameters['DO_ACCOUNT_ID']
     projects_res = get_projects(manager)
-    projects_resources = get_projects_resources(projects_res)
-    projects = transform_projects(projects_res, account_id)
-    load_projects(neo4j_session, projects, digitalocean_update_tag)
-    cleanup_projects(neo4j_session, common_job_parameters)
+    projects = transform_projects(projects_res)
+    load_projects(neo4j_session, projects, account_id, update_tag)
+    cleanup(neo4j_session, common_job_parameters)
 
-    return projects_resources
+    return get_projects_resources(projects_res)
 
 
 @timeit
@@ -49,77 +42,48 @@ def get_projects_resources(projects_res: list) -> dict:
     for p in projects_res:
         resources = p.get_all_resources()
         result[p.id] = resources
-
     return result
 
 
 @timeit
-def transform_projects(project_res: list, account_id: str) -> list:
+def transform_projects(project_res: list) -> list:
     result = list()
     for p in project_res:
         project = {
-            'id': p.id,
-            'name': p.name,
-            'owner_uuid': p.owner_uuid,
-            'description': p.description,
-            'environment': p.environment,
-            'is_default': p.is_default,
-            'created_at': p.created_at,
-            'updated_at': p.updated_at,
-            'account_id': account_id,
+            "id": p.id,
+            "name": p.name,
+            "owner_uuid": p.owner_uuid,
+            "description": p.description,
+            "environment": p.environment,
+            "is_default": p.is_default,
+            "created_at": p.created_at,
+            "updated_at": p.updated_at,
         }
         result.append(project)
     return result
 
 
 @timeit
-def load_projects(neo4j_session: neo4j.Session, data: list, digitalocean_update_tag: int) -> None:
-    query = """
-        MERGE (a:DOAccount{id:$AccountId})
-        ON CREATE SET a.firstseen = timestamp()
-        SET a.lastupdated = $digitalocean_update_tag
-
-        MERGE (p:DOProject{id:$ProjectId})
-        ON CREATE SET p.firstseen = timestamp()
-        SET p.account_id = $AccountId,
-        p.name = $Name,
-        p.owner_uuid = $OwnerUuid,
-        p.description = $Description,
-        p.environment = $Environment,
-        p.is_default = $IsDefault,
-        p.created_at = $CreatedAt,
-        p.updated_at = $UpdatedAt,
-        p.lastupdated = $digitalocean_update_tag
-        WITH p, a
-
-        MERGE (a)-[r:RESOURCE]->(p)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $digitalocean_update_tag
-        """
-    for project in data:
-        neo4j_session.run(
-            query,
-            AccountId=project['account_id'],
-            ProjectId=project['id'],
-            Name=project['name'],
-            OwnerUuid=project['owner_uuid'],
-            Description=project['description'],
-            Environment=project['environment'],
-            IsDefault=project['is_default'],
-            CreatedAt=project['created_at'],
-            UpdatedAt=project['updated_at'],
-            digitalocean_update_tag=digitalocean_update_tag,
-        )
-    return
+def load_projects(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    account_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        DOProjectSchema(),
+        data,
+        lastupdated=update_tag,
+        ACCOUNT_ID=str(account_id),
+    )
 
 
 @timeit
-def cleanup_projects(neo4j_session: neo4j.Session, common_job_parameters: dict) -> None:
-    """
-        Delete out-of-date DigitalOcean projects 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_project_cleanup.json', neo4j_session, common_job_parameters)
-    return
+def cleanup(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    GraphJob.from_node_schema(DOProjectSchema(), common_job_parameters).run(
+        neo4j_session,
+    )