cartography 0.104.0rc2__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/kubernetes/namespaces.py

@@ -1,82 +1,84 @@
 import logging
-from typing import Dict
-from typing import List
-from typing import Tuple
+from typing import Any
 
-from neo4j import Session
+import neo4j
+from kubernetes.client.models import V1Namespace
 
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
 from cartography.intel.kubernetes.util import get_epoch
+from cartography.intel.kubernetes.util import k8s_paginate
 from cartography.intel.kubernetes.util import K8sClient
-from cartography.stats import get_stats_client
-from cartography.util import merge_module_sync_metadata
+from cartography.models.kubernetes.namespaces import KubernetesNamespaceSchema
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
-stat_handler = get_stats_client(__name__)
 
 
 @timeit
-def sync_namespaces(session: Session, client: K8sClient, update_tag: int) -> Dict:
-    cluster, namespaces = get_namespaces(client)
-    load_namespaces(session, cluster, namespaces, update_tag)
-    merge_module_sync_metadata(
-        session,
-        group_type="KubernetesCluster",
-        group_id=cluster["uid"],
-        synced_type="KubernetesCluster",
-        update_tag=update_tag,
-        stat_handler=stat_handler,
-    )
-    return cluster
+def get_namespaces(client: K8sClient) -> list[V1Namespace]:
+    items = k8s_paginate(client.core.list_namespace)
+    return items
 
 
-@timeit
-def get_namespaces(client: K8sClient) -> Tuple[Dict, List[Dict]]:
-    cluster = dict()
-    namespaces = list()
-    for namespace in client.core.list_namespace().items:
-        namespaces.append(
+def transform_namespaces(namespaces: list[V1Namespace]) -> list[dict[str, Any]]:
+    transformed_namespaces = []
+    for namespace in namespaces:
+        transformed_namespaces.append(
             {
                 "uid": namespace.metadata.uid,
                 "name": namespace.metadata.name,
                 "creation_timestamp": get_epoch(namespace.metadata.creation_timestamp),
                 "deletion_timestamp": get_epoch(namespace.metadata.deletion_timestamp),
-            },
+                "status_phase": namespace.status.phase if namespace.status else None,
+            }
         )
-        if namespace.metadata.name == "kube-system":
-            cluster = {"uid": namespace.metadata.uid, "name": client.name}
-    return cluster, namespaces
+    return transformed_namespaces
 
 
 def load_namespaces(
-    session: Session,
-    cluster: Dict,
-    data: List[Dict],
+    session: neo4j.Session,
+    namespaces: list[dict[str, Any]],
     update_tag: int,
+    cluster_name: str,
+    cluster_id: str,
+) -> None:
+    logger.info(f"Loading {len(namespaces)} kubernetes namespaces.")
+    load(
+        session,
+        KubernetesNamespaceSchema(),
+        namespaces,
+        lastupdated=update_tag,
+        cluster_name=cluster_name,
+        CLUSTER_ID=cluster_id,
+    )
+
+
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
 ) -> None:
