cartography 0.112.0__py3-none-any.whl → 0.114.0__py3-none-any.whl

cartography/intel/gcp/dns.py

@@ -7,28 +7,20 @@ import neo4j
 from googleapiclient.discovery import HttpError
 from googleapiclient.discovery import Resource
 
-from cartography.util import run_cleanup_job
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.gcp.dns import GCPDNSZoneSchema
+from cartography.models.gcp.dns import GCPRecordSetSchema
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
 
 
 @timeit
-def get_dns_zones(dns: Resource, project_id: str) -> List[Resource]:
-    """
-    Returns a list of DNS zones within the given project.
-
-    :type dns: The GCP DNS resource object
-    :param dns: The DNS resource object created by googleapiclient.discovery.build()
-
-    :type project_id: str
-    :param project_id: Current Google Project Id
-
-    :rtype: list
-    :return: List of DNS zones
-    """
+def get_dns_zones(dns: Resource, project_id: str) -> List[Dict]:
+    """Returns a list of DNS zones within the given project."""
     try:
-        zones = []
+        zones: List[Dict] = []
         request = dns.managedZones().list(project=project_id)
         while request is not None:
             response = request.execute()
@@ -47,40 +39,22 @@ def get_dns_zones(dns: Resource, project_id: str) -> List[Resource]:
         ):
             logger.warning(
                 (
-                    "Could not retrieve DNS zones on project %s due to permissions issues. Code: %s, Message: %s"
+                    "Could not retrieve DNS zones on project %s due to permissions issues. "
+                    "Code: %s, Message: %s"
                 ),
                 project_id,
                 err["code"],
                 err["message"],
             )
             return []
-        else:
-            raise
+        raise
 
 
 @timeit
-def get_dns_rrs(
-    dns: Resource,
-    dns_zones: List[Dict],
-    project_id: str,
-) -> List[Resource]:
-    """
-    Returns a list of DNS Resource Record Sets within the given project.
-
-    :type dns: The GCP DNS resource object
-    :param dns: The DNS resource object created by googleapiclient.discovery.build()
-
-    :type dns_zones: list
-    :param dns_zones: List of DNS zones for the project
-
-    :type project_id: str
-    :param project_id: Current Google Project Id
-
-    :rtype: list
-    :return: List of Resource Record Sets
-    """
+def get_dns_rrs(dns: Resource, dns_zones: List[Dict], project_id: str) -> List[Dict]:
+    """Returns a list of DNS Resource Record Sets within the given project."""
     try:
