feat(api): make Attack Paths sink selectable between Neo4j and Neptune (#11524)

This commit is contained in:
Josema Camacho authored and GitHub committed 2026-06-26 10:22:29 +02:00
1 parent 9b8b77cec0
commit 5793cd7e38
48 files changed
+9928 -3210

No files matched your search

+24 -8
View File
@@ -6,6 +6,7 @@ from typing import Any
import aioboto3
import boto3
import botocore
import neo4j
from api.models import (
AttackPathsScan as ProwlerAPIAttackPathsScan,
@@ -73,13 +74,28 @@ def start_aws_ingestion(
# Adding an extra field
common_job_parameters["AWS_ID"] = prowler_api_provider.uid
cartography_aws._autodiscover_accounts(
neo4j_session,
boto3_session,
prowler_api_provider.uid,
cartography_config.update_tag,
common_job_parameters,
)
# AWS Organizations account autodiscovery. Inlined from Cartography's removed
# `_autodiscover_accounts` (deleted in `0.137.0`), as `load_aws_accounts` is still public.
try:
org_client = boto3_session.client("organizations")
paginator = org_client.get_paginator("list_accounts")
discovered = []
for page in paginator.paginate():
discovered.extend(page["Accounts"])
active_accounts = {
a["Name"]: a["Id"] for a in discovered if a["Status"] == "ACTIVE"
}
cartography_aws.organizations.load_aws_accounts(
neo4j_session,
active_accounts,
cartography_config.update_tag,
common_job_parameters,
)
except botocore.exceptions.ClientError:
logger.warning(
f"Account {prowler_api_provider.uid} lacks permissions for AWS "
"Organizations autodiscovery."
)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 4)
failed_syncs = sync_aws_account(
@@ -277,7 +293,7 @@ def sync_aws_account(
sync_args: dict[str, Any],
attack_paths_scan: ProwlerAPIAttackPathsScan,
) -> dict[str, str]:
current_progress = 4 # `cartography_aws._autodiscover_accounts`
current_progress = 4 # AWS Organizations account autodiscovery
max_progress = (
87 # `cartography_aws.RESOURCE_FUNCTIONS["permission_relationships"]` - 1
)
@@ -8,7 +8,7 @@ from celery import states
from celery.utils.log import get_task_logger
from config.django.base import ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES
from tasks.jobs.attack_paths.db_utils import (
_mark_scan_finished,
mark_scan_finished,
recover_graph_data_ready,
)
from tasks.jobs.orphan_recovery import is_worker_alive as _is_worker_alive
@@ -87,7 +87,7 @@ def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]:
else:
reason = "Worker dead — cleaned up by periodic task"
else:
# No worker recorded — time-based heuristic only
# No worker recorded, time-based heuristic only
if scan.started_at and scan.started_at >= cutoff:
continue
reason = (
@@ -160,7 +160,7 @@ def _cleanup_scan(scan, task_result, reason: str) -> bool:
"""
scan_id_str = str(scan.id)
# 1. Drop temp Neo4j database
# Drop temp Neo4j database
tmp_db_name = graph_database.get_database_name(scan.id, temporary=True)
try:
graph_database.drop_database(tmp_db_name)
@@ -225,6 +225,6 @@ def _finalize_failed_scan(scan, expected_state: str, reason: str):
logger.info(f"Scan {scan_id_str} is now {fresh_scan.state}, skipping")
return None
_mark_scan_finished(fresh_scan, StateChoices.FAILED, {"global_error": reason})
mark_scan_finished(fresh_scan, StateChoices.FAILED, {"global_error": reason})
return fresh_scan
@@ -1,9 +1,14 @@
from collections.abc import Callable
from dataclasses import dataclass
from uuid import UUID
from config.env import env
from tasks.jobs.attack_paths import aws
from tasks.jobs.attack_paths import provider_config as _provider_config
# Re-export provider config objects so existing imports keep working.
AWS_CONFIG = _provider_config.AWS_CONFIG
NormalizedList = _provider_config.NormalizedList
PROVIDER_CONFIGS = _provider_config.PROVIDER_CONFIGS
ProviderConfig = _provider_config.ProviderConfig
# Batch size for Neo4j write operations (resource labeling, cleanup)
BATCH_SIZE = env.int("ATTACK_PATHS_BATCH_SIZE", 1000)
@@ -21,42 +26,12 @@ PROWLER_FINDING_LABEL = "ProwlerFinding"
PROVIDER_RESOURCE_LABEL = "_ProviderResource"
# Dynamic isolation labels that contain entity UUIDs and are added to every synced node during sync
# Format: _Tenant_{uuid_no_hyphens}, _Provider_{uuid_no_hyphens}
# Format: `_Tenant_{uuid_no_hyphens}`, `_Provider_{uuid_no_hyphens}`
TENANT_LABEL_PREFIX = "_Tenant_"
PROVIDER_LABEL_PREFIX = "_Provider_"
DYNAMIC_ISOLATION_PREFIXES = [TENANT_LABEL_PREFIX, PROVIDER_LABEL_PREFIX]
@dataclass(frozen=True)
class ProviderConfig:
"""Configuration for a cloud provider's Attack Paths integration."""
name: str
root_node_label: str # e.g., "AWSAccount"
uid_field: str # e.g., "arn"
# Label for resources connected to the account node, enabling indexed finding lookups.
resource_label: str # e.g., "_AWSResource"
ingestion_function: Callable
# Maps a Postgres resource UID (e.g. full ARN) to the short-id form Cartography stores on some node types (e.g. `i-xxx` for EC2Instance).
short_uid_extractor: Callable[[str], str]
# Provider Configurations
# -----------------------
AWS_CONFIG = ProviderConfig(
name="aws",
root_node_label="AWSAccount",
uid_field="arn",
resource_label="_AWSResource",
ingestion_function=aws.start_aws_ingestion,
short_uid_extractor=aws.extract_short_uid,
)
PROVIDER_CONFIGS: dict[str, ProviderConfig] = {
"aws": AWS_CONFIG,
}
# Labels added by Prowler that should be filtered from API responses
# Derived from provider configs + common internal labels
INTERNAL_LABELS: list[str] = [
@@ -87,7 +62,6 @@ INTERNAL_PROPERTIES: list[str] = [
# Provider Config Accessors
# -------------------------
def is_provider_available(provider_type: str) -> bool:
@@ -135,7 +109,6 @@ def get_short_uid_extractor(provider_type: str) -> Callable[[str], str]:
# Dynamic Isolation Label Helpers
# --------------------------------
def _normalize_uuid(value: str | UUID) -> str:
@@ -8,6 +8,8 @@ from api.models import Provider as ProwlerAPIProvider
from api.models import StateChoices
from cartography.config import Config as CartographyConfig
from celery.utils.log import get_task_logger
from django.conf import settings
from django.db.models import Case, IntegerField, Value, When
from tasks.jobs.attack_paths.config import is_provider_available
logger = get_task_logger(__name__)
@@ -29,13 +31,33 @@ def create_attack_paths_scan(
return None
with rls_transaction(tenant_id):
# Inherit graph_data_ready from the previous scan for this provider,
# so queries remain available while the new scan runs.
previous_data_ready = ProwlerAPIAttackPathsScan.objects.filter(
tenant_id=tenant_id,
provider_id=provider_id,
graph_data_ready=True,
).exists()
# Inherit metadata from the previous ready scan for this provider so
# queries remain available while the new scan runs. The new row only
# flips to the target sink after its own graph sync succeeds.
active_sink_backend = settings.ATTACK_PATHS_SINK_DATABASE
previous_ready = (
ProwlerAPIAttackPathsScan.objects.filter(
tenant_id=tenant_id,
provider_id=provider_id,
graph_data_ready=True,
)
.annotate(
active_sink_rank=Case(
When(sink_backend=active_sink_backend, then=Value(0)),
default=Value(1),
output_field=IntegerField(),
)
)
.order_by("active_sink_rank", "-inserted_at")
.first()
)
previous_data_ready = previous_ready is not None
inherited_is_migrated = previous_ready.is_migrated if previous_ready else False
inherited_sink_backend = (
previous_ready.sink_backend
if previous_ready
else ProwlerAPIAttackPathsScan.SinkBackendChoices.NEO4J
)
attack_paths_scan = ProwlerAPIAttackPathsScan.objects.create(
tenant_id=tenant_id,
@@ -44,6 +66,8 @@ def create_attack_paths_scan(
state=StateChoices.SCHEDULED,
started_at=datetime.now(tz=UTC),
graph_data_ready=previous_data_ready,
is_migrated=inherited_is_migrated,
sink_backend=inherited_sink_backend,
)
attack_paths_scan.save()
@@ -114,7 +138,7 @@ def starting_attack_paths_scan(
return True
def _mark_scan_finished(
def mark_scan_finished(
attack_paths_scan: ProwlerAPIAttackPathsScan,
state: StateChoices,
ingestion_exceptions: dict[str, Any],
@@ -148,7 +172,7 @@ def finish_attack_paths_scan(
ingestion_exceptions: dict[str, Any],
) -> None:
with rls_transaction(attack_paths_scan.tenant_id):
_mark_scan_finished(attack_paths_scan, state, ingestion_exceptions)
mark_scan_finished(attack_paths_scan, state, ingestion_exceptions)
def update_attack_paths_scan_progress(
@@ -169,19 +193,45 @@ def set_graph_data_ready(
attack_paths_scan.save(update_fields=["graph_data_ready"])
def set_scan_migrated(
attack_paths_scan: ProwlerAPIAttackPathsScan,
migrated: bool,
sink_backend: str | None = None,
) -> None:
"""Mark the scan as written with the current (migrated) schema.
Called after a successful sync so the read catalog and sink backend only
switch once the new graph is actually live.
# TODO: drop after Neptune cutover
"""
with rls_transaction(attack_paths_scan.tenant_id):
attack_paths_scan.is_migrated = migrated
update_fields = ["is_migrated"]
if sink_backend is not None:
attack_paths_scan.sink_backend = sink_backend
update_fields.append("sink_backend")
attack_paths_scan.save(update_fields=update_fields)
def set_provider_graph_data_ready(
attack_paths_scan: ProwlerAPIAttackPathsScan,
ready: bool,
sink_backend: str | None = None,
) -> None:
"""
Set `graph_data_ready` for ALL scans of the same provider.
Set `graph_data_ready` for scans of the same provider in one sink.
Used before drop/sync so that older scan IDs cannot bypass the query gate while the graph is being replaced.
Used before drop/sync so that older scan IDs in the target sink cannot
bypass the query gate while that sink's graph is being replaced. Scans
preserved in another sink stay queryable for rollback.
"""
target_sink_backend = sink_backend or attack_paths_scan.sink_backend
with rls_transaction(attack_paths_scan.tenant_id):
ProwlerAPIAttackPathsScan.objects.filter(
tenant_id=attack_paths_scan.tenant_id,
provider_id=attack_paths_scan.provider_id,
sink_backend=target_sink_backend,
).update(graph_data_ready=ready)
attack_paths_scan.refresh_from_db(fields=["graph_data_ready"])
@@ -202,10 +252,15 @@ def recover_graph_data_ready(
next successful scan) is a worse outcome for the user.
"""
try:
from api.attack_paths import sink as sink_module
tenant_db = graph_database.get_database_name(attack_paths_scan.tenant_id)
if graph_database.has_provider_data(
tenant_db, str(attack_paths_scan.provider_id)
):
# TODO: drop after Neptune cutover
# Check the backend that actually holds this scan's data, not the
# currently configured sink, a stale `EXECUTING` scan from before a
# backend switch must still be recoverable
backend = sink_module.get_backend_for_scan(attack_paths_scan)
if backend.has_provider_data(tenant_db, str(attack_paths_scan.provider_id)):
set_provider_graph_data_ready(attack_paths_scan, True)
logger.info(
f"Recovered `graph_data_ready` for provider {attack_paths_scan.provider_id}"
@@ -247,6 +302,6 @@ def fail_attack_paths_scan(
return
if fresh.state in (StateChoices.COMPLETED, StateChoices.FAILED):
return
_mark_scan_finished(fresh, StateChoices.FAILED, {"global_error": error})
mark_scan_finished(fresh, StateChoices.FAILED, {"global_error": error})
recover_graph_data_ready(fresh)
@@ -82,7 +82,6 @@ def _to_neo4j_dict(
# Public API
# ----------
def analysis(
@@ -196,7 +195,6 @@ def load_findings(
# Findings Streaming (Generator-based)
# -------------------------------------
def stream_findings_with_resources(
@@ -275,7 +273,6 @@ def _fetch_findings_batch(
# Batch Enrichment
# -----------------
def _enrich_batch_with_resources(
@@ -1,5 +1,6 @@
import neo4j
from cartography.client.core.tx import run_write_query
from cartography.intel import create_indexes as cartography_create_indexes
from celery.utils.log import get_task_logger
from tasks.jobs.attack_paths.config import (
INTERNET_NODE_LABEL,
@@ -30,14 +31,34 @@ SYNC_INDEX_STATEMENTS = [
def create_findings_indexes(neo4j_session: neo4j.Session) -> None:
"""Create indexes for Prowler findings and resource lookups."""
"""Create indexes for Prowler findings and resource lookups.
Runs `CREATE INDEX`, so the caller must only invoke this against a Neo4j
session (the temp ingest DB or a Neo4j sink). Neptune auto-manages indexes
and rejects `CREATE INDEX`, so callers skip it for the Neptune sink.
"""
logger.info("Creating indexes for Prowler Findings node types")
for statement in FINDINGS_INDEX_STATEMENTS:
run_write_query(neo4j_session, statement)
def create_cartography_indexes(neo4j_session: neo4j.Session, config) -> None:
"""Create Cartography's standard indexes for the session's database.
Runs `CREATE INDEX`, so the caller must only invoke this against a Neo4j
session (the temp ingest DB or a Neo4j sink). Neptune auto-manages indexes
and rejects `CREATE INDEX`, so callers skip it for the Neptune sink.
"""
cartography_create_indexes.run(neo4j_session, config)
def create_sync_indexes(neo4j_session: neo4j.Session) -> None:
"""Create indexes for provider resource sync operations."""
"""Create indexes for provider resource sync operations.
Runs `CREATE INDEX`, so the caller must only invoke this against a Neo4j
session (the temp ingest DB or a Neo4j sink). Neptune auto-manages indexes
and rejects `CREATE INDEX`, so callers skip it for the Neptune sink.
"""
logger.info("Ensuring ProviderResource indexes exist")
for statement in SYNC_INDEX_STATEMENTS:
neo4j_session.run(statement)
@@ -0,0 +1,413 @@
"""
Provider-level Attack Paths configuration.
Each `ProviderConfig` carries the cloud provider's ingestion entry point and
the catalog of list-typed node properties (`normalized_lists`). The sync
layer reads this catalog and materialises each list element as a child node
connected to the parent by a typed edge, so queries traverse the graph
instead of working on serialised list values. Both Neo4j and Neptune sinks
write the same shape and queries are portable across them.
"""
from collections.abc import Callable
from dataclasses import dataclass, field
from tasks.jobs.attack_paths import aws
@dataclass(frozen=True)
class NormalizedList:
"""Catalog entry for a list-typed node property.
Describes how the sync layer materialises a parent node's list-typed
property as a set of child item nodes connected by a typed edge.
Conventions (mechanical, do not invent):
- `child_label`: `<SourceLabel><PropertyPascal>Item`
e.g. AWSPolicyStatement.resource -> AWSPolicyStatementResourceItem
- `rel_type`: `HAS_<PROPERTY_UPPER>`
e.g. resource -> HAS_RESOURCE
- child node property:
* `field_map = []` (scalar list, ~95% case) -> child stores `value: str`
* `field_map = [(src_key, child_field), ...]` (list of dicts, rare)
-> child stores those fields
"""
source_label: str
source_property: str
child_label: str
rel_type: str
field_map: list[tuple[str, str]] = field(default_factory=list)
def __post_init__(self) -> None:
if self.field_map:
child_fields = [dst for _, dst in self.field_map]
if "value" in child_fields:
raise ValueError(
f"NormalizedList {self.source_label}.{self.source_property}: "
"`value` is reserved for scalar mode; do not map a source key to it"
)
src_keys = [src for src, _ in self.field_map]
if len(set(src_keys)) != len(src_keys):
raise ValueError(
f"NormalizedList {self.source_label}.{self.source_property}: "
"duplicate source key in field_map"
)
if len(set(child_fields)) != len(child_fields):
raise ValueError(
f"NormalizedList {self.source_label}.{self.source_property}: "
"duplicate child field in field_map"
)
@dataclass(frozen=True)
class ProviderConfig:
"""Configuration for a cloud provider's Attack Paths integration."""
name: str
root_node_label: str # e.g., "AWSAccount"
uid_field: str # e.g., "arn"
# Label for resources connected to the account node, enabling indexed finding lookups
resource_label: str # e.g., "_AWSResource"
ingestion_function: Callable
# Maps a Postgres resource UID (e.g. full ARN) to the short-id form Cartography stores on some node types (e.g. `i-xxx` for EC2Instance)
short_uid_extractor: Callable[[str], str]
# List-typed properties to materialise as child nodes + edges at sync time.
# Mandatory (may be []). Without an entry here, a list-typed property falls
# back to comma-string flatten and emits a one-time warning.
normalized_lists: list[NormalizedList]
# AWS list-typed property catalog.
# One entry per Cartography node property whose runtime value is a list. The
# sync layer materialises each element as a `<child_label>` node and links it
# to the parent with a `<rel_type>` edge; see the `NormalizedList` docstring
# above for the naming conventions.
AWS_NORMALIZED_LISTS: list[NormalizedList] = [
# AWSPolicyStatement - the hot path driving the 53-query perf fix.
NormalizedList(
"AWSPolicyStatement", "action", "AWSPolicyStatementActionItem", "HAS_ACTION"
),
NormalizedList(
"AWSPolicyStatement",
"notaction",
"AWSPolicyStatementNotactionItem",
"HAS_NOTACTION",
),
NormalizedList(
"AWSPolicyStatement",
"resource",
"AWSPolicyStatementResourceItem",
"HAS_RESOURCE",
),
NormalizedList(
"AWSPolicyStatement",
"notresource",
"AWSPolicyStatementNotresourceItem",
"HAS_NOTRESOURCE",
),
# S3PolicyStatement - same shape as IAM policies; AWS allows list or string.
NormalizedList(
"S3PolicyStatement", "action", "S3PolicyStatementActionItem", "HAS_ACTION"
),
NormalizedList(
"S3PolicyStatement", "resource", "S3PolicyStatementResourceItem", "HAS_RESOURCE"
),
# IAM / Cognito / KMS / Secrets
NormalizedList(
"CognitoIdentityPool", "roles", "CognitoIdentityPoolRolesItem", "HAS_ROLES"
),
NormalizedList(
"KMSKey",
"encryption_algorithms",
"KMSKeyEncryptionAlgorithmsItem",
"HAS_ENCRYPTION_ALGORITHMS",
),
NormalizedList(
"KMSKey",
"signing_algorithms",
"KMSKeySigningAlgorithmsItem",
"HAS_SIGNING_ALGORITHMS",
),
NormalizedList(
"KMSKey",
"anonymous_actions",
"KMSKeyAnonymousActionsItem",
"HAS_ANONYMOUS_ACTIONS",
),
NormalizedList(
"KMSGrant", "operations", "KMSGrantOperationsItem", "HAS_OPERATIONS"
),
NormalizedList(
"SecretsManagerSecretVersion",
"version_stages",
"SecretsManagerSecretVersionVersionStagesItem",
"HAS_VERSION_STAGES",
),
NormalizedList(
"SecretsManagerSecretVersion",
"kms_key_ids",
"SecretsManagerSecretVersionKmsKeyIdsItem",
"HAS_KMS_KEY_IDS",
),
NormalizedList(
"SecretsManagerSecretVersion",
"tags",
"SecretsManagerSecretVersionTagsItem",
"HAS_TAGS",
field_map=[("Key", "key"), ("Value", "value_")],
# `value` is reserved for scalar mode; map `Value` to `value_` to keep dict shape.
),
# Lambda / Compute
NormalizedList(
"AWSLambda", "architectures", "AWSLambdaArchitecturesItem", "HAS_ARCHITECTURES"
),
NormalizedList(
"AWSLambda",
"anonymous_actions",
"AWSLambdaAnonymousActionsItem",
"HAS_ANONYMOUS_ACTIONS",
),
NormalizedList(
"CodeBuildProject",
"environment_variables",
"CodeBuildProjectEnvironmentVariablesItem",
"HAS_ENVIRONMENT_VARIABLES",
),
# ECS family
NormalizedList(
"ECSCluster",
"capacity_providers",
"ECSClusterCapacityProvidersItem",
"HAS_CAPACITY_PROVIDERS",
),
NormalizedList(
"ECSTaskDefinition",
"compatibilities",
"ECSTaskDefinitionCompatibilitiesItem",
"HAS_COMPATIBILITIES",
),
NormalizedList(
"ECSTaskDefinition",
"requires_compatibilities",
"ECSTaskDefinitionRequiresCompatibilitiesItem",
"HAS_REQUIRES_COMPATIBILITIES",
),
NormalizedList(
"ECSContainerDefinition",
"links",
"ECSContainerDefinitionLinksItem",
"HAS_LINKS",
),
NormalizedList(
"ECSContainerDefinition",
"entry_point",
"ECSContainerDefinitionEntryPointItem",
"HAS_ENTRY_POINT",
),
NormalizedList(
"ECSContainerDefinition",
"command",
"ECSContainerDefinitionCommandItem",
"HAS_COMMAND",
),
NormalizedList(
"ECSContainerDefinition",
"dns_servers",
"ECSContainerDefinitionDnsServersItem",
"HAS_DNS_SERVERS",
),
NormalizedList(
"ECSContainerDefinition",
"dns_search_domains",
"ECSContainerDefinitionDnsSearchDomainsItem",
"HAS_DNS_SEARCH_DOMAINS",
),
NormalizedList(
"ECSContainerDefinition",
"docker_security_options",
"ECSContainerDefinitionDockerSecurityOptionsItem",
"HAS_DOCKER_SECURITY_OPTIONS",
),
NormalizedList("ECSContainer", "gpu_ids", "ECSContainerGpuIdsItem", "HAS_GPU_IDS"),
# ECR
NormalizedList(
"ECRImage", "layer_diff_ids", "ECRImageLayerDiffIdsItem", "HAS_LAYER_DIFF_IDS"
),
NormalizedList(
"ECRImage",
"child_image_digests",
"ECRImageChildImageDigestsItem",
"HAS_CHILD_IMAGE_DIGESTS",
),
# EC2 / Networking
NormalizedList(
"EC2Instance",
"exposed_internet_type",
"EC2InstanceExposedInternetTypeItem",
"HAS_EXPOSED_INTERNET_TYPE",
),
NormalizedList(
"AutoScalingGroup",
"exposed_internet_type",
"AutoScalingGroupExposedInternetTypeItem",
"HAS_EXPOSED_INTERNET_TYPE",
),
NormalizedList(
"LaunchConfiguration",
"security_groups",
"LaunchConfigurationSecurityGroupsItem",
"HAS_SECURITY_GROUPS",
),
NormalizedList(
"LaunchTemplateVersion",
"security_group_ids",
"LaunchTemplateVersionSecurityGroupIdsItem",
"HAS_SECURITY_GROUP_IDS",
),
NormalizedList(
"LaunchTemplateVersion",
"security_groups",
"LaunchTemplateVersionSecurityGroupsItem",
"HAS_SECURITY_GROUPS",
),
NormalizedList(
"ELBListener", "policy_names", "ELBListenerPolicyNamesItem", "HAS_POLICY_NAMES"
),
# CloudFront / Route53 / CloudWatch / CloudTrail
NormalizedList(
"CloudFrontDistribution",
"aliases",
"CloudFrontDistributionAliasesItem",
"HAS_ALIASES",
),
NormalizedList(
"CloudFrontDistribution",
"geo_restriction_locations",
"CloudFrontDistributionGeoRestrictionLocationsItem",
"HAS_GEO_RESTRICTION_LOCATIONS",
),
NormalizedList(
"CloudWatchLogGroup",
"inherited_properties",
"CloudWatchLogGroupInheritedPropertiesItem",
"HAS_INHERITED_PROPERTIES",
),
# RDS / Storage
NormalizedList(
"RDSCluster",
"availability_zones",
"RDSClusterAvailabilityZonesItem",
"HAS_AVAILABILITY_ZONES",
),
NormalizedList(
"RDSEventSubscription",
"event_categories",
"RDSEventSubscriptionEventCategoriesItem",
"HAS_EVENT_CATEGORIES",
),
NormalizedList(
"RDSEventSubscription",
"source_ids",
"RDSEventSubscriptionSourceIdsItem",
"HAS_SOURCE_IDS",
),
NormalizedList(
"S3Bucket",
"anonymous_actions",
"S3BucketAnonymousActionsItem",
"HAS_ANONYMOUS_ACTIONS",
),
# Inspector / Config / SSM / ACM / APIGateway / Glue / SageMaker / Bedrock
NormalizedList(
"AWSInspectorFinding",
"referenceurls",
"AWSInspectorFindingReferenceurlsItem",
"HAS_REFERENCEURLS",
),
NormalizedList(
"AWSInspectorFinding",
"relatedvulnerabilities",
"AWSInspectorFindingRelatedvulnerabilitiesItem",
"HAS_RELATEDVULNERABILITIES",
),
NormalizedList(
"AWSInspectorFinding",
"vulnerablepackageids",
"AWSInspectorFindingVulnerablepackageidsItem",
"HAS_VULNERABLEPACKAGEIDS",
),
NormalizedList(
"AWSConfigurationRecorder",
"recording_group_resource_types",
"AWSConfigurationRecorderRecordingGroupResourceTypesItem",
"HAS_RECORDING_GROUP_RESOURCE_TYPES",
),
NormalizedList(
"AWSConfigRule",
"scope_compliance_resource_types",
"AWSConfigRuleScopeComplianceResourceTypesItem",
"HAS_SCOPE_COMPLIANCE_RESOURCE_TYPES",
),
NormalizedList(
"AWSConfigRule",
"source_details",
"AWSConfigRuleSourceDetailsItem",
"HAS_SOURCE_DETAILS",
),
NormalizedList(
"SSMInstancePatch", "cve_ids", "SSMInstancePatchCveIdsItem", "HAS_CVE_IDS"
),
NormalizedList(
"ACMCertificate", "in_use_by", "ACMCertificateInUseByItem", "HAS_IN_USE_BY"
),
NormalizedList(
"APIGatewayRestAPI",
"anonymous_actions",
"APIGatewayRestAPIAnonymousActionsItem",
"HAS_ANONYMOUS_ACTIONS",
),
NormalizedList(
"GlueJob", "connections", "GlueJobConnectionsItem", "HAS_CONNECTIONS"
),
NormalizedList(
"AWSBedrockFoundationModel",
"input_modalities",
"AWSBedrockFoundationModelInputModalitiesItem",
"HAS_INPUT_MODALITIES",
),
NormalizedList(
"AWSBedrockFoundationModel",
"output_modalities",
"AWSBedrockFoundationModelOutputModalitiesItem",
"HAS_OUTPUT_MODALITIES",
),
NormalizedList(
"AWSBedrockFoundationModel",
"customizations_supported",
"AWSBedrockFoundationModelCustomizationsSupportedItem",
"HAS_CUSTOMIZATIONS_SUPPORTED",
),
NormalizedList(
"AWSBedrockFoundationModel",
"inference_types_supported",
"AWSBedrockFoundationModelInferenceTypesSupportedItem",
"HAS_INFERENCE_TYPES_SUPPORTED",
),
]
AWS_CONFIG = ProviderConfig(
name="aws",
root_node_label="AWSAccount",
uid_field="arn",
resource_label="_AWSResource",
ingestion_function=aws.start_aws_ingestion,
short_uid_extractor=aws.extract_short_uid,
normalized_lists=AWS_NORMALIZED_LISTS,
)
PROVIDER_CONFIGS: dict[str, ProviderConfig] = {
"aws": AWS_CONFIG,
}
@@ -1,8 +1,6 @@
# Cypher query templates for Attack Paths operations
from tasks.jobs.attack_paths.config import (
INTERNET_NODE_LABEL,
PROVIDER_ELEMENT_ID_PROPERTY,
PROVIDER_RESOURCE_LABEL,
PROWLER_FINDING_LABEL,
)
@@ -21,7 +19,6 @@ def render_cypher_template(template: str, replacements: dict[str, str]) -> str:
# Findings queries (used by findings.py)
# ---------------------------------------
ADD_RESOURCE_LABEL_TEMPLATE = """
MATCH (account:__ROOT_LABEL__ {id: $provider_uid})-->(r)
@@ -88,7 +85,6 @@ INSERT_FINDING_TEMPLATE = f"""
"""
# Internet queries (used by internet.py)
# ---------------------------------------
CREATE_INTERNET_NODE = f"""
MERGE (internet:{INTERNET_NODE_LABEL} {{id: 'Internet'}})
@@ -118,8 +114,8 @@ CREATE_CAN_ACCESS_RELATIONSHIPS_TEMPLATE = f"""
RETURN COUNT(r) AS relationships_merged
"""
# Sync queries (used by sync.py)
# -------------------------------
# Sync queries (used by sync.py to fetch from the cartography temp Neo4j DB)
# The write side of sync lives in each sink (`api/attack_paths/sink/`).
NODE_FETCH_QUERY = """
MATCH (n)
@@ -143,17 +139,3 @@ RELATIONSHIPS_FETCH_QUERY = """
ORDER BY internal_id
LIMIT $batch_size
"""
NODE_SYNC_TEMPLATE = f"""
UNWIND $rows AS row
MERGE (n:__NODE_LABELS__ {{{PROVIDER_ELEMENT_ID_PROPERTY}: row.provider_element_id}})
SET n += row.props
"""
RELATIONSHIP_SYNC_TEMPLATE = f"""
UNWIND $rows AS row
MATCH (s:{PROVIDER_RESOURCE_LABEL} {{{PROVIDER_ELEMENT_ID_PROPERTY}: row.start_element_id}})
MATCH (t:{PROVIDER_RESOURCE_LABEL} {{{PROVIDER_ELEMENT_ID_PROPERTY}: row.end_element_id}})
MERGE (s)-[r:__REL_TYPE__ {{{PROVIDER_ELEMENT_ID_PROPERTY}: row.provider_element_id}}]->(t)
SET r += row.props
"""
+91 -32
View File
@@ -39,8 +39,8 @@ Pipeline steps:
7. Sync the temp database into the tenant database:
- Drop the old provider subgraph (matched by dynamic _Provider_{uuid} label).
graph_data_ready is set to False for all scans of this provider while
the swap happens so the API doesn't serve partial data.
graph_data_ready is set to False for scans of this provider in the
target sink while the swap happens so the API doesn't serve partial data.
- Copy nodes and relationships in batches. Every synced node gets a
_ProviderResource label and dynamic _Tenant_{uuid} / _Provider_{uuid}
isolation labels, plus a _provider_element_id property for MERGE keys.
@@ -64,10 +64,17 @@ from api.models import StateChoices
from api.utils import initialize_prowler_provider
from cartography.config import Config as CartographyConfig
from cartography.intel import analysis as cartography_analysis
from cartography.intel import create_indexes as cartography_create_indexes
from cartography.intel import ontology as cartography_ontology
from celery.utils.log import get_task_logger
from tasks.jobs.attack_paths import db_utils, findings, indexes, internet, sync, utils
from django.conf import settings
from tasks.jobs.attack_paths import (
db_utils,
findings,
indexes,
internet,
sync,
utils,
)
from tasks.jobs.attack_paths.config import get_cartography_ingestion_function
# Without this Celery goes crazy with Cartography logging
@@ -96,7 +103,7 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
attack_paths_scan = db_utils.retrieve_attack_paths_scan(tenant_id, scan_id)
# Idempotency guard: cleanup may have flipped this row to a terminal state
# while the message was still in flight. Bail out before touching state.
# while the message was still in flight. Bail out before touching state
if attack_paths_scan and attack_paths_scan.state in (
StateChoices.FAILED,
StateChoices.COMPLETED,
@@ -125,7 +132,7 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
else:
if not attack_paths_scan:
# Safety net for in-flight messages or direct task invocations; dispatcher normally pre-creates the row.
# Safety net for in-flight messages or direct task invocations; dispatcher normally pre-creates the row
logger.warning(
f"No Attack Paths Scan found for scan {scan_id} and tenant {tenant_id}, let's create it then"
)
@@ -143,10 +150,18 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
tenant_database_name = graph_database.get_database_name(
prowler_api_provider.tenant_id
)
target_sink_backend = settings.ATTACK_PATHS_SINK_DATABASE
target_description = (
f"tenant Neo4j database {tenant_database_name}"
if target_sink_backend == "neo4j"
else f"{target_sink_backend} sink"
)
# While creating the Cartography configuration, attributes `neo4j_user` and `neo4j_password` are not really needed in this config object
tmp_cartography_config = CartographyConfig(
neo4j_uri=graph_database.get_uri(),
# The temp ingest database is always Neo4j, so use the ingest URI here
# rather than the sink URI (which points at Neptune when configured).
neo4j_uri=graph_database.get_ingest_uri(),
neo4j_database=tmp_database_name,
update_tag=int(time.time()),
)
@@ -168,7 +183,8 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
scan_t0 = time.perf_counter()
logger.info(
f"Starting Attack Paths scan ({attack_paths_scan.id}) for "
f"{prowler_api_provider.provider.upper()} provider {prowler_api_provider.id}"
f"{prowler_api_provider.provider.upper()} provider {prowler_api_provider.id} "
f"(staging=Neo4j database {tmp_database_name}, target={target_description})"
)
subgraph_dropped = False
@@ -177,7 +193,8 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
try:
logger.info(
f"Creating Neo4j database {tmp_cartography_config.neo4j_database} for tenant {prowler_api_provider.tenant_id}"
f"Creating staging Neo4j database {tmp_cartography_config.neo4j_database} "
f"for tenant {prowler_api_provider.tenant_id}"
)
graph_database.create_database(tmp_cartography_config.neo4j_database)
@@ -191,7 +208,9 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
tmp_cartography_config.neo4j_database
) as tmp_neo4j_session:
# Indexes creation
cartography_create_indexes.run(tmp_neo4j_session, tmp_cartography_config)
indexes.create_cartography_indexes(
tmp_neo4j_session, tmp_cartography_config
)
indexes.create_findings_indexes(tmp_neo4j_session)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 2)
@@ -223,7 +242,7 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
cartography_analysis.run(tmp_neo4j_session, tmp_cartography_config)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 95)
# Creating Internet node and CAN_ACCESS relationships
# Creating Internet node and `CAN_ACCESS` relationships
logger.info(
f"Creating Internet graph for AWS account {prowler_api_provider.uid}"
)
@@ -247,23 +266,41 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 97)
logger.info(
f"Clearing Neo4j cache for database {tmp_cartography_config.neo4j_database}"
f"Clearing Neo4j cache for staging database {tmp_cartography_config.neo4j_database}"
)
graph_database.clear_cache(tmp_cartography_config.neo4j_database)
t0 = time.perf_counter()
logger.info(
f"Ensuring tenant database {tenant_database_name}, and its indexes, exists for tenant {prowler_api_provider.tenant_id}"
f"Preparing target {target_description} for tenant {prowler_api_provider.tenant_id}"
)
graph_database.create_database(tenant_database_name)
with graph_database.get_session(tenant_database_name) as tenant_neo4j_session:
cartography_create_indexes.run(
tenant_neo4j_session, tenant_cartography_config
)
indexes.create_findings_indexes(tenant_neo4j_session)
indexes.create_sync_indexes(tenant_neo4j_session)
# Sink-side index creation: Neptune auto-manages indexes and rejects
# `CREATE INDEX`, so only run it when the sink is Neo4j
# The temp ingest DB is always Neo4j and is always indexed above
if target_sink_backend != "neptune":
logger.info(f"Ensuring indexes exist for {target_description}")
with graph_database.get_session(
tenant_database_name
) as tenant_neo4j_session:
indexes.create_cartography_indexes(
tenant_neo4j_session, tenant_cartography_config
)
indexes.create_findings_indexes(tenant_neo4j_session)
indexes.create_sync_indexes(tenant_neo4j_session)
else:
logger.info("Skipping tenant database indexes for neptune sink")
logger.info(
f"Prepared target {target_description} in {time.perf_counter() - t0:.3f}s"
)
logger.info(f"Deleting existing provider graph in {tenant_database_name}")
db_utils.set_provider_graph_data_ready(attack_paths_scan, False)
logger.info(
f"Deleting existing provider graph from {target_description} "
f"(tenant={prowler_api_provider.tenant_id}, provider={prowler_api_provider.id})"
)
db_utils.set_provider_graph_data_ready(
attack_paths_scan, False, target_sink_backend
)
provider_gated = True
t0 = time.perf_counter()
@@ -272,14 +309,17 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
provider_id=str(prowler_api_provider.id),
)
logger.info(
f"Deleted existing provider graph in {time.perf_counter() - t0:.3f}s "
f"(deleted_nodes={deleted_nodes})"
f"Deleted existing provider graph from {target_description} "
f"in {time.perf_counter() - t0:.3f}s (deleted_nodes={deleted_nodes})"
)
subgraph_dropped = True
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 98)
logger.info(
f"Syncing graph from {tmp_database_name} into {tenant_database_name}"
f"Syncing staging graph {tmp_database_name} into {target_description} "
f"for provider {prowler_api_provider.id} "
f"(tenant {prowler_api_provider.tenant_id}, "
f"type {prowler_api_provider.provider})"
)
t0 = time.perf_counter()
sync_result = sync.sync_graph(
@@ -287,17 +327,34 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
target_database=tenant_database_name,
tenant_id=str(prowler_api_provider.tenant_id),
provider_id=str(prowler_api_provider.id),
provider_type=prowler_api_provider.provider,
)
elapsed = time.perf_counter() - t0
total_nodes = sync_result["nodes"] + sync_result["child_nodes"]
elements = total_nodes + sync_result["relationships"]
rate = elements / elapsed if elapsed else 0
logger.info(
f"Synced graph in {time.perf_counter() - t0:.3f}s "
f"(nodes={sync_result['nodes']}, relationships={sync_result['relationships']})"
f"Synced staging graph into {target_description} in {elapsed:.3f}s - "
f"nodes={total_nodes} (source={sync_result['nodes']}, "
f"items={sync_result['child_nodes']}), "
f"relationships={sync_result['relationships']} "
f"(structural={sync_result['structural_relationships']}, "
f"items={sync_result['item_relationships']}), "
f"~{rate:.0f} elem/s"
)
sync_completed = True
# Flip metadata only now: the new schema is live in the target sink, so
# reads can switch to the current catalog/backend. The target-sink gate
# is already closed, so the switch is atomic from the API's view.
db_utils.set_scan_migrated(attack_paths_scan, True, target_sink_backend)
db_utils.set_graph_data_ready(attack_paths_scan, True)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 99)
logger.info(f"Clearing Neo4j cache for database {tenant_database_name}")
graph_database.clear_cache(tenant_database_name)
if target_sink_backend == "neptune":
logger.info("Skipping cache clear for neptune sink")
else:
logger.info(f"Clearing Neo4j cache for target {target_description}")
graph_database.clear_cache(tenant_database_name)
logger.info(f"Dropping temporary Neo4j database {tmp_database_name}")
graph_database.drop_database(tmp_database_name)
@@ -316,14 +373,16 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
logger.exception(exception_message)
ingestion_exceptions["global_error"] = exception_message
# Recover graph_data_ready based on how far the swap got.
# Partial drop (mid-batch failure) may leave `subgraph_dropped=False`
# with data partially deleted, so we prefer that over permanently blocked queries.
# Recover `graph_data_ready` based on how far the swap got
# Partial drop (mid-batch failure) may leave `subgraph_dropped=False` with data partially deleted,
# so we prefer that over permanently blocked queries
try:
if sync_completed:
db_utils.set_graph_data_ready(attack_paths_scan, True)
elif provider_gated and not subgraph_dropped:
db_utils.set_provider_graph_data_ready(attack_paths_scan, True)
db_utils.set_provider_graph_data_ready(
attack_paths_scan, True, target_sink_backend
)
except Exception:
logger.error(
+356 -44
View File
@@ -1,40 +1,57 @@
"""
Graph sync operations for Attack Paths.
This module handles syncing graph data from temporary scan databases
to the tenant database, adding provider isolation labels and properties.
Reads nodes and relationships out of the cartography temp database (always
Neo4j) and hands them to the configured sink (Neo4j or Neptune) in batches.
Backend-specific Cypher (MERGE shape, ID strategy, indexes) lives in each
sink; this module owns the source read loop, per-batch grouping, and the
list-property materialisation policy (see `NormalizedList`).
Each list-typed node property that appears in the provider's
`normalized_lists` catalog becomes a set of child item nodes connected to
the parent by a typed edge. A list-typed property that is not in the
catalog is serialised to a comma-delimited string and emits a one-time
warning per (label, property), surfacing Cartography fields that should be
added to the catalog.
"""
import json
import time
from collections import defaultdict
from typing import Any
import neo4j
from api.attack_paths import database as graph_database
from api.attack_paths import sink as sink_module
from celery.utils.log import get_task_logger
from tasks.jobs.attack_paths.config import (
PROVIDER_CONFIGS,
PROVIDER_ISOLATION_PROPERTIES,
PROVIDER_RESOURCE_LABEL,
SYNC_BATCH_SIZE,
NormalizedList,
get_provider_label,
get_tenant_label,
)
from tasks.jobs.attack_paths.queries import (
NODE_FETCH_QUERY,
NODE_SYNC_TEMPLATE,
RELATIONSHIP_SYNC_TEMPLATE,
RELATIONSHIPS_FETCH_QUERY,
render_cypher_template,
)
logger = get_task_logger(__name__)
# (label, property) tuples for which we've already emitted the
# "unnormalised list" warning. Module-level so the warning fires once per
# process, not once per node.
_WARNED_UNNORMALIZED: set[tuple[str, str]] = set()
def sync_graph(
source_database: str,
target_database: str,
tenant_id: str,
provider_id: str,
provider_type: str,
) -> dict[str, int]:
"""
Sync all nodes and relationships from source to target database.
@@ -44,25 +61,38 @@ def sync_graph(
`target_database`: The tenant database
`tenant_id`: The tenant ID for isolation
`provider_id`: The provider ID for isolation
`provider_type`: Provider type key (e.g. "aws"), used to resolve the
`NormalizedList` catalog from `PROVIDER_CONFIGS`.
Returns:
Dict with counts of synced nodes and relationships
Dict with counts of synced nodes, child item nodes, and relationships.
"""
nodes_synced = sync_nodes(
sink = sink_module.get_backend()
sink.ensure_sync_indexes(target_database)
normalized_lists = _resolve_normalized_lists(provider_type)
node_result = sync_nodes(
source_database,
target_database,
tenant_id,
provider_id,
sink,
normalized_lists,
)
relationships_synced = sync_relationships(
source_database,
target_database,
provider_id,
sink,
)
return {
"nodes": nodes_synced,
"relationships": relationships_synced,
"nodes": node_result["parents"],
"child_nodes": node_result["children"],
"relationships": relationships_synced + node_result["parent_child_rels"],
"structural_relationships": relationships_synced,
"item_relationships": node_result["parent_child_rels"],
}
@@ -71,22 +101,35 @@ def sync_nodes(
target_database: str,
tenant_id: str,
provider_id: str,
) -> int:
sink: Any,
normalized_lists: list[NormalizedList],
) -> dict[str, int]:
"""
Sync nodes from source to target database.
Sync nodes from source to target database, exploding catalogued list
properties into child nodes + parent->child edges.
Adds `_ProviderResource` label and dynamic `_Tenant_{id}` and `_Provider_{id}`
isolation labels to all nodes.
isolation labels to all nodes (parents and children alike).
Source and target sessions are opened sequentially per batch to avoid
holding two Bolt connections simultaneously for the entire sync duration.
"""
t0 = time.perf_counter()
last_id = -1
total_synced = 0
parents_synced = 0
children_synced = 0
parent_child_rels = 0
catalog = _build_catalog_index(normalized_lists)
extra_labels = _build_extra_labels(tenant_id, provider_id)
while True:
grouped: dict[tuple[str, ...], list[dict[str, Any]]] = defaultdict(list)
tb = time.perf_counter()
prev_children = children_synced
prev_rels = parent_child_rels
parent_groups: dict[tuple[str, ...], list[dict[str, Any]]] = defaultdict(list)
child_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
rel_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
batch_count = 0
with graph_database.get_session(source_database) as source_session:
@@ -97,43 +140,65 @@ def sync_nodes(
for record in result:
batch_count += 1
last_id = record["internal_id"]
key, value = _node_to_sync_dict(record, provider_id)
grouped[key].append(value)
key, parent_dict, children, rels = _node_to_sync_dict(
record, provider_id, catalog
)
parent_groups[key].append(parent_dict)
for child in children:
child_groups[child["_child_label"]].append(child["row"])
for rel in rels:
rel_groups[rel["rel_type"]].append(rel["row"])
if batch_count == 0:
break
with graph_database.get_session(target_database) as target_session:
for labels, batch in grouped.items():
label_set = set(labels)
label_set.add(PROVIDER_RESOURCE_LABEL)
label_set.add(get_tenant_label(tenant_id))
label_set.add(get_provider_label(provider_id))
node_labels = ":".join(f"`{label}`" for label in sorted(label_set))
for labels, batch in parent_groups.items():
sink.write_nodes(
target_database, _render_labels(labels, extra_labels), batch
)
query = render_cypher_template(
NODE_SYNC_TEMPLATE, {"__NODE_LABELS__": node_labels}
)
target_session.run(query, {"rows": batch})
for child_label, batch in child_groups.items():
sink.write_nodes(
target_database,
_render_labels((child_label,), extra_labels),
batch,
)
children_synced += len(batch)
total_synced += batch_count
for rel_type, batch in rel_groups.items():
sink.write_relationships(target_database, rel_type, provider_id, batch)
parent_child_rels += len(batch)
parents_synced += batch_count
batch_dt = time.perf_counter() - tb
batch_elements = (
batch_count
+ (children_synced - prev_children)
+ (parent_child_rels - prev_rels)
)
rate = batch_elements / batch_dt if batch_dt else 0
logger.info(
f"Synced {total_synced} nodes from {source_database} to {target_database} in {time.perf_counter() - t0:.3f}s"
f"[sync nodes] {parents_synced} source (+{children_synced} items, "
f"+{parent_child_rels} item rels) · batch {batch_dt:.1f}s · "
f"elapsed {time.perf_counter() - t0:.1f}s · ~{rate:.0f} elem/s"
)
return total_synced
return {
"parents": parents_synced,
"children": children_synced,
"parent_child_rels": parent_child_rels,
}
def sync_relationships(
source_database: str,
target_database: str,
provider_id: str,
sink: Any,
) -> int:
"""
Sync relationships from source to target database.
Matches source and target nodes by `_provider_element_id` in the tenant database.
Source and target sessions are opened sequentially per batch to avoid
holding two Bolt connections simultaneously for the entire sync duration.
"""
@@ -142,6 +207,7 @@ def sync_relationships(
total_synced = 0
while True:
tb = time.perf_counter()
grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
batch_count = 0
@@ -159,32 +225,197 @@ def sync_relationships(
if batch_count == 0:
break
with graph_database.get_session(target_database) as target_session:
for rel_type, batch in grouped.items():
query = render_cypher_template(
RELATIONSHIP_SYNC_TEMPLATE, {"__REL_TYPE__": rel_type}
)
target_session.run(query, {"rows": batch})
for rel_type, batch in grouped.items():
sink.write_relationships(target_database, rel_type, provider_id, batch)
total_synced += batch_count
batch_dt = time.perf_counter() - tb
rate = batch_count / batch_dt if batch_dt else 0
logger.info(
f"Synced {total_synced} relationships from {source_database} to {target_database} in {time.perf_counter() - t0:.3f}s"
f"[sync rels] {total_synced} structural · batch {batch_dt:.1f}s · "
f"elapsed {time.perf_counter() - t0:.1f}s · ~{rate:.0f}/s"
)
return total_synced
def _node_to_sync_dict(
record: neo4j.Record, provider_id: str
) -> tuple[tuple[str, ...], dict[str, Any]]:
"""Transform a source node record into a (grouping_key, sync_dict) pair."""
record: neo4j.Record,
provider_id: str,
catalog: dict[tuple[str, str], NormalizedList],
) -> tuple[
tuple[str, ...],
dict[str, Any],
list[dict[str, Any]],
list[dict[str, Any]],
]:
"""Transform a source node record into a (grouping_key, sync_dict, children, rels) tuple.
Catalogued list properties are popped from `props` and emitted as child
nodes + parent->child relationships.
"""
props = dict(record["props"] or {})
_strip_internal_properties(props)
labels = tuple(sorted(set(record["labels"] or [])))
return labels, {
"provider_element_id": f"{provider_id}:{record['element_id']}",
parent_element_id = f"{provider_id}:{record['element_id']}"
children, rels = _explode_catalogued_lists(
labels, props, catalog, provider_id, parent_element_id
)
_normalize_sink_properties(props, labels)
parent = {
"provider_element_id": parent_element_id,
"props": props,
}
return labels, parent, children, rels
def _explode_catalogued_lists(
labels: tuple[str, ...],
props: dict[str, Any],
catalog: dict[tuple[str, str], NormalizedList],
provider_id: str,
parent_element_id: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""Pop catalogued list properties from `props` and produce child + rel emits.
A node may carry multiple labels (e.g. `AWSPolicyStatement` plus
`_AWSResource`); we check each label for catalog matches independently.
Returns:
- children: list of {"_child_label": str, "row": <node row>} dicts.
- rels: list of {"rel_type": str, "row": <rel row>} dicts.
"""
children: list[dict[str, Any]] = []
rels: list[dict[str, Any]] = []
for label in labels:
for key in list(props.keys()):
spec = catalog.get((label, key))
if spec is None:
continue
value = props.pop(key)
if value is None:
continue
if not isinstance(value, list):
# Catalogued but not actually a list this scan - fall back to
# the generic normaliser so we don't lose the value.
props[key] = value
continue
for item in value:
child_value_key, child_props = _build_child_props(spec, item)
if child_value_key is None:
continue
child_element_id = _build_child_id(
provider_id, spec.child_label, child_value_key
)
children.append(
{
"_child_label": spec.child_label,
"row": {
"provider_element_id": child_element_id,
"props": child_props,
},
}
)
rels.append(
{
"rel_type": spec.rel_type,
"row": {
"start_element_id": parent_element_id,
"end_element_id": child_element_id,
"provider_element_id": (
f"{parent_element_id}::{spec.rel_type}::"
f"{child_element_id}"
),
"props": {},
},
}
)
return children, rels
def _build_child_props(
spec: NormalizedList, item: Any
) -> tuple[str | None, dict[str, Any]]:
"""Translate one list element into a child node's prop dict.
Returns (dedup_key, props). The dedup_key is what makes two child nodes
equal within (tenant, provider) - used to build `_provider_element_id`.
For scalar mode, the dedup key is the value itself. For dict mode it is
a stable concatenation of the mapped fields in `field_map` order.
"""
if not spec.field_map:
if isinstance(item, (dict, list)):
# Defensive: caller marked this list as scalar but elements are
# structured. Convert to a stable string so the value survives.
value_str = json.dumps(item, sort_keys=True, default=str)
else:
value_str = str(item)
return value_str, {"value": value_str}
if not isinstance(item, dict):
# Catalogued as dict-shape but got a scalar. Skip - caller will see
# the value go missing and can fix the field_map.
return None, {}
props: dict[str, Any] = {}
dedup_parts: list[str] = []
for src_key, child_field in spec.field_map:
raw = item.get(src_key)
value_str = _to_sink_property_value(raw) if raw is not None else ""
props[child_field] = value_str
dedup_parts.append(f"{child_field}={value_str}")
return "::".join(dedup_parts), props
def _build_child_id(provider_id: str, child_label: str, value_key: str) -> str:
"""Deterministic `_provider_element_id` for a list-item child node.
Dedupes within (tenant, provider): multiple parents referencing the same
value share one child node via the existing MERGE-on-_provider_element_id
index in both sinks.
"""
return f"{provider_id}::{child_label}::{value_key}"
def _build_catalog_index(
normalized_lists: list[NormalizedList],
) -> dict[tuple[str, str], NormalizedList]:
"""Index the catalog by (source_label, source_property) for O(1) lookup."""
return {
(spec.source_label, spec.source_property): spec for spec in normalized_lists
}
def _build_extra_labels(tenant_id: str, provider_id: str) -> tuple[str, ...]:
return (
PROVIDER_RESOURCE_LABEL,
get_tenant_label(tenant_id),
get_provider_label(provider_id),
)
def _render_labels(base_labels: tuple[str, ...], extra_labels: tuple[str, ...]) -> str:
"""Render the Cypher label string for a node-write batch."""
label_set = set(base_labels) | set(extra_labels)
return ":".join(f"`{label}`" for label in sorted(label_set))
def _resolve_normalized_lists(provider_type: str) -> list[NormalizedList]:
config = PROVIDER_CONFIGS.get(provider_type)
if config is None:
# Unknown provider: empty catalog. Any list-typed property will be
# serialised to a comma-delimited string with one warning per
# (label, property).
logger.warning(
"Provider type %s not in PROVIDER_CONFIGS; no normalized_lists active",
provider_type,
)
return []
return config.normalized_lists
def _rel_to_sync_dict(
@@ -193,7 +424,11 @@ def _rel_to_sync_dict(
"""Transform a source relationship record into a (grouping_key, sync_dict) pair."""
props = dict(record["props"] or {})
_strip_internal_properties(props)
# Relationship properties go through the same primitive coercion as
# nodes; catalog-driven materialisation applies to node properties only.
_normalize_sink_properties(props, labels=None)
rel_type = record["rel_type"]
return rel_type, {
"start_element_id": f"{provider_id}:{record['start_element_id']}",
"end_element_id": f"{provider_id}:{record['end_element_id']}",
@@ -206,3 +441,80 @@ def _strip_internal_properties(props: dict[str, Any]) -> None:
"""Remove provider isolation properties before the += spread in sync templates."""
for key in PROVIDER_ISOLATION_PROPERTIES:
props.pop(key, None)
def _normalize_sink_properties(
props: dict[str, Any], labels: tuple[str, ...] | None
) -> None:
"""Normalize property values to primitive Cypher literals for either sink.
Attack-paths node and relationship properties are written as primitive
scalars regardless of the active sink (Neo4j or Neptune). The convention
is driven by Neptune's openCypher type restrictions, which reject list,
map, temporal and spatial property values, but it is applied uniformly
so that custom and predefined queries are portable across sinks without
runtime rewriting.
Concretely:
- Temporal values (neo4j.time.{DateTime,Date,Time,Duration}) become
their ISO-8601 string representation.
- Spatial values (neo4j.spatial.Point and subclasses) become their
WKT-style string representation.
- Maps / dicts become a JSON-encoded string, read back with `CONTAINS`
substring checks inside queries.
- Lists become a comma-delimited string. Catalogued list properties
are materialised as child item nodes upstream in
`_explode_catalogued_lists` and never reach this point; any list
seen here is uncatalogued, so we log a one-time warning per
(label, property) to surface Cartography fields that should be
added to the catalog.
`labels` is only used for the warning message; pass `None` for
relationship props (no label context).
"""
for key, value in list(props.items()):
if isinstance(value, list) and labels is not None:
_warn_unnormalized_list(labels, key)
props[key] = _to_sink_property_value(value)
def _warn_unnormalized_list(labels: tuple[str, ...], key: str) -> None:
"""Warn once per (label, property), on the real label(s) only.
Every synced node also carries internal isolation labels (`_AWSResource`,
`_ProviderResource`, `_Tenant_*`, `_Provider_*`); warning on those just
doubles the noise, so skip them and point at the actionable Cartography
label. Falls back to all labels if only internal ones are present.
"""
real_labels = [label for label in labels if not label.startswith("_")]
for label in real_labels or labels:
token = (label, key)
if token in _WARNED_UNNORMALIZED:
continue
_WARNED_UNNORMALIZED.add(token)
logger.warning(
"Unnormalized list property %s.%s reached sink as comma-string; "
"add a NormalizedList entry to the provider catalog to explode it",
label,
key,
)
def _to_sink_property_value(value: Any) -> Any:
if hasattr(value, "iso_format") and callable(value.iso_format):
return value.iso_format()
if type(value).__module__.startswith("neo4j.spatial"):
return str(value)
if isinstance(value, dict):
# openCypher `SET` rejects map property values: encode as JSON so the structured payload
# survives the round-trip and is queryable with `CONTAINS` substring checks
return json.dumps(value, sort_keys=True, default=str)
if isinstance(value, list):
# openCypher `SET` rejects list/array property values: encode as a
# delimited string read back with split() inside queries
return ",".join(str(_to_sink_property_value(v)) for v in value)
return value
+14 -1
View File
@@ -1,4 +1,5 @@
from api.attack_paths import database as graph_database
from api.attack_paths import sink as sink_module
from api.db_router import MainRouter
from api.db_utils import batch_delete, rls_transaction
from api.models import (
@@ -76,6 +77,12 @@ def delete_provider(tenant_id: str, pk: str):
"id", flat=True
)
)
attack_paths_sink_backends = list(
AttackPathsScan.all_objects.filter(provider=instance)
.values_list("sink_backend", flat=True)
.distinct()
.order_by("sink_backend")
)
deletion_steps = [
("Scan Summaries", ScanSummary.all_objects.filter(scan__provider=instance)),
@@ -97,7 +104,13 @@ def delete_provider(tenant_id: str, pk: str):
# Delete the Attack Paths' graph data related to the provider from the tenant database
tenant_database_name = graph_database.get_database_name(tenant_id)
try:
graph_database.drop_subgraph(tenant_database_name, str(pk))
if attack_paths_sink_backends:
for sink_backend in attack_paths_sink_backends:
sink_module.get_backend_for_name(sink_backend).drop_subgraph(
tenant_database_name, str(pk)
)
else:
graph_database.drop_subgraph(tenant_database_name, str(pk))
except graph_database.GraphDatabaseQueryException as gdb_error:
logger.error(f"Error deleting Provider graph data: {gdb_error}")
@@ -23,6 +23,14 @@ from tasks.jobs.attack_paths import internet as internet_module
from tasks.jobs.attack_paths import sync as sync_module
from tasks.jobs.attack_paths.scan import run as attack_paths_run
SYNC_RESULT_EMPTY = {
"nodes": 0,
"child_nodes": 0,
"relationships": 0,
"structural_relationships": 0,
"item_relationships": 0,
}
@pytest.mark.django_db
class TestAttackPathsRun:
@@ -32,6 +40,7 @@ class TestAttackPathsRun:
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.db_utils.set_scan_migrated")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@@ -39,7 +48,7 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
return_value=SYNC_RESULT_EMPTY,
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_subgraph", return_value=0)
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@@ -48,11 +57,11 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
"tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri",
return_value="bolt://neo4j",
)
@patch(
@@ -66,7 +75,7 @@ class TestAttackPathsRun:
def test_run_success_flow(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
@@ -83,6 +92,7 @@ class TestAttackPathsRun:
mock_finish,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_set_scan_migrated,
mock_event_loop,
mock_drop_db,
tenants_fixture,
@@ -159,6 +169,7 @@ class TestAttackPathsRun:
target_database="tenant-db",
tenant_id=str(provider.tenant_id),
provider_id=str(provider.id),
provider_type="aws",
)
mock_get_ingestion.assert_called_once_with(provider.provider)
mock_event_loop.assert_called_once()
@@ -172,9 +183,12 @@ class TestAttackPathsRun:
attack_paths_scan, StateChoices.COMPLETED, ingestion_result
)
mock_set_provider_graph_data_ready.assert_called_once_with(
attack_paths_scan, False
attack_paths_scan, False, "neo4j"
)
mock_set_graph_data_ready.assert_called_once_with(attack_paths_scan, True)
# is_migrated is flipped to True only after the sync succeeds, so reads
# don't switch to the new catalog/sink before the graph is live.
mock_set_scan_migrated.assert_called_once_with(attack_paths_scan, True, "neo4j")
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
@@ -194,13 +208,13 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
)
@patch("tasks.jobs.attack_paths.scan.graph_database.get_uri")
@patch("tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri")
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
@@ -212,7 +226,7 @@ class TestAttackPathsRun:
def test_run_failure_marks_scan_failed(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_get_db_name,
mock_create_db,
mock_cartography_indexes,
@@ -293,13 +307,13 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
)
@patch("tasks.jobs.attack_paths.scan.graph_database.get_uri")
@patch("tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri")
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
@@ -311,7 +325,7 @@ class TestAttackPathsRun:
def test_failure_before_gate_does_not_flip_graph_data_ready_true(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_get_db_name,
mock_create_db,
mock_cartography_indexes,
@@ -396,13 +410,13 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
)
@patch("tasks.jobs.attack_paths.scan.graph_database.get_uri")
@patch("tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri")
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
@@ -414,7 +428,7 @@ class TestAttackPathsRun:
def test_run_failure_marks_scan_failed_even_when_drop_database_fails(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_get_db_name,
mock_create_db,
mock_cartography_indexes,
@@ -493,7 +507,7 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
return_value=SYNC_RESULT_EMPTY,
)
@patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_subgraph",
@@ -505,11 +519,11 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
"tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri",
return_value="bolt://neo4j",
)
@patch(
@@ -523,7 +537,7 @@ class TestAttackPathsRun:
def test_failure_after_gate_before_drop_restores_graph_data_ready(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
@@ -589,8 +603,8 @@ class TestAttackPathsRun:
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
assert mock_set_provider_graph_data_ready.call_args_list == [
call(attack_paths_scan, False),
call(attack_paths_scan, True),
call(attack_paths_scan, False, "neo4j"),
call(attack_paths_scan, True, "neo4j"),
]
@patch(
@@ -618,11 +632,11 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
"tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri",
return_value="bolt://neo4j",
)
@patch(
@@ -636,7 +650,7 @@ class TestAttackPathsRun:
def test_failure_after_drop_before_sync_leaves_graph_data_ready_false(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
@@ -703,7 +717,7 @@ class TestAttackPathsRun:
# Only called with False (gate), never with True (no recovery for partial data)
mock_set_provider_graph_data_ready.assert_called_once_with(
attack_paths_scan, False
attack_paths_scan, False, "neo4j"
)
@patch(
@@ -716,6 +730,7 @@ class TestAttackPathsRun:
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_scan_migrated")
@patch(
"tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready",
side_effect=[RuntimeError("flag failed"), None],
@@ -725,7 +740,7 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
return_value=SYNC_RESULT_EMPTY,
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_subgraph")
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@@ -734,11 +749,11 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
"tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri",
return_value="bolt://neo4j",
)
@patch(
@@ -752,7 +767,7 @@ class TestAttackPathsRun:
def test_failure_after_sync_restores_graph_data_ready(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
@@ -768,6 +783,7 @@ class TestAttackPathsRun:
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_set_scan_migrated,
mock_finish,
mock_drop_db,
mock_event_loop,
@@ -824,8 +840,11 @@ class TestAttackPathsRun:
]
# set_provider_graph_data_ready only called once with False (the gate)
mock_set_provider_graph_data_ready.assert_called_once_with(
attack_paths_scan, False
attack_paths_scan, False, "neo4j"
)
# is_migrated is flipped once after the sync and is not touched again by
# the failure-recovery branch
mock_set_scan_migrated.assert_called_once_with(attack_paths_scan, True, "neo4j")
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
@@ -843,7 +862,7 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
return_value=SYNC_RESULT_EMPTY,
)
@patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_subgraph",
@@ -855,11 +874,11 @@ class TestAttackPathsRun:
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.indexes.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
"tasks.jobs.attack_paths.scan.graph_database.get_ingest_uri",
return_value="bolt://neo4j",
)
@patch(
@@ -873,7 +892,7 @@ class TestAttackPathsRun:
def test_recovery_failure_does_not_suppress_original_exception(
self,
mock_init_provider,
mock_get_uri,
mock_get_ingest_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
@@ -1116,7 +1135,7 @@ class TestFailAttackPathsScan:
fail_attack_paths_scan(str(tenant.id), "nonexistent", "setup exploded")
def test_fail_recovers_graph_data_ready_when_data_exists(
self, tenants_fixture, providers_fixture, scans_fixture
self, tenants_fixture, providers_fixture, scans_fixture, sink_backend_stub
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
@@ -1135,16 +1154,18 @@ class TestFailAttackPathsScan:
state=StateChoices.EXECUTING,
)
# `recover_graph_data_ready` routes `has_provider_data` through
# `sink_module.get_backend_for_scan(scan)`. With `is_migrated=False`
# and the default `ATTACK_PATHS_SINK_DATABASE=neo4j`, the factory
# returns the active backend, which `sink_backend_stub` replaces.
sink_backend_stub.has_provider_data.return_value = True
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch("tasks.jobs.attack_paths.db_utils.graph_database.drop_database"),
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.has_provider_data",
return_value=True,
),
patch(
"tasks.jobs.attack_paths.db_utils.set_provider_graph_data_ready"
) as mock_set_ready,
@@ -1154,7 +1175,7 @@ class TestFailAttackPathsScan:
mock_set_ready.assert_called_once_with(attack_paths_scan, True)
def test_fail_leaves_graph_data_ready_false_when_no_data(
self, tenants_fixture, providers_fixture, scans_fixture
self, tenants_fixture, providers_fixture, scans_fixture, sink_backend_stub
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
@@ -1173,16 +1194,14 @@ class TestFailAttackPathsScan:
state=StateChoices.EXECUTING,
)
sink_backend_stub.has_provider_data.return_value = False
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch("tasks.jobs.attack_paths.db_utils.graph_database.drop_database"),
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.has_provider_data",
return_value=False,
),
patch(
"tasks.jobs.attack_paths.db_utils.set_provider_graph_data_ready"
) as mock_set_ready,
@@ -1271,6 +1290,20 @@ class TestAttackPathsFindingsHelpers:
[call(mock_session, stmt) for stmt in FINDINGS_INDEX_STATEMENTS]
)
def test_create_findings_indexes_runs_even_when_sink_is_neptune(self, settings):
# The index helpers run against the temp ingest DB, which is always
# Neo4j regardless of the configured sink. A Neptune sink must not
# suppress index creation on that DB (regression for the dropped
# in-helper sink gate).
settings.ATTACK_PATHS_SINK_DATABASE = "neptune"
mock_session = MagicMock()
with patch("tasks.jobs.attack_paths.indexes.run_write_query") as mock_run_write:
indexes_module.create_findings_indexes(mock_session)
from tasks.jobs.attack_paths.indexes import FINDINGS_INDEX_STATEMENTS
assert mock_run_write.call_count == len(FINDINGS_INDEX_STATEMENTS)
def test_load_findings_batches_requests(self, providers_fixture):
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
@@ -1802,7 +1835,7 @@ def _make_session_ctx(session, call_order=None, name=None):
class TestSyncNodes:
def test_sync_nodes_adds_private_label(self):
def test_sync_nodes_passes_isolation_labels_to_sink(self):
row = {
"internal_id": 1,
"element_id": "elem-1",
@@ -1812,29 +1845,32 @@ class TestSyncNodes:
mock_source_1 = MagicMock()
mock_source_1.run.return_value = [row]
mock_target = MagicMock()
mock_source_2 = MagicMock()
mock_source_2.run.return_value = []
sink = MagicMock()
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(mock_source_1),
_make_session_ctx(mock_target),
_make_session_ctx(mock_source_2),
],
):
total = sync_module.sync_nodes(
"source-db", "target-db", "tenant-1", "prov-1"
result = sync_module.sync_nodes(
"source-db", "target-db", "tenant-1", "prov-1", sink, []
)
assert total == 1
query = mock_target.run.call_args.args[0]
assert "_ProviderResource" in query
assert "_Tenant_tenant1" in query
assert "_Provider_prov1" in query
assert result["parents"] == 1
sink.write_nodes.assert_called_once()
target_db, labels, batch = sink.write_nodes.call_args.args
assert target_db == "target-db"
assert "_ProviderResource" in labels
assert "_Tenant_tenant1" in labels
assert "_Provider_prov1" in labels
assert batch[0]["provider_element_id"] == "prov-1:elem-1"
assert batch[0]["props"] == {"key": "value"}
def test_sync_nodes_source_closes_before_target_opens(self):
def test_sync_nodes_writes_after_source_session_closes(self):
row = {
"internal_id": 1,
"element_id": "elem-1",
@@ -1846,21 +1882,23 @@ class TestSyncNodes:
src_1 = MagicMock()
src_1.run.return_value = [row]
tgt = MagicMock()
src_2 = MagicMock()
src_2.run.return_value = []
sink = MagicMock()
sink.write_nodes.side_effect = lambda *_a, **_kw: call_order.append(
"sink:write"
)
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1, call_order, "source1"),
_make_session_ctx(tgt, call_order, "target"),
_make_session_ctx(src_2, call_order, "source2"),
],
):
sync_module.sync_nodes("src-db", "tgt-db", "t-1", "p-1")
sync_module.sync_nodes("src-db", "tgt-db", "t-1", "p-1", sink, [])
assert call_order.index("source1:exit") < call_order.index("target:enter")
assert call_order.index("source1:exit") < call_order.index("sink:write")
def test_sync_nodes_pagination_with_batch_size_1(self):
row_a = {
@@ -1882,44 +1920,44 @@ class TestSyncNodes:
src_2.run.return_value = [row_b]
src_3 = MagicMock()
src_3.run.return_value = []
tgt_1 = MagicMock()
tgt_2 = MagicMock()
sink = MagicMock()
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(tgt_1),
_make_session_ctx(src_2),
_make_session_ctx(tgt_2),
_make_session_ctx(src_3),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1),
):
total = sync_module.sync_nodes("src", "tgt", "t-1", "p-1")
result = sync_module.sync_nodes("src", "tgt", "t-1", "p-1", sink, [])
assert total == 2
assert result["parents"] == 2
assert sink.write_nodes.call_count == 2
assert src_1.run.call_args.args[1]["last_id"] == -1
assert src_2.run.call_args.args[1]["last_id"] == 1
def test_sync_nodes_empty_source_returns_zero(self):
src = MagicMock()
src.run.return_value = []
sink = MagicMock()
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[_make_session_ctx(src)],
) as mock_get_session:
total = sync_module.sync_nodes("src", "tgt", "t-1", "p-1")
result = sync_module.sync_nodes("src", "tgt", "t-1", "p-1", sink, [])
assert total == 0
assert result["parents"] == 0
assert mock_get_session.call_count == 1
sink.write_nodes.assert_not_called()
class TestSyncRelationships:
def test_sync_relationships_source_closes_before_target_opens(self):
def test_sync_relationships_writes_after_source_session_closes(self):
row = {
"internal_id": 1,
"rel_type": "HAS",
@@ -1932,21 +1970,23 @@ class TestSyncRelationships:
src_1 = MagicMock()
src_1.run.return_value = [row]
tgt = MagicMock()
src_2 = MagicMock()
src_2.run.return_value = []
sink = MagicMock()
sink.write_relationships.side_effect = lambda *_a, **_kw: call_order.append(
"sink:write"
)
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1, call_order, "source1"),
_make_session_ctx(tgt, call_order, "target"),
_make_session_ctx(src_2, call_order, "source2"),
],
):
sync_module.sync_relationships("src", "tgt", "p-1")
sync_module.sync_relationships("src", "tgt", "p-1", sink)
assert call_order.index("source1:exit") < call_order.index("target:enter")
assert call_order.index("source1:exit") < call_order.index("sink:write")
def test_sync_relationships_pagination_with_batch_size_1(self):
row_a = {
@@ -1970,40 +2010,40 @@ class TestSyncRelationships:
src_2.run.return_value = [row_b]
src_3 = MagicMock()
src_3.run.return_value = []
tgt_1 = MagicMock()
tgt_2 = MagicMock()
sink = MagicMock()
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(tgt_1),
_make_session_ctx(src_2),
_make_session_ctx(tgt_2),
_make_session_ctx(src_3),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1),
):
total = sync_module.sync_relationships("src", "tgt", "p-1")
total = sync_module.sync_relationships("src", "tgt", "p-1", sink)
assert total == 2
assert sink.write_relationships.call_count == 2
assert src_1.run.call_args.args[1]["last_id"] == -1
assert src_2.run.call_args.args[1]["last_id"] == 1
def test_sync_relationships_empty_source_returns_zero(self):
src = MagicMock()
src.run.return_value = []
sink = MagicMock()
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[_make_session_ctx(src)],
) as mock_get_session:
total = sync_module.sync_relationships("src", "tgt", "p-1")
total = sync_module.sync_relationships("src", "tgt", "p-1", sink)
assert total == 0
assert mock_get_session.call_count == 1
sink.write_relationships.assert_not_called()
class TestInternetAnalysis:
@@ -2075,6 +2115,8 @@ class TestAttackPathsDbUtilsGraphDataReady:
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is False
assert attack_paths_scan.is_migrated is False
assert attack_paths_scan.sink_backend == "neo4j"
def test_create_attack_paths_scan_inherits_true_from_previous(
self, tenants_fixture, providers_fixture, scans_fixture
@@ -2095,6 +2137,8 @@ class TestAttackPathsDbUtilsGraphDataReady:
scan=scan,
state=StateChoices.COMPLETED,
graph_data_ready=True,
is_migrated=True,
sink_backend="neptune",
)
new_scan = Scan.objects.create(
@@ -2115,6 +2159,109 @@ class TestAttackPathsDbUtilsGraphDataReady:
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is True
# is_migrated tracks the data being served: inherited from the ready scan
assert attack_paths_scan.is_migrated is True
assert attack_paths_scan.sink_backend == "neptune"
def test_create_attack_paths_scan_prefers_active_sink_ready_scan(
self, tenants_fixture, providers_fixture, scans_fixture, settings
):
from tasks.jobs.attack_paths.db_utils import create_attack_paths_scan
settings.ATTACK_PATHS_SINK_DATABASE = "neo4j"
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.COMPLETED,
graph_data_ready=True,
is_migrated=False,
sink_backend="neo4j",
)
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.COMPLETED,
graph_data_ready=True,
is_migrated=True,
sink_backend="neptune",
)
new_scan = Scan.objects.create(
name="New Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.AVAILABLE,
tenant_id=tenant.id,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
attack_paths_scan = create_attack_paths_scan(
str(tenant.id), str(new_scan.id), provider.id
)
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is True
assert attack_paths_scan.is_migrated is False
assert attack_paths_scan.sink_backend == "neo4j"
def test_create_attack_paths_scan_inherits_is_migrated_false_from_legacy_ready(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import create_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
# Previous scan is ready but pre-cutover (legacy Neo4j graph shape)
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.COMPLETED,
graph_data_ready=True,
is_migrated=False,
sink_backend="neo4j",
)
new_scan = Scan.objects.create(
name="New Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.AVAILABLE,
tenant_id=tenant.id,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
attack_paths_scan = create_attack_paths_scan(
str(tenant.id), str(new_scan.id), provider.id
)
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is True
# Reads stay on the legacy catalog/backend until this scan's own sync
assert attack_paths_scan.is_migrated is False
assert attack_paths_scan.sink_backend == "neo4j"
def test_create_attack_paths_scan_inherits_false_when_no_previous_ready(
self, tenants_fixture, providers_fixture, scans_fixture
@@ -2135,6 +2282,7 @@ class TestAttackPathsDbUtilsGraphDataReady:
scan=scan,
state=StateChoices.FAILED,
graph_data_ready=False,
sink_backend="neptune",
)
new_scan = Scan.objects.create(
@@ -2155,6 +2303,8 @@ class TestAttackPathsDbUtilsGraphDataReady:
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is False
assert attack_paths_scan.is_migrated is False
assert attack_paths_scan.sink_backend == "neo4j"
def test_set_graph_data_ready_updates_field(
self, tenants_fixture, providers_fixture, scans_fixture
@@ -2261,7 +2411,7 @@ class TestAttackPathsDbUtilsGraphDataReady:
assert attack_paths_scan.state == StateChoices.FAILED
assert attack_paths_scan.graph_data_ready is True
def test_set_provider_graph_data_ready_updates_all_scans_for_provider(
def test_set_provider_graph_data_ready_updates_all_scans_for_provider_sink(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import set_provider_graph_data_ready
@@ -2289,6 +2439,7 @@ class TestAttackPathsDbUtilsGraphDataReady:
scan=scan_a,
state=StateChoices.COMPLETED,
graph_data_ready=True,
sink_backend="neptune",
)
new_ap_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
@@ -2296,6 +2447,7 @@ class TestAttackPathsDbUtilsGraphDataReady:
scan=scan_b,
state=StateChoices.EXECUTING,
graph_data_ready=True,
sink_backend="neptune",
)
with patch(
@@ -2309,6 +2461,48 @@ class TestAttackPathsDbUtilsGraphDataReady:
assert old_ap_scan.graph_data_ready is False
assert new_ap_scan.graph_data_ready is False
def test_set_provider_graph_data_ready_preserves_other_sink_scans(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import set_provider_graph_data_ready
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
legacy_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.COMPLETED,
graph_data_ready=True,
sink_backend="neo4j",
)
neptune_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
graph_data_ready=True,
sink_backend="neptune",
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
set_provider_graph_data_ready(neptune_scan, False)
legacy_scan.refresh_from_db()
neptune_scan.refresh_from_db()
assert legacy_scan.graph_data_ready is True
assert neptune_scan.graph_data_ready is False
def test_set_provider_graph_data_ready_does_not_affect_other_providers(
self, tenants_fixture, providers_fixture, scans_fixture
):
@@ -2871,3 +3065,57 @@ class TestCleanupStaleAttackPathsScans:
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.SCHEDULED
mock_revoke.assert_not_called()
class TestNormalizeSinkProperties:
"""Coerce Cartography-emitted property values into sink-portable primitives.
Lists become comma-strings, dicts become JSON strings, temporals become
ISO strings, spatials become their stringified form. The same coercion
runs regardless of the active sink so queries are portable.
"""
@pytest.mark.parametrize(
"raw, expected",
[
(
{"a": "x", "b": 1, "c": 1.5, "d": True, "e": None},
{"a": "x", "b": 1, "c": 1.5, "d": True, "e": None},
),
(
{"actions": ["s3:GetObject", "s3:PutObject"], "tags": []},
{"actions": "s3:GetObject,s3:PutObject", "tags": ""},
),
(
{"condition": {"StringEquals": {"aws:SourceAccount": "123456789012"}}},
{
"condition": '{"StringEquals": {"aws:SourceAccount": "123456789012"}}'
},
),
],
)
def test_primitive_list_and_dict_branches(self, raw, expected):
sync_module._normalize_sink_properties(raw, labels=None)
assert raw == expected
def test_temporal_and_spatial_become_strings(self):
class FakeDateTime:
def iso_format(self) -> str:
return "2026-05-13T10:00:00+00:00"
class FakeSpatialPoint:
def __str__(self) -> str:
return "POINT(1.0 2.0)"
# The spatial branch is detected by module prefix, not by base class.
FakeSpatialPoint.__module__ = "neo4j.spatial.fake"
props = {
"created_at": FakeDateTime(),
"location": FakeSpatialPoint(),
}
sync_module._normalize_sink_properties(props, labels=None)
assert props == {
"created_at": "2026-05-13T10:00:00+00:00",
"location": "POINT(1.0 2.0)",
}
+50 -3
View File
@@ -1,4 +1,4 @@
from unittest.mock import call, patch
from unittest.mock import MagicMock, call, patch
import pytest
from api.attack_paths import database as graph_database
@@ -60,10 +60,12 @@ class TestDeleteProvider:
aps1 = create_attack_paths_scan(instance)
aps2 = create_attack_paths_scan(instance)
backend = MagicMock()
with (
patch(
"tasks.jobs.deletion.graph_database.drop_subgraph",
"tasks.jobs.deletion.sink_module.get_backend_for_name",
return_value=backend,
),
patch(
"tasks.jobs.deletion.graph_database.drop_database",
@@ -72,12 +74,55 @@ class TestDeleteProvider:
result = delete_provider(tenant_id, instance.id)
assert result
backend.drop_subgraph.assert_called_once_with(
graph_database.get_database_name(tenant_id), str(instance.id)
)
expected_tmp_calls = [
call(f"db-tmp-scan-{str(aps1.id).lower()}"),
call(f"db-tmp-scan-{str(aps2.id).lower()}"),
]
mock_drop_database.assert_has_calls(expected_tmp_calls, any_order=True)
def test_delete_provider_drops_graph_data_from_all_recorded_sinks(
self, providers_fixture, create_attack_paths_scan
):
instance = providers_fixture[0]
tenant_id = str(instance.tenant_id)
create_attack_paths_scan(instance, sink_backend="neo4j")
create_attack_paths_scan(instance, sink_backend="neptune")
neo4j_backend = MagicMock()
neptune_backend = MagicMock()
def get_backend_for_name(name):
return {
"neo4j": neo4j_backend,
"neptune": neptune_backend,
}[name]
with (
patch(
"tasks.jobs.deletion.graph_database.get_database_name",
return_value="tenant-db",
),
patch(
"tasks.jobs.deletion.sink_module.get_backend_for_name",
side_effect=get_backend_for_name,
) as mock_get_backend_for_name,
patch("tasks.jobs.deletion.graph_database.drop_database"),
):
result = delete_provider(tenant_id, instance.id)
assert result
mock_get_backend_for_name.assert_has_calls(
[call("neo4j"), call("neptune")], any_order=True
)
neo4j_backend.drop_subgraph.assert_called_once_with(
"tenant-db", str(instance.id)
)
neptune_backend.drop_subgraph.assert_called_once_with(
"tenant-db", str(instance.id)
)
def test_delete_provider_continues_when_temp_db_drop_fails(
self, providers_fixture, create_attack_paths_scan
):
@@ -85,10 +130,12 @@ class TestDeleteProvider:
tenant_id = str(instance.tenant_id)
create_attack_paths_scan(instance)
backend = MagicMock()
with (
patch(
"tasks.jobs.deletion.graph_database.drop_subgraph",
"tasks.jobs.deletion.sink_module.get_backend_for_name",
return_value=backend,
),
patch(
"tasks.jobs.deletion.graph_database.drop_database",