-    ingestion_cypher_query = """
-    MERGE (cluster:KubernetesCluster {id: $cluster_id})
-    ON CREATE SET cluster.firstseen = timestamp()
-    SET cluster.name = $cluster_name,
-        cluster.lastupdated = $update_tag
-    WITH cluster
-    UNWIND $namespaces as namespace
-        MERGE (space:KubernetesNamespace {id: namespace.uid})
-        ON CREATE SET space.firstseen = timestamp()
-        SET space.lastupdated = $update_tag,
-            space.name = namespace.name,
-            space.created_at = namespace.creation_timestamp,
-            space.deleted_at = namespace.deletion_timestamp
-        WITH cluster, space
-        MERGE (cluster)-[rel1:HAS_NAMESPACE]->(space)
-        ON CREATE SET rel1.firstseen = timestamp()
-        SET rel1.lastupdated = $update_tag
-    """
-    logger.info(f"Loading {len(data)} kubernetes namespaces.")
-    session.run(
-        ingestion_cypher_query,
-        namespaces=data,
-        cluster_id=cluster["uid"],
-        cluster_name=cluster["name"],
-        update_tag=update_tag,
+    logger.debug("Running cleanup job for KubernetesNamespace")
+    cleanup_job = GraphJob.from_node_schema(
+        KubernetesNamespaceSchema(), common_job_parameters
+    )
+    cleanup_job.run(neo4j_session)
+
+
+@timeit
+def sync_namespaces(
+    session: neo4j.Session,
+    client: K8sClient,
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    namespaces = get_namespaces(client)
+    transformed_namespaces = transform_namespaces(namespaces)
+    cluster_id: str = common_job_parameters["CLUSTER_ID"]
+    load_namespaces(
+        session,
+        transformed_namespaces,
+        update_tag,
+        client.name,
+        cluster_id,
     )
+    cleanup(session, common_job_parameters)

cartography/intel/kubernetes/pods.py

@@ -1,39 +1,65 @@
+import json
 import logging
-from typing import Dict
-from typing import List
+from typing import Any
 
-from neo4j import Session
+import neo4j
+from kubernetes.client.models import V1Container
+from kubernetes.client.models import V1Pod
 
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
 from cartography.intel.kubernetes.util import get_epoch
+from cartography.intel.kubernetes.util import k8s_paginate
 from cartography.intel.kubernetes.util import K8sClient
+from cartography.models.kubernetes.containers import KubernetesContainerSchema
+from cartography.models.kubernetes.pods import KubernetesPodSchema
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
 
 
-@timeit
-def sync_pods(
-    session: Session,
-    client: K8sClient,
-    update_tag: int,
-    cluster: Dict,
-) -> List[Dict]:
-    pods = get_pods(client, cluster)
-    load_pods(session, pods, update_tag)
-    return pods
+def _extract_pod_containers(pod: V1Pod) -> dict[str, Any]:
+    pod_containers: list[V1Container] = pod.spec.containers
+    containers = dict()
+    for container in pod_containers:
+        containers[container.name] = {
+            "uid": f"{pod.metadata.uid}-{container.name}",
+            "name": container.name,
+            "image": container.image,
+            "namespace": pod.metadata.namespace,
+            "pod_id": pod.metadata.uid,
+            "imagePullPolicy": container.image_pull_policy,
+        }
 
+        # Extract resource requests and limits
+        if container.resources:
+            if container.resources.requests:
+                containers[container.name]["memory_request"] = (
+                    container.resources.requests.get("memory")
+                )
+                containers[container.name]["cpu_request"] = (
+                    container.resources.requests.get("cpu")
+                )
+            else:
+                containers[container.name]["memory_request"] = None
+                containers[container.name]["cpu_request"] = None
+
+            if container.resources.limits:
+                containers[container.name]["memory_limit"] = (
+                    container.resources.limits.get("memory")
+                )
+                containers[container.name]["cpu_limit"] = (
+                    container.resources.limits.get("cpu")
+                )
+            else:
+                containers[container.name]["memory_limit"] = None
+                containers[container.name]["cpu_limit"] = None
+        else:
+            containers[container.name]["memory_request"] = None
+            containers[container.name]["cpu_request"] = None
+            containers[container.name]["memory_limit"] = None
+            containers[container.name]["cpu_limit"] = None
 
-@timeit
-def get_pods(client: K8sClient, cluster: Dict) -> List[Dict]:
-    pods = list()
-    for pod in client.core.list_pod_for_all_namespaces().items:
-        containers = {}
-        for container in pod.spec.containers:
-            containers[container.name] = {
-                "name": container.name,
-                "image": container.image,
-                "uid": f"{pod.metadata.uid}-{container.name}",
-            }
         if pod.status and pod.status.container_statuses:
             for status in pod.status.container_statuses:
                 if status.name in containers:
@@ -46,14 +72,31 @@ def get_pods(client: K8sClient, cluster: Dict) -> List[Dict]:
                         image_sha = status.image_id.split("@")[1]
                     except IndexError:
                         image_sha = None
-                    containers[status.name]["status"] = {
-                        "image_id": status.image_id,
-                        "image_sha": image_sha,
-                        "ready": status.ready,
-                        "started": status.started,
-                        "state": _state,
-                    }
-        pods.append(
+
+                    containers[status.name]["status_image_id"] = status.image_id
+                    containers[status.name]["status_image_sha"] = image_sha
+                    containers[status.name]["status_ready"] = status.ready
+                    containers[status.name]["status_started"] = status.started
+                    containers[status.name]["status_state"] = _state
+    return containers
+
+
+@timeit
+def get_pods(client: K8sClient) -> list[V1Pod]:
+    items = k8s_paginate(client.core.list_pod_for_all_namespaces)
+    return items
+
+
+def _format_pod_labels(labels: dict[str, str]) -> str:
+    return json.dumps(labels)
+
+
+def transform_pods(pods: list[V1Pod]) -> list[dict[str, Any]]:
+    transformed_pods = []
+
+    for pod in pods:
+        containers = _extract_pod_containers(pod)
+        transformed_pods.append(
             {
                 "uid": pod.metadata.uid,
                 "name": pod.metadata.name,
@@ -62,49 +105,99 @@ def get_pods(client: K8sClient, cluster: Dict) -> List[Dict]:
                 "deletion_timestamp": get_epoch(pod.metadata.deletion_timestamp),
                 "namespace": pod.metadata.namespace,
                 "node": pod.spec.node_name,
-                "cluster_uid": cluster["uid"],
-                "labels": pod.metadata.labels,
+                "labels": _format_pod_labels(pod.metadata.labels),
                 "containers": list(containers.values()),
             },
         )
-    return pods
-
-
-def load_pods(session: Session, data: List[Dict], update_tag: int) -> None:
-    ingestion_cypher_query = """
-    UNWIND $pods as k8pod
-        MERGE (pod:KubernetesPod {id: k8pod.uid})
-        ON CREATE SET pod.firstseen = timestamp()
-        SET pod.lastupdated = $update_tag,
-            pod.name = k8pod.name,
-            pod.status_phase = k8pod.status_phase,
-            pod.created_at = k8pod.creation_timestamp,
-            pod.deleted_at = k8pod.deletion_timestamp
-        WITH pod, k8pod.namespace as ns, k8pod.cluster_uid as cuid, k8pod.containers as k8containers
-        MATCH (cluster:KubernetesCluster {id: cuid})-[:HAS_NAMESPACE]->(space:KubernetesNamespace {name: ns})
-        MERGE (space)-[rel1:HAS_POD]->(pod)
-        ON CREATE SET rel1.firstseen = timestamp()
-        SET rel1.lastupdated = $update_tag
-        WITH pod, space, cluster, k8containers
-        UNWIND k8containers as k8container
-            MERGE (container: KubernetesContainer {id: k8container.uid})
-            ON CREATE SET container.firstseen = timestamp()
-            SET container.image = k8container.image,
-                container.status_image_id = k8container.status.image_id,
-                container.status_image_sha = k8container.status.image_sha,
-                container.status_ready = k8container.status.ready,
-                container.status_started = k8container.status.started,
-                container.status_state = k8container.status.state,
-                container.name = k8container.name,
-                container.lastupdated = $update_tag
-            WITH pod, space, cluster, container
-            MERGE (pod)-[rel2:HAS_CONTAINER]->(container)
-            ON CREATE SET rel2.firstseen = timestamp()
-            SET rel2.lastupdated = $update_tag
-            WITH pod, space, container
-            MERGE (cluster)-[rel3:HAS_POD]->(pod)
-            ON CREATE SET rel3.firstseen = timestamp()
-            SET rel3.lastupdated = $update_tag
-    """
-    logger.info(f"Loading {len(data)} kubernetes pods.")
-    session.run(ingestion_cypher_query, pods=data, update_tag=update_tag)
+    return transformed_pods
+
+
+@timeit
+def load_pods(
+    session: neo4j.Session,
+    pods: list[dict[str, Any]],
+    update_tag: int,
+    cluster_id: str,
+    cluster_name: str,
+) -> None:
+    logger.info(f"Loading {len(pods)} kubernetes pods.")
+    load(
+        session,
+        KubernetesPodSchema(),
+        pods,
+        lastupdated=update_tag,
+        CLUSTER_ID=cluster_id,
+        CLUSTER_NAME=cluster_name,
+    )
+
+
+def transform_containers(pods: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    containers = []
+    for pod in pods:
+        containers.extend(pod.get("containers", []))
+    return containers
+
+
+@timeit
+def load_containers(
+    session: neo4j.Session,
+    containers: list[dict[str, Any]],
+    update_tag: int,
+    cluster_id: str,
+    cluster_name: str,
+) -> None:
+    logger.info(f"Loading {len(containers)} kubernetes containers.")
+    load(
+        session,
+        KubernetesContainerSchema(),
+        containers,
+        lastupdated=update_tag,
+        CLUSTER_ID=cluster_id,
+        CLUSTER_NAME=cluster_name,
+    )
+
+
+@timeit
+def cleanup(session: neo4j.Session, common_job_parameters: dict[str, Any]) -> None:
+    logger.debug("Running cleanup job for KubernetesContainer")
+    cleanup_job = GraphJob.from_node_schema(
+        KubernetesContainerSchema(), common_job_parameters
+    )
+    cleanup_job.run(session)
+
+    logger.debug("Running cleanup job for KubernetesPod")
+    cleanup_job = GraphJob.from_node_schema(
+        KubernetesPodSchema(), common_job_parameters
+    )
+    cleanup_job.run(session)
+
+
+@timeit
+def sync_pods(
+    session: neo4j.Session,
+    client: K8sClient,
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
+) -> list[dict[str, Any]]:
+    pods = get_pods(client)
+
+    transformed_pods = transform_pods(pods)
+    load_pods(
+        session=session,
+        pods=transformed_pods,
+        update_tag=update_tag,
+        cluster_id=common_job_parameters["CLUSTER_ID"],
+        cluster_name=client.name,
+    )
+
+    transformed_containers = transform_containers(transformed_pods)
+    load_containers(
+        session=session,
+        containers=transformed_containers,
+        update_tag=update_tag,
+        cluster_id=common_job_parameters["CLUSTER_ID"],
+        cluster_name=client.name,
+    )
+
+    cleanup(session, common_job_parameters)
+    return transformed_pods