-        rrs: List[Resource] = []
+        rrs: List[Dict] = []
         for zone in dns_zones:
             request = dns.resourceRecordSets().list(
                 project=project_id,
@@ -104,16 +78,53 @@ def get_dns_rrs(
         ):
             logger.warning(
                 (
-                    "Could not retrieve DNS RRS on project %s due to permissions issues. Code: %s, Message: %s"
+                    "Could not retrieve DNS RRS on project %s due to permissions issues. "
+                    "Code: %s, Message: %s"
                 ),
                 project_id,
                 err["code"],
                 err["message"],
             )
             return []
-        else:
-            raise
-        raise e
+        raise
+
+
+@timeit
+def transform_dns_zones(dns_zones: List[Dict]) -> List[Dict]:
+    """Transform raw DNS zone responses into Neo4j-ready dicts."""
+    zones: List[Dict] = []
+    for z in dns_zones:
+        zones.append(
+            {
+                "id": z["id"],
+                "name": z.get("name"),
+                "dns_name": z.get("dnsName"),
+                "description": z.get("description"),
+                "visibility": z.get("visibility"),
+                "kind": z.get("kind"),
+                "nameservers": z.get("nameServers"),
+                "created_at": z.get("creationTime"),
+            }
+        )
+    return zones
+
+
+@timeit
+def transform_dns_rrs(dns_rrs: List[Dict]) -> List[Dict]:
+    """Transform raw DNS record set responses into Neo4j-ready dicts."""
+    records: List[Dict] = []
+    for r in dns_rrs:
+        records.append(
+            {
+                "id": r["name"],
+                "name": r["name"],
+                "type": r.get("type"),
+                "ttl": r.get("ttl"),
+                "data": r.get("rrdatas"),
+                "zone_id": r.get("zone"),
+            }
+        )
+    return records
 
 
 @timeit
@@ -123,102 +134,30 @@ def load_dns_zones(
     project_id: str,
     gcp_update_tag: int,
 ) -> None:
-    """
-    Ingest GCP DNS Zones into Neo4j
-
-    :type neo4j_session: Neo4j session object
-    :param neo4j session: The Neo4j session object
-
-    :type dns_resp: Dict
-    :param dns_resp: A DNS response object from the GKE API
-
-    :type project_id: str
-    :param project_id: Current Google Project Id
-
-    :type gcp_update_tag: timestamp
-    :param gcp_update_tag: The timestamp value to set our new Neo4j nodes with
-
-    :rtype: NoneType
-    :return: Nothing
-    """
-
-    ingest_records = """
-    UNWIND $records as record
-    MERGE(zone:GCPDNSZone{id:record.id})
-    ON CREATE SET
-        zone.firstseen = timestamp(),
-        zone.created_at = record.creationTime
-    SET
-        zone.name = record.name,
-        zone.dns_name = record.dnsName,
-        zone.description = record.description,
-        zone.visibility = record.visibility,
-        zone.kind = record.kind,
-        zone.nameservers = record.nameServers,
-        zone.lastupdated = $gcp_update_tag
-    WITH zone
-    MATCH (owner:GCPProject{id:$ProjectId})
-    MERGE (owner)-[r:RESOURCE]->(zone)
-    ON CREATE SET
-        r.firstseen = timestamp(),
-        r.lastupdated = $gcp_update_tag
-    """
-    neo4j_session.run(
-        ingest_records,
-        records=dns_zones,
-        ProjectId=project_id,
-        gcp_update_tag=gcp_update_tag,
+    """Ingest GCP DNS Zones into Neo4j."""
+    load(
+        neo4j_session,
+        GCPDNSZoneSchema(),
+        dns_zones,
+        lastupdated=gcp_update_tag,
+        PROJECT_ID=project_id,
     )
 
 
 @timeit
 def load_rrs(
     neo4j_session: neo4j.Session,
-    dns_rrs: List[Resource],
+    dns_rrs: List[Dict],
     project_id: str,
     gcp_update_tag: int,
 ) -> None:
-    """
-    Ingest GCP RRS into Neo4j
-
-    :type neo4j_session: Neo4j session object
-    :param neo4j session: The Neo4j session object
-
-    :type dns_rrs: list
-    :param dns_rrs: A list of RRS
-
-    :type project_id: str
-    :param project_id: Current Google Project Id
-
-    :type gcp_update_tag: timestamp
-    :param gcp_update_tag: The timestamp value to set our new Neo4j nodes with
-
-    :rtype: NoneType
-    :return: Nothing
-    """
-
-    ingest_records = """
-    UNWIND $records as record
-    MERGE(rrs:GCPRecordSet{id:record.name})
-    ON CREATE SET
-        rrs.firstseen = timestamp()
-    SET
-        rrs.name = record.name,
-        rrs.type = record.type,
-        rrs.ttl = record.ttl,
-        rrs.data = record.rrdatas,
-        rrs.lastupdated = $gcp_update_tag
-    WITH rrs, record
-    MATCH (zone:GCPDNSZone{id:record.zone})
-    MERGE (zone)-[r:HAS_RECORD]->(rrs)
-    ON CREATE SET
-        r.firstseen = timestamp(),
-        r.lastupdated = $gcp_update_tag
-    """
-    neo4j_session.run(
-        ingest_records,
-        records=dns_rrs,
-        gcp_update_tag=gcp_update_tag,
+    """Ingest GCP DNS Resource Record Sets into Neo4j."""
+    load(
+        neo4j_session,
+        GCPRecordSetSchema(),
+        dns_rrs,
+        lastupdated=gcp_update_tag,
+        PROJECT_ID=project_id,
     )
 
 
@@ -227,19 +166,14 @@ def cleanup_dns_records(
     neo4j_session: neo4j.Session,
     common_job_parameters: Dict,
 ) -> None:
-    """
-    Delete out-of-date GCP DNS Zones and RRS nodes and relationships
-
-    :type neo4j_session: The Neo4j session object
-    :param neo4j_session: The Neo4j session
-
-    :type common_job_parameters: dict
-    :param common_job_parameters: Dictionary of other job parameters to pass to Neo4j
-
-    :rtype: NoneType
-    :return: Nothing
-    """
-    run_cleanup_job("gcp_dns_cleanup.json", neo4j_session, common_job_parameters)
+    """Delete out-of-date GCP DNS Zones and Record Sets nodes and relationships."""
+    # Record sets depend on zones, so clean them up first.
+    GraphJob.from_node_schema(GCPRecordSetSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(GCPDNSZoneSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
 
 
 @timeit
@@ -250,33 +184,12 @@ def sync(
     gcp_update_tag: int,
     common_job_parameters: Dict,
 ) -> None:
-    """
-    Get GCP DNS Zones and Resource Record Sets using the DNS resource object, ingest to Neo4j, and clean up old data.
-
-    :type neo4j_session: The Neo4j session object
-    :param neo4j_session: The Neo4j session
-
-    :type dns: The DNS resource object created by googleapiclient.discovery.build()
-    :param dns: The GCP DNS resource object
-
-    :type project_id: str
-    :param project_id: The project ID of the corresponding project
-
-    :type gcp_update_tag: timestamp
-    :param gcp_update_tag: The timestamp value to set our new Neo4j nodes with
-
-    :type common_job_parameters: dict
-    :param common_job_parameters: Dictionary of other job parameters to pass to Neo4j
-
-    :rtype: NoneType
-    :return: Nothing
-    """
+    """Get GCP DNS Zones and Record Sets, load them into Neo4j, and clean up old data."""
     logger.info("Syncing DNS records for project %s.", project_id)
-    # DNS ZONES
-    dns_zones = get_dns_zones(dns, project_id)
+    dns_zones_resp = get_dns_zones(dns, project_id)
+    dns_zones = transform_dns_zones(dns_zones_resp)
     load_dns_zones(neo4j_session, dns_zones, project_id, gcp_update_tag)
-    # RECORD SETS
-    dns_rrs = get_dns_rrs(dns, dns_zones, project_id)
+    dns_rrs_resp = get_dns_rrs(dns, dns_zones_resp, project_id)
+    dns_rrs = transform_dns_rrs(dns_rrs_resp)
     load_rrs(neo4j_session, dns_rrs, project_id, gcp_update_tag)
-    # TODO scope the cleanup to the current project - https://github.com/cartography-cncf/cartography/issues/381
     cleanup_dns_records(neo4j_session, common_job_parameters)

cartography/intel/gcp/gke.py

@@ -1,12 +1,16 @@
 import json
 import logging
+from typing import Any
 from typing import Dict
+from typing import List
 
 import neo4j
 from googleapiclient.discovery import HttpError
 from googleapiclient.discovery import Resource
 
-from cartography.util import run_cleanup_job
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.gcp.gke import GCPGKEClusterSchema
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
@@ -56,105 +60,20 @@ def load_gke_clusters(
     gcp_update_tag: int,
 ) -> None:
     """
-    Ingest GCP GKE Clusters to Neo4j
-
-    :type neo4j_session: Neo4j session object
-    :param neo4j session: The Neo4j session object
-
-    :type cluster_resp: Dict
-    :param cluster_resp: A cluster response object from the GKE API
-
-    :type gcp_update_tag: timestamp
-    :param gcp_update_tag: The timestamp value to set our new Neo4j nodes with
-
-    :rtype: NoneType
-    :return: Nothing
+    Ingest GCP GKE clusters using the data model loader.
     """
+    clusters: List[Dict[str, Any]] = transform_gke_clusters(cluster_resp)
 
-    query = """
-    MERGE(cluster:GKECluster{id:$ClusterSelfLink})
-    ON CREATE SET
-        cluster.firstseen = timestamp(),
-        cluster.created_at = $ClusterCreateTime
-    SET
-        cluster.name = $ClusterName,
-        cluster.self_link = $ClusterSelfLink,
-        cluster.description = $ClusterDescription,
-        cluster.logging_service = $ClusterLoggingService,
-        cluster.monitoring_service = $ClusterMonitoringService,
-        cluster.network = $ClusterNetwork,
-        cluster.subnetwork = $ClusterSubnetwork,
-        cluster.cluster_ipv4cidr = $ClusterIPv4Cidr,
-        cluster.zone = $ClusterZone,
-        cluster.location = $ClusterLocation,
-        cluster.endpoint = $ClusterEndpoint,
-        cluster.initial_version = $ClusterInitialVersion,
-        cluster.current_master_version = $ClusterMasterVersion,
-        cluster.status = $ClusterStatus,
-        cluster.services_ipv4cidr = $ClusterServicesIPv4Cidr,
-        cluster.database_encryption = $ClusterDatabaseEncryption,
-        cluster.network_policy = $ClusterNetworkPolicy,
-        cluster.master_authorized_networks = $ClusterMasterAuthorizedNetworks,
-        cluster.legacy_abac = $ClusterAbac,
-        cluster.shielded_nodes = $ClusterShieldedNodes,
-        cluster.private_nodes = $ClusterPrivateNodes,
-        cluster.private_endpoint_enabled = $ClusterPrivateEndpointEnabled,
-        cluster.private_endpoint = $ClusterPrivateEndpoint,
-        cluster.public_endpoint = $ClusterPublicEndpoint,
-        cluster.masterauth_username = $ClusterMasterUsername,
-        cluster.masterauth_password = $ClusterMasterPassword
-    WITH cluster
-    MATCH (owner:GCPProject{id:$ProjectId})
-    MERGE (owner)-[r:RESOURCE]->(cluster)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $gcp_update_tag
-    """
-    for cluster in cluster_resp.get("clusters", []):
-        neo4j_session.run(
-            query,
-            ProjectId=project_id,
-            ClusterSelfLink=cluster["selfLink"],
-            ClusterCreateTime=cluster["createTime"],
-            ClusterName=cluster["name"],
-            ClusterDescription=cluster.get("description"),
-            ClusterLoggingService=cluster.get("loggingService"),
-            ClusterMonitoringService=cluster.get("monitoringService"),
-            ClusterNetwork=cluster.get("network"),
-            ClusterSubnetwork=cluster.get("subnetwork"),
-            ClusterIPv4Cidr=cluster.get("clusterIpv4Cidr"),
-            ClusterZone=cluster.get("zone"),
-            ClusterLocation=cluster.get("location"),
-            ClusterEndpoint=cluster.get("endpoint"),
-            ClusterInitialVersion=cluster.get("initialClusterVersion"),
-            ClusterMasterVersion=cluster.get("currentMasterVersion"),
-            ClusterStatus=cluster.get("status"),
-            ClusterServicesIPv4Cidr=cluster.get("servicesIpv4Cidr"),
-            ClusterDatabaseEncryption=cluster.get("databaseEncryption", {}).get(
-                "state",
-            ),
-            ClusterNetworkPolicy=_process_network_policy(cluster),
-            ClusterMasterAuthorizedNetworks=cluster.get(
-                "masterAuthorizedNetworksConfig",
-                {},
-            ).get("enabled"),
-            ClusterAbac=cluster.get("legacyAbac", {}).get("enabled"),
-            ClusterShieldedNodes=cluster.get("shieldedNodes", {}).get("enabled"),
-            ClusterPrivateNodes=cluster.get("privateClusterConfig", {}).get(
-                "enablePrivateNodes",
-            ),
-            ClusterPrivateEndpointEnabled=cluster.get("privateClusterConfig", {}).get(
-                "enablePrivateEndpoint",
-            ),
-            ClusterPrivateEndpoint=cluster.get("privateClusterConfig", {}).get(
-                "privateEndpoint",
-            ),
-            ClusterPublicEndpoint=cluster.get("privateClusterConfig", {}).get(
-                "publicEndpoint",
-            ),
-            ClusterMasterUsername=cluster.get("masterAuth", {}).get("username"),
-            ClusterMasterPassword=cluster.get("masterAuth", {}).get("password"),
-            gcp_update_tag=gcp_update_tag,
-        )
+    if not clusters:
+        return
+
+    load(
+        neo4j_session,
+        GCPGKEClusterSchema(),
+        clusters,
+        lastupdated=gcp_update_tag,
+        PROJECT_ID=project_id,
+    )
 
 
 def _process_network_policy(cluster: Dict) -> bool:
@@ -175,21 +94,10 @@ def cleanup_gke_clusters(
     common_job_parameters: Dict,
 ) -> None:
     """
-    Delete out-of-date GCP GKE Clusters nodes and relationships
-
-    :type neo4j_session: The Neo4j session object
-    :param neo4j_session: The Neo4j session
-
-    :type common_job_parameters: dict
-    :param common_job_parameters: Dictionary of other job parameters to pass to Neo4j
-
-    :rtype: NoneType
-    :return: Nothing
+    Scoped cleanup for GKE clusters based on the project sub-resource relationship.
     """
-    run_cleanup_job(
-        "gcp_gke_cluster_cleanup.json",
+    GraphJob.from_node_schema(GCPGKEClusterSchema(), common_job_parameters).run(
         neo4j_session,
-        common_job_parameters,
     )
 
 
@@ -222,8 +130,59 @@ def sync_gke_clusters(
     :rtype: NoneType
     :return: Nothing
     """
-    logger.info("Syncing Compute objects for project %s.", project_id)
+    logger.info("Syncing GKE clusters for project %s.", project_id)
     gke_res = get_gke_clusters(container, project_id)
     load_gke_clusters(neo4j_session, gke_res, project_id, gcp_update_tag)
-    # TODO scope the cleanup to the current project - https://github.com/cartography-cncf/cartography/issues/381
     cleanup_gke_clusters(neo4j_session, common_job_parameters)
+
+
+def transform_gke_clusters(api_result: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """
+    Transform GKE API response into a list of dicts suitable for the data model loader.
+    """
+    result: List[Dict[str, Any]] = []
+    for c in api_result.get("clusters", []):
+        transformed: Dict[str, Any] = {
+            # Required fields
+            "id": c["selfLink"],
+            "self_link": c["selfLink"],
+            "name": c["name"],
+            "created_at": c.get("createTime"),
+            # Optional fields
+            "description": c.get("description"),
+            "logging_service": c.get("loggingService"),
+            "monitoring_service": c.get("monitoringService"),
+            "network": c.get("network"),
+            "subnetwork": c.get("subnetwork"),
+            "cluster_ipv4cidr": c.get("clusterIpv4Cidr"),
+            "zone": c.get("zone"),
+            "location": c.get("location"),
+            "endpoint": c.get("endpoint"),
+            "initial_version": c.get("initialClusterVersion"),
+            "current_master_version": c.get("currentMasterVersion"),
+            "status": c.get("status"),
+            "services_ipv4cidr": c.get("servicesIpv4Cidr"),
+            "database_encryption": (c.get("databaseEncryption", {}) or {}).get("state"),
+            "network_policy": _process_network_policy(c),
+            "master_authorized_networks": (
+                c.get("masterAuthorizedNetworksConfig", {}) or {}
+            ).get("enabled"),
+            "legacy_abac": (c.get("legacyAbac", {}) or {}).get("enabled"),
+            "shielded_nodes": (c.get("shieldedNodes", {}) or {}).get("enabled"),
+            "private_nodes": (c.get("privateClusterConfig", {}) or {}).get(
+                "enablePrivateNodes"
+            ),
+            "private_endpoint_enabled": (c.get("privateClusterConfig", {}) or {}).get(
+                "enablePrivateEndpoint"
+            ),
+            "private_endpoint": (c.get("privateClusterConfig", {}) or {}).get(
+                "privateEndpoint"
+            ),
+            "public_endpoint": (c.get("privateClusterConfig", {}) or {}).get(
+                "publicEndpoint"
+            ),
+            "masterauth_username": (c.get("masterAuth", {}) or {}).get("username"),
+            "masterauth_password": (c.get("masterAuth", {}) or {}).get("password"),
+        }
+        result.append(transformed)
+    return result