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

cartography/intel/spacelift/util.py

@@ -0,0 +1,122 @@
+"""
+Utility functions for Spacelift GraphQL API interactions.
+"""
+
+import logging
+from typing import Any
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+# Timeout for API calls: (connection timeout, read timeout) in seconds
+_TIMEOUT = (60, 60)
+
+# GraphQL mutation to exchange API key ID and secret for a JWT token
+_TOKEN_EXCHANGE_MUTATION = """
+mutation GetSpaceliftToken($id: ID!, $secret: String!) {
+  apiKeyUser(id: $id, secret: $secret) {
+    jwt
+  }
+}
+"""
+
+
+def call_spacelift_api(
+    session: requests.Session,
+    api_endpoint: str,
+    query: str,
+    variables: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Make a GraphQL query to the Spacelift API.
+    """
+    logger.debug(f"Making GraphQL request to {api_endpoint}")
+
+    # Prepare the GraphQL request payload
+    payload: dict[str, Any] = {"query": query}
+    if variables:
+        payload["variables"] = variables
+
+    # Make the POST request to the GraphQL endpoint
+    response = session.post(
+        api_endpoint,
+        json=payload,
+        timeout=_TIMEOUT,
+    )
+
+    # Raise an exception for HTTP errors (4xx, 5xx)
+    response.raise_for_status()
+
+    # Parse the JSON response
+    result = response.json()
+
+    # Check for GraphQL errors in the response
+    if "errors" in result:
+        error_messages = [
+            error.get("message", "Unknown error") for error in result["errors"]
+        ]
+        error_string = "; ".join(error_messages)
+        raise ValueError(f"GraphQL query failed: {error_string}")
+
+    return result
+
+
+def get_spacelift_token(api_endpoint: str, key_id: str, key_secret: str) -> str:
+    """
+    Exchange Spacelift API key ID and secret for a JWT bearer token.
+
+    This uses the Spacelift GraphQL API's apiKeyUser mutation to obtain a JWT token
+    that can be used for authenticated requests. This is the recommended authentication
+    method as it avoids storing long-lived tokens.
+
+    :param api_endpoint: Spacelift GraphQL API endpoint (e.g., https://your-account.app.spacelift.io/graphql)
+    :param key_id: Spacelift API key ID (26-character ULID)
+    :param key_secret: Spacelift API key secret
+    :return: JWT bearer token string
+    :raises: requests.exceptions.RequestException if the API request fails
+    :raises: ValueError if the response doesn't contain a valid token
+    """
+    logger.info("Exchanging Spacelift API key for JWT token")
+
+    payload = {
+        "query": _TOKEN_EXCHANGE_MUTATION,
+        "variables": {
+            "id": key_id,
+            "secret": key_secret,
+        },
+    }
+
+    try:
+        response = requests.post(
+            api_endpoint,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=_TIMEOUT,
+        )
+        response.raise_for_status()
+
+        data: dict[str, Any] = response.json()
+
+        # Check for GraphQL errors
+        if "errors" in data:
+            error_messages = [
+                error.get("message", "Unknown error") for error in data["errors"]
+            ]
+            error_string = "; ".join(error_messages)
+            raise ValueError(f"Token exchange failed: {error_string}")
+
+        # Extract JWT token from response
+        jwt = data.get("data", {}).get("apiKeyUser", {}).get("jwt")
+        if not jwt:
+            raise ValueError("Token exchange response did not contain a JWT token")
+
+        logger.info("Successfully obtained JWT token from Spacelift API")
+        return jwt
+
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to exchange Spacelift API key for token: {e}")
+        raise
+    except (KeyError, ValueError) as e:
+        logger.error(f"Invalid response from Spacelift token exchange: {e}")
+        raise

cartography/intel/spacelift/workerpools.py

@@ -0,0 +1,131 @@
+import logging
+from typing import Any
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.spacelift.util import call_spacelift_api
+from cartography.models.spacelift.workerpool import SpaceliftWorkerPoolSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+# GraphQL query to fetch all worker pools (direct array, no pagination)
+GET_WORKER_POOLS_QUERY = """
+query {
+    workerPools {
+        id
+        name
+        description
+        space
+    }
+}
+"""
+
+
+@timeit
+def get_worker_pools(
+    session: requests.Session, api_endpoint: str
+) -> list[dict[str, Any]]:
+
+    logger.info("Fetching Spacelift worker pools")
+
+    response = call_spacelift_api(session, api_endpoint, GET_WORKER_POOLS_QUERY)
+    worker_pools_data = response.get("data", {}).get("workerPools", [])
+
+    logger.info(f"Retrieved {len(worker_pools_data)} Spacelift worker pools")
+    return worker_pools_data
+
+
+def transform_worker_pools(
+    worker_pools_data: list[dict[str, Any]], account_id: str
+) -> list[dict[str, Any]]:
+    result: list[dict[str, Any]] = []
+
+    for pool in worker_pools_data:
+        transformed_pool = {
+            "id": pool["id"],
+            "name": pool.get("name"),
+            "description": pool.get("description"),
+            "pool_type": None,  # WorkerPool type doesn't have a 'type' field in the API
+            "space_id": pool.get("space"),
+            "spacelift_account_id": account_id,
+        }
+
+        result.append(transformed_pool)
+
+    logger.info(f"Transformed {len(result)} worker pools")
+    return result
+
+
+def load_worker_pools(
+    neo4j_session: neo4j.Session,
+    worker_pools_data: list[dict[str, Any]],
+    update_tag: int,
+    account_id: str,
+) -> None:
+    """
+    Load Spacelift worker pools data into Neo4j using the data model.
+    """
+    load(
+        neo4j_session,
+        SpaceliftWorkerPoolSchema(),
+        worker_pools_data,
+        lastupdated=update_tag,
+        spacelift_account_id=account_id,
+    )
+
+    logger.info(f"Loaded {len(worker_pools_data)} Spacelift worker pools")
+
+
+@timeit
+def cleanup_worker_pools(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    """
+    Remove stale Spacelift worker pool data from Neo4j.
+    """
+    logger.debug("Running SpaceliftWorkerPool cleanup job")
+    GraphJob.from_node_schema(SpaceliftWorkerPoolSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+
+
+@timeit
+def sync_worker_pools(
+    neo4j_session: neo4j.Session,
+    spacelift_session: requests.Session,
+    api_endpoint: str,
+    account_id: str,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    """
+    Sync Spacelift worker pools data to Neo4j.
+
+    :param neo4j_session: Neo4j session
+    :param spacelift_session: Authenticated requests session for Spacelift API
+    :param api_endpoint: Spacelift GraphQL API endpoint
+    :param account_id: The Spacelift account ID
+    :param common_job_parameters: Common job parameters containing UPDATE_TAG
+    """
+    # 1. GET - Fetch data from Spacelift API
+    worker_pools_raw_data = get_worker_pools(spacelift_session, api_endpoint)
+
+    # 2. TRANSFORM - Shape data for ingestion
+    transformed_worker_pools = transform_worker_pools(worker_pools_raw_data, account_id)
+
+    # 3. LOAD - Ingest to Neo4j using data model
+    load_worker_pools(
+        neo4j_session,
+        transformed_worker_pools,
+        common_job_parameters["UPDATE_TAG"],
+        account_id,
+    )
+
+    # 4. CLEANUP - Remove stale data
+    cleanup_worker_pools(neo4j_session, common_job_parameters)
+
+    logger.info(f"Synced {len(transformed_worker_pools)} Spacelift worker pools")

cartography/intel/spacelift/workers.py

@@ -0,0 +1,128 @@
+import logging
+from typing import Any
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.spacelift.util import call_spacelift_api
+from cartography.models.spacelift.worker import SpaceliftWorkerSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+# GraphQL query to fetch workers nested under workerPools
+# Note: Workers don't have a top-level query, they're nested under workerPools
+# Note: Worker type doesn't have a 'name' field, only 'id' and 'status'
+GET_WORKERS_QUERY = """
+query {
+    workerPools {
+        id
+        workers {
+            id
+            status
+        }
+    }
+}
+"""
+
+
+@timeit
+def get_workers(session: requests.Session, api_endpoint: str) -> list[dict[str, Any]]:
+    """
+    Fetch all Spacelift workers from the API.
+    Workers are nested under workerPools, so we query workerPools and flatten the workers.
+    """
+    logger.info("Fetching Spacelift workers")
+
+    response = call_spacelift_api(session, api_endpoint, GET_WORKERS_QUERY)
+    worker_pools = response.get("data", {}).get("workerPools", [])
+
+    # Flatten workers from all pools and add the pool ID to each worker
+    all_workers = []
+    for pool in worker_pools:
+        pool_id = pool["id"]
+        for worker in pool.get("workers", []):
+            # Add the pool ID to each worker
+            worker["workerPool"] = pool_id
+            all_workers.append(worker)
+
+    logger.info(
+        f"Retrieved {len(all_workers)} Spacelift workers from {len(worker_pools)} worker pools"
+    )
+    return all_workers
+
+
+def transform_workers(
+    workers_data: list[dict[str, Any]], account_id: str
+) -> list[dict[str, Any]]:
+    result: list[dict[str, Any]] = []
+
+    for worker in workers_data:
+        worker_id = worker["id"]
+
+        transformed_worker = {
+            "id": worker_id,
+            "name": worker_id,  # Use ID as name since Worker type doesn't have a name field
+            "status": worker.get("status"),
+            "worker_pool_id": worker.get("workerPool"),
+            "spacelift_account_id": account_id,
+        }
+
+        result.append(transformed_worker)
+
+    logger.info(f"Transformed {len(result)} workers")
+    return result
+
+
+def load_workers(
+    neo4j_session: neo4j.Session,
+    workers_data: list[dict[str, Any]],
+    update_tag: int,
+    account_id: str,
+) -> None:
+
+    load(
+        neo4j_session,
+        SpaceliftWorkerSchema(),
+        workers_data,
+        lastupdated=update_tag,
+        spacelift_account_id=account_id,
+    )
+
+    logger.info(f"Loaded {len(workers_data)} Spacelift workers")
+
+
+@timeit
+def cleanup_workers(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+
+    logger.debug("Running SpaceliftWorker cleanup job")
+    GraphJob.from_node_schema(SpaceliftWorkerSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+
+
+@timeit
+def sync_workers(
+    neo4j_session: neo4j.Session,
+    spacelift_session: requests.Session,
+    api_endpoint: str,
+    account_id: str,
+    common_job_parameters: dict[str, Any],
+) -> None:
+
+    workers_raw_data = get_workers(spacelift_session, api_endpoint)
+    transformed_workers = transform_workers(workers_raw_data, account_id)
+    load_workers(
+        neo4j_session,
+        transformed_workers,
+        common_job_parameters["UPDATE_TAG"],
+        account_id,
+    )
+    cleanup_workers(neo4j_session, common_job_parameters)
+
+    logger.info(f"Synced {len(transformed_workers)} Spacelift workers")

cartography/intel/trivy/__init__.py

@@ -0,0 +1,272 @@
+import json
+import logging
+from typing import Any
+
+import boto3
+from neo4j import Session
+
+from cartography.client.aws import list_accounts
+from cartography.client.aws.ecr import get_ecr_images
+from cartography.config import Config
+from cartography.intel.trivy.scanner import cleanup
+from cartography.intel.trivy.scanner import get_json_files_in_dir
+from cartography.intel.trivy.scanner import get_json_files_in_s3
+from cartography.intel.trivy.scanner import sync_single_image
+from cartography.stats import get_stats_client
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+stat_handler = get_stats_client("trivy.scanner")
+
+
+def _get_scan_targets_and_aliases(
+    neo4j_session: Session,
+    account_ids: list[str] | None = None,
+) -> tuple[set[str], dict[str, str]]:
+    """
+    Return tag URIs and a mapping of digest-qualified URIs to tag URIs.
+    """
+    if not account_ids:
+        aws_accounts = list_accounts(neo4j_session)
+    else:
+        aws_accounts = account_ids
+
+    image_uris: set[str] = set()
+    digest_aliases: dict[str, str] = {}
+
+    for account_id in aws_accounts:
+        for _, _, image_uri, _, digest in get_ecr_images(neo4j_session, account_id):
+            if not image_uri:
+                continue
+            image_uris.add(image_uri)
+            if digest:
+                # repo URI is everything before the trailing ":" (if present)
+                repo_uri = image_uri.rsplit(":", 1)[0]
+                digest_uri = f"{repo_uri}@{digest}"
+                digest_aliases[digest_uri] = image_uri
+
+    return image_uris, digest_aliases
+
+
+@timeit
+def get_scan_targets(
+    neo4j_session: Session,
+    account_ids: list[str] | None = None,
+) -> set[str]:
+    """
+    Return list of ECR images from all accounts in the graph.
+    """
+    image_uris, _ = _get_scan_targets_and_aliases(neo4j_session, account_ids)
+    return image_uris
+
+
+def _prepare_trivy_data(
+    trivy_data: dict[str, Any],
+    image_uris: set[str],
+    digest_aliases: dict[str, str],
+    source: str,
+) -> tuple[dict[str, Any], str] | None:
+    """
+    Determine the tag URI that corresponds to this Trivy payload.
+
+    Returns (trivy_data, display_uri) if the payload can be linked to an image present
+    in the graph; otherwise returns None so the caller can skip ingestion.
+    """
+
+    artifact_name = (trivy_data.get("ArtifactName") or "").strip()
+    metadata = trivy_data.get("Metadata") or {}
+    candidates: list[str] = []
+
+    if artifact_name:
+        candidates.append(artifact_name)
+
+    repo_tags = metadata.get("RepoTags", [])
+    repo_digests = metadata.get("RepoDigests", [])
+    stripped_tags_digests = [item.strip() for item in repo_tags + repo_digests]
+    candidates.extend(stripped_tags_digests)
+
+    display_uri: str | None = None
+
+    for candidate in candidates:
+        if not candidate:
+            continue
+        if candidate in image_uris:
+            display_uri = candidate
+            break
+        alias = digest_aliases.get(candidate)
+        if alias:
+            display_uri = alias
+            break
+
+    if not display_uri:
+        logger.debug(
+            "Skipping Trivy results for %s because no matching image URI was found in the graph",
+            source,
+        )
+        return None
+
+    return trivy_data, display_uri
+
+
+@timeit
+def sync_trivy_aws_ecr_from_s3(
+    neo4j_session: Session,
+    trivy_s3_bucket: str,
+    trivy_s3_prefix: str,
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
+    boto3_session: boto3.Session,
+) -> None:
+    """
+    Sync Trivy scan results from S3 for AWS ECR images.
+
+    Args:
+        neo4j_session: Neo4j session for database operations
+        trivy_s3_bucket: S3 bucket containing scan results
+        trivy_s3_prefix: S3 prefix path containing scan results
+        update_tag: Update tag for tracking
+        common_job_parameters: Common job parameters for cleanup
+        boto3_session: boto3 session for S3 operations
+    """
+    logger.info(
+        f"Using Trivy scan results from s3://{trivy_s3_bucket}/{trivy_s3_prefix}"
+    )
+
+    image_uris, digest_aliases = _get_scan_targets_and_aliases(neo4j_session)
+    json_files: set[str] = get_json_files_in_s3(
+        trivy_s3_bucket, trivy_s3_prefix, boto3_session
+    )
+
+    if len(json_files) == 0:
+        logger.error(
+            f"Trivy sync was configured, but there are no ECR images with S3 json scan results in bucket "
+            f"'{trivy_s3_bucket}' with prefix '{trivy_s3_prefix}'. "
+            "Skipping Trivy sync to avoid potential data loss. "
+            "Please check the S3 bucket and prefix configuration. We expect the json files in s3 to be named "
+            f"`<image_uri>.json` and to be in the same bucket and prefix as the scan results. If the prefix is "
+            "a folder, it MUST end with a trailing slash '/'. "
+        )
+        raise ValueError("No ECR images with S3 json scan results found.")
+
+    logger.info(f"Processing {len(json_files)} Trivy result files from S3")
+    s3_client = boto3_session.client("s3")
+    for s3_object_key in json_files:
+        logger.debug(
+            f"Reading scan results from S3: s3://{trivy_s3_bucket}/{s3_object_key}"
+        )
+        response = s3_client.get_object(Bucket=trivy_s3_bucket, Key=s3_object_key)
+        scan_data_json = response["Body"].read().decode("utf-8")
+        trivy_data = json.loads(scan_data_json)
+
+        prepared = _prepare_trivy_data(
+            trivy_data,
+            image_uris=image_uris,
+            digest_aliases=digest_aliases,
+            source=f"s3://{trivy_s3_bucket}/{s3_object_key}",
+        )
+        if prepared is None:
+            continue
+
+        prepared_data, display_uri = prepared
+        sync_single_image(
+            neo4j_session,
+            prepared_data,
+            display_uri,
+            update_tag,
+        )
+
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def sync_trivy_aws_ecr_from_dir(
+    neo4j_session: Session,
+    results_dir: str,
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    """Sync Trivy scan results from local files for AWS ECR images."""
+    logger.info(f"Using Trivy scan results from {results_dir}")
+
+    image_uris, digest_aliases = _get_scan_targets_and_aliases(neo4j_session)
+    json_files: set[str] = get_json_files_in_dir(results_dir)
+
+    if not json_files:
+        logger.error(
+            f"Trivy sync was configured, but no json files were found in {results_dir}."
+        )
+        raise ValueError("No Trivy json results found on disk")
+
+    logger.info(f"Processing {len(json_files)} local Trivy result files")
+
+    for file_path in json_files:
+        try:
+            with open(file_path, encoding="utf-8") as f:
+                trivy_data = json.load(f)
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to read Trivy data from {file_path}: {e}")
+            continue
+
+        prepared = _prepare_trivy_data(
+            trivy_data,
+            image_uris=image_uris,
+            digest_aliases=digest_aliases,
+            source=file_path,
+        )
+        if prepared is None:
+            continue
+
+        prepared_data, display_uri = prepared
+        sync_single_image(
+            neo4j_session,
+            prepared_data,
+            display_uri,
+            update_tag,
+        )
+
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def start_trivy_ingestion(neo4j_session: Session, config: Config) -> None:
+    """Start Trivy scan ingestion from S3 or local files.
+
+    Args:
+        neo4j_session: Neo4j session for database operations
+        config: Configuration object containing S3 or directory paths
+    """
+    if not config.trivy_s3_bucket and not config.trivy_results_dir:
+        logger.info("Trivy configuration not provided. Skipping Trivy ingestion.")
+        return
+
+    if config.trivy_results_dir:
+        common_job_parameters = {
+            "UPDATE_TAG": config.update_tag,
+        }
+        sync_trivy_aws_ecr_from_dir(
+            neo4j_session,
+            config.trivy_results_dir,
+            config.update_tag,
+            common_job_parameters,
+        )
+        return
+
+    if config.trivy_s3_prefix is None:
+        config.trivy_s3_prefix = ""
+
+    common_job_parameters = {
+        "UPDATE_TAG": config.update_tag,
+    }
+
+    boto3_session = boto3.Session()
+
+    sync_trivy_aws_ecr_from_s3(
+        neo4j_session,
+        config.trivy_s3_bucket,
+        config.trivy_s3_prefix,
+        config.update_tag,
+        common_job_parameters,
+        boto3_session,
+    )
+
+    # Support other Trivy resource types here e.g. if Google Cloud has images.