mirror of
https://github.com/prowler-cloud/prowler.git
synced 2026-10-24 04:14:18 +00:00
807 lines
29 KiB
Python
807 lines
29 KiB
Python
from contextlib import nullcontext
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pytest
|
|
from tasks.jobs.attack_paths import findings as findings_module
|
|
from tasks.jobs.attack_paths import internet as internet_module
|
|
from tasks.jobs.attack_paths.scan import run as attack_paths_run
|
|
|
|
from api.models import (
|
|
AttackPathsScan,
|
|
Finding,
|
|
Provider,
|
|
Resource,
|
|
ResourceFindingMapping,
|
|
Scan,
|
|
StateChoices,
|
|
StatusChoices,
|
|
)
|
|
from prowler.lib.check.models import Severity
|
|
|
|
|
|
@pytest.mark.django_db
|
|
class TestAttackPathsRun:
|
|
# Patching with decorators as we got a `SyntaxError: too many statically nested blocks` error if we use context managers
|
|
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
|
|
@patch(
|
|
"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.get_old_attack_paths_scans",
|
|
return_value=[],
|
|
)
|
|
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
|
|
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
|
|
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
|
|
@patch("tasks.jobs.attack_paths.scan.sync.sync_graph")
|
|
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_subgraph")
|
|
@patch("tasks.jobs.attack_paths.scan.sync.create_sync_indexes")
|
|
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
|
|
@patch("tasks.jobs.attack_paths.scan.findings.analysis")
|
|
@patch("tasks.jobs.attack_paths.scan.findings.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.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",
|
|
return_value="bolt://neo4j",
|
|
)
|
|
@patch(
|
|
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
|
|
return_value=MagicMock(_enabled_regions=["us-east-1"]),
|
|
)
|
|
@patch(
|
|
"tasks.jobs.attack_paths.scan.rls_transaction",
|
|
new=lambda *args, **kwargs: nullcontext(),
|
|
)
|
|
def test_run_success_flow(
|
|
self,
|
|
mock_init_provider,
|
|
mock_get_uri,
|
|
mock_create_db,
|
|
mock_clear_cache,
|
|
mock_cartography_indexes,
|
|
mock_cartography_analysis,
|
|
mock_cartography_ontology,
|
|
mock_findings_indexes,
|
|
mock_findings_analysis,
|
|
mock_internet_analysis,
|
|
mock_sync_indexes,
|
|
mock_drop_subgraph,
|
|
mock_sync,
|
|
mock_starting,
|
|
mock_update_progress,
|
|
mock_finish,
|
|
mock_get_old_scans,
|
|
mock_event_loop,
|
|
mock_drop_db,
|
|
tenants_fixture,
|
|
providers_fixture,
|
|
scans_fixture,
|
|
):
|
|
tenant = tenants_fixture[0]
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
scan = scans_fixture[0]
|
|
scan.provider = provider
|
|
scan.save()
|
|
|
|
attack_paths_scan = AttackPathsScan.objects.create(
|
|
tenant_id=tenant.id,
|
|
provider=provider,
|
|
scan=scan,
|
|
state=StateChoices.SCHEDULED,
|
|
)
|
|
|
|
mock_session = MagicMock()
|
|
session_ctx = MagicMock()
|
|
session_ctx.__enter__.return_value = mock_session
|
|
session_ctx.__exit__.return_value = False
|
|
ingestion_result = {"organizations": "warning"}
|
|
ingestion_fn = MagicMock(return_value=ingestion_result)
|
|
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
|
|
side_effect=["db-scan-id", "tenant-db"],
|
|
) as mock_get_db_name,
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.graph_database.get_session",
|
|
return_value=session_ctx,
|
|
) as mock_get_session,
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
|
|
return_value=attack_paths_scan,
|
|
) as mock_retrieve_scan,
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
|
|
return_value=ingestion_fn,
|
|
) as mock_get_ingestion,
|
|
):
|
|
result = attack_paths_run(str(tenant.id), str(scan.id), "task-123")
|
|
|
|
assert result == ingestion_result
|
|
mock_retrieve_scan.assert_called_once_with(str(tenant.id), str(scan.id))
|
|
mock_starting.assert_called_once()
|
|
config = mock_starting.call_args[0][2]
|
|
assert config.neo4j_database == "tenant-db"
|
|
mock_get_db_name.assert_has_calls(
|
|
[call(attack_paths_scan.id, temporary=True), call(provider.tenant_id)]
|
|
)
|
|
|
|
mock_create_db.assert_has_calls([call("db-scan-id"), call("tenant-db")])
|
|
mock_get_session.assert_has_calls([call("db-scan-id"), call("tenant-db")])
|
|
assert mock_cartography_indexes.call_count == 2
|
|
mock_findings_indexes.assert_has_calls([call(mock_session), call(mock_session)])
|
|
mock_sync_indexes.assert_called_once_with(mock_session)
|
|
# These use tmp_cartography_config (neo4j_database="db-scan-id")
|
|
mock_cartography_analysis.assert_called_once()
|
|
mock_cartography_ontology.assert_called_once()
|
|
mock_internet_analysis.assert_called_once()
|
|
mock_findings_analysis.assert_called_once()
|
|
mock_drop_subgraph.assert_called_once_with(
|
|
database="tenant-db",
|
|
provider_id=str(provider.id),
|
|
)
|
|
mock_sync.assert_called_once_with(
|
|
source_database="db-scan-id",
|
|
target_database="tenant-db",
|
|
provider_id=str(provider.id),
|
|
)
|
|
mock_get_ingestion.assert_called_once_with(provider.provider)
|
|
mock_event_loop.assert_called_once()
|
|
mock_update_progress.assert_any_call(attack_paths_scan, 1)
|
|
mock_update_progress.assert_any_call(attack_paths_scan, 2)
|
|
mock_update_progress.assert_any_call(attack_paths_scan, 95)
|
|
mock_update_progress.assert_any_call(attack_paths_scan, 97)
|
|
mock_update_progress.assert_any_call(attack_paths_scan, 98)
|
|
mock_update_progress.assert_any_call(attack_paths_scan, 99)
|
|
mock_finish.assert_called_once_with(
|
|
attack_paths_scan, StateChoices.COMPLETED, ingestion_result
|
|
)
|
|
|
|
def test_run_failure_marks_scan_failed(
|
|
self, tenants_fixture, providers_fixture, scans_fixture
|
|
):
|
|
tenant = tenants_fixture[0]
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
scan = scans_fixture[0]
|
|
scan.provider = provider
|
|
scan.save()
|
|
|
|
attack_paths_scan = AttackPathsScan.objects.create(
|
|
tenant_id=tenant.id,
|
|
provider=provider,
|
|
scan=scan,
|
|
state=StateChoices.SCHEDULED,
|
|
)
|
|
|
|
mock_session = MagicMock()
|
|
session_ctx = MagicMock()
|
|
session_ctx.__enter__.return_value = mock_session
|
|
session_ctx.__exit__.return_value = False
|
|
ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom"))
|
|
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.rls_transaction",
|
|
new=lambda *args, **kwargs: nullcontext(),
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
|
|
return_value=MagicMock(_enabled_regions=["us-east-1"]),
|
|
),
|
|
patch("tasks.jobs.attack_paths.scan.graph_database.get_uri"),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
|
|
return_value="db-scan-id",
|
|
),
|
|
patch("tasks.jobs.attack_paths.scan.graph_database.create_database"),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.graph_database.get_session",
|
|
return_value=session_ctx,
|
|
),
|
|
patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run"),
|
|
patch("tasks.jobs.attack_paths.scan.cartography_analysis.run"),
|
|
patch("tasks.jobs.attack_paths.scan.findings.create_findings_indexes"),
|
|
patch("tasks.jobs.attack_paths.scan.internet.analysis"),
|
|
patch("tasks.jobs.attack_paths.scan.findings.analysis"),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
|
|
return_value=attack_paths_scan,
|
|
),
|
|
patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan"),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress"
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan"
|
|
) as mock_finish,
|
|
patch("tasks.jobs.attack_paths.scan.graph_database.drop_database"),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
|
|
return_value=ingestion_fn,
|
|
),
|
|
patch(
|
|
"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.utils.stringify_exception",
|
|
return_value="Cartography failed: ingestion boom",
|
|
),
|
|
):
|
|
with pytest.raises(RuntimeError, match="ingestion boom"):
|
|
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
|
|
|
|
failure_args = mock_finish.call_args[0]
|
|
assert failure_args[0] is attack_paths_scan
|
|
assert failure_args[1] == StateChoices.FAILED
|
|
assert failure_args[2] == {
|
|
"global_cartography_error": "Cartography failed: ingestion boom"
|
|
}
|
|
|
|
def test_run_returns_early_for_unsupported_provider(self, tenants_fixture):
|
|
tenant = tenants_fixture[0]
|
|
provider = Provider.objects.create(
|
|
provider=Provider.ProviderChoices.GCP,
|
|
uid="gcp-account",
|
|
alias="gcp",
|
|
tenant_id=tenant.id,
|
|
)
|
|
scan = Scan.objects.create(
|
|
name="GCP Scan",
|
|
provider=provider,
|
|
trigger=Scan.TriggerChoices.MANUAL,
|
|
state=StateChoices.AVAILABLE,
|
|
tenant_id=tenant.id,
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.rls_transaction",
|
|
new=lambda *args, **kwargs: nullcontext(),
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
|
|
return_value=MagicMock(),
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
|
|
return_value=None,
|
|
) as mock_get_ingestion,
|
|
patch(
|
|
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan"
|
|
) as mock_retrieve,
|
|
):
|
|
mock_retrieve.return_value = None
|
|
result = attack_paths_run(str(tenant.id), str(scan.id), "task-789")
|
|
|
|
assert result == {
|
|
"global_error": "Provider gcp is not supported for Attack Paths scans"
|
|
}
|
|
mock_get_ingestion.assert_called_once_with(provider.provider)
|
|
mock_retrieve.assert_called_once_with(str(tenant.id), str(scan.id))
|
|
|
|
|
|
@pytest.mark.django_db
|
|
class TestAttackPathsFindingsHelpers:
|
|
def test_create_findings_indexes_executes_all_statements(self):
|
|
mock_session = MagicMock()
|
|
with patch("tasks.jobs.attack_paths.indexes.run_write_query") as mock_run_write:
|
|
findings_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)
|
|
mock_run_write.assert_has_calls(
|
|
[call(mock_session, stmt) for stmt in FINDINGS_INDEX_STATEMENTS]
|
|
)
|
|
|
|
def test_load_findings_batches_requests(self, providers_fixture):
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
|
|
# Create mock Finding objects with to_dict() method
|
|
mock_finding_1 = MagicMock()
|
|
mock_finding_1.to_dict.return_value = {"id": "1", "resource_uid": "r-1"}
|
|
mock_finding_2 = MagicMock()
|
|
mock_finding_2.to_dict.return_value = {"id": "2", "resource_uid": "r-2"}
|
|
|
|
# Create a generator that yields two batches of Finding instances
|
|
def findings_generator():
|
|
yield [mock_finding_1]
|
|
yield [mock_finding_2]
|
|
|
|
config = SimpleNamespace(update_tag=12345)
|
|
mock_session = MagicMock()
|
|
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.get_root_node_label",
|
|
return_value="AWSAccount",
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.get_node_uid_field",
|
|
return_value="arn",
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.get_provider_resource_label",
|
|
return_value="AWSResource",
|
|
),
|
|
):
|
|
findings_module.load_findings(
|
|
mock_session, findings_generator(), provider, config
|
|
)
|
|
|
|
assert mock_session.run.call_count == 2
|
|
for call_args in mock_session.run.call_args_list:
|
|
params = call_args.args[1]
|
|
assert params["provider_uid"] == str(provider.uid)
|
|
assert params["last_updated"] == config.update_tag
|
|
assert "findings_data" in params
|
|
|
|
def test_cleanup_findings_runs_batches(self, providers_fixture):
|
|
provider = providers_fixture[0]
|
|
config = SimpleNamespace(update_tag=1024)
|
|
mock_session = MagicMock()
|
|
|
|
first_batch = MagicMock()
|
|
first_batch.single.return_value = {"deleted_findings_count": 3}
|
|
second_batch = MagicMock()
|
|
second_batch.single.return_value = {"deleted_findings_count": 0}
|
|
mock_session.run.side_effect = [first_batch, second_batch]
|
|
|
|
findings_module.cleanup_findings(mock_session, provider, config)
|
|
|
|
assert mock_session.run.call_count == 2
|
|
params = mock_session.run.call_args.args[1]
|
|
assert params["provider_uid"] == str(provider.uid)
|
|
assert params["last_updated"] == config.update_tag
|
|
|
|
def test_stream_findings_with_resources_returns_latest_scan_data(
|
|
self,
|
|
tenants_fixture,
|
|
providers_fixture,
|
|
):
|
|
tenant = tenants_fixture[0]
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
|
|
resource = Resource.objects.create(
|
|
tenant_id=tenant.id,
|
|
provider=provider,
|
|
uid="resource-uid",
|
|
name="Resource",
|
|
region="us-east-1",
|
|
service="ec2",
|
|
type="instance",
|
|
)
|
|
|
|
older_scan = Scan.objects.create(
|
|
name="Older",
|
|
provider=provider,
|
|
trigger=Scan.TriggerChoices.MANUAL,
|
|
state=StateChoices.COMPLETED,
|
|
tenant_id=tenant.id,
|
|
)
|
|
old_finding = Finding.objects.create(
|
|
tenant_id=tenant.id,
|
|
uid="older-finding",
|
|
scan=older_scan,
|
|
delta=Finding.DeltaChoices.NEW,
|
|
status=StatusChoices.PASS,
|
|
status_extended="ok",
|
|
severity=Severity.low,
|
|
impact=Severity.low,
|
|
impact_extended="",
|
|
raw_result={},
|
|
check_id="check-old",
|
|
check_metadata={"checktitle": "Old"},
|
|
first_seen_at=older_scan.inserted_at,
|
|
)
|
|
ResourceFindingMapping.objects.create(
|
|
tenant_id=tenant.id,
|
|
resource=resource,
|
|
finding=old_finding,
|
|
)
|
|
|
|
latest_scan = Scan.objects.create(
|
|
name="Latest",
|
|
provider=provider,
|
|
trigger=Scan.TriggerChoices.MANUAL,
|
|
state=StateChoices.COMPLETED,
|
|
tenant_id=tenant.id,
|
|
)
|
|
finding = Finding.objects.create(
|
|
tenant_id=tenant.id,
|
|
uid="finding-uid",
|
|
scan=latest_scan,
|
|
delta=Finding.DeltaChoices.NEW,
|
|
status=StatusChoices.FAIL,
|
|
status_extended="failed",
|
|
severity=Severity.high,
|
|
impact=Severity.high,
|
|
impact_extended="",
|
|
raw_result={},
|
|
check_id="check-1",
|
|
check_metadata={"checktitle": "Check title"},
|
|
first_seen_at=latest_scan.inserted_at,
|
|
)
|
|
ResourceFindingMapping.objects.create(
|
|
tenant_id=tenant.id,
|
|
resource=resource,
|
|
finding=finding,
|
|
)
|
|
|
|
latest_scan.refresh_from_db()
|
|
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.rls_transaction",
|
|
new=lambda *args, **kwargs: nullcontext(),
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
|
|
"default",
|
|
),
|
|
):
|
|
# Generator yields batches, collect all findings from all batches
|
|
findings_batches = findings_module.stream_findings_with_resources(
|
|
provider,
|
|
str(latest_scan.id),
|
|
)
|
|
findings_data = []
|
|
for batch in findings_batches:
|
|
findings_data.extend(batch)
|
|
|
|
assert len(findings_data) == 1
|
|
finding_result = findings_data[0]
|
|
assert finding_result.id == str(finding.id)
|
|
assert finding_result.resource_uid == resource.uid
|
|
assert finding_result.check_title == "Check title"
|
|
assert finding_result.scan_id == str(latest_scan.id)
|
|
|
|
def test_enrich_batch_with_resources_single_resource(
|
|
self,
|
|
tenants_fixture,
|
|
providers_fixture,
|
|
):
|
|
"""One finding + one resource = one output Finding instance"""
|
|
tenant = tenants_fixture[0]
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
|
|
resource = Resource.objects.create(
|
|
tenant_id=tenant.id,
|
|
provider=provider,
|
|
uid="resource-uid-1",
|
|
name="Resource 1",
|
|
region="us-east-1",
|
|
service="ec2",
|
|
type="instance",
|
|
)
|
|
|
|
scan = Scan.objects.create(
|
|
name="Test Scan",
|
|
provider=provider,
|
|
trigger=Scan.TriggerChoices.MANUAL,
|
|
state=StateChoices.COMPLETED,
|
|
tenant_id=tenant.id,
|
|
)
|
|
|
|
finding = Finding.objects.create(
|
|
tenant_id=tenant.id,
|
|
uid="finding-uid",
|
|
scan=scan,
|
|
delta=Finding.DeltaChoices.NEW,
|
|
status=StatusChoices.FAIL,
|
|
status_extended="failed",
|
|
severity=Severity.high,
|
|
impact=Severity.high,
|
|
impact_extended="",
|
|
raw_result={},
|
|
check_id="check-1",
|
|
check_metadata={"checktitle": "Check title"},
|
|
first_seen_at=scan.inserted_at,
|
|
)
|
|
ResourceFindingMapping.objects.create(
|
|
tenant_id=tenant.id,
|
|
resource=resource,
|
|
finding=finding,
|
|
)
|
|
|
|
# Simulate the dict returned by .values()
|
|
finding_dict = {
|
|
"id": finding.id,
|
|
"uid": finding.uid,
|
|
"inserted_at": finding.inserted_at,
|
|
"updated_at": finding.updated_at,
|
|
"first_seen_at": finding.first_seen_at,
|
|
"scan_id": scan.id,
|
|
"delta": finding.delta,
|
|
"status": finding.status,
|
|
"status_extended": finding.status_extended,
|
|
"severity": finding.severity,
|
|
"check_id": finding.check_id,
|
|
"check_metadata__checktitle": finding.check_metadata["checktitle"],
|
|
"muted": finding.muted,
|
|
"muted_reason": finding.muted_reason,
|
|
}
|
|
|
|
# _enrich_batch_with_resources queries ResourceFindingMapping directly
|
|
# No RLS mock needed - test DB doesn't enforce RLS policies
|
|
with patch(
|
|
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
|
|
"default",
|
|
):
|
|
result = findings_module._enrich_batch_with_resources(
|
|
[finding_dict], str(tenant.id)
|
|
)
|
|
|
|
assert len(result) == 1
|
|
assert result[0].resource_uid == resource.uid
|
|
assert result[0].id == str(finding.id)
|
|
assert result[0].status == "FAIL"
|
|
|
|
def test_enrich_batch_with_resources_multiple_resources(
|
|
self,
|
|
tenants_fixture,
|
|
providers_fixture,
|
|
):
|
|
"""One finding + three resources = three output Finding instances"""
|
|
tenant = tenants_fixture[0]
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
|
|
resources = []
|
|
for i in range(3):
|
|
resource = Resource.objects.create(
|
|
tenant_id=tenant.id,
|
|
provider=provider,
|
|
uid=f"resource-uid-{i}",
|
|
name=f"Resource {i}",
|
|
region="us-east-1",
|
|
service="ec2",
|
|
type="instance",
|
|
)
|
|
resources.append(resource)
|
|
|
|
scan = Scan.objects.create(
|
|
name="Test Scan",
|
|
provider=provider,
|
|
trigger=Scan.TriggerChoices.MANUAL,
|
|
state=StateChoices.COMPLETED,
|
|
tenant_id=tenant.id,
|
|
)
|
|
|
|
finding = Finding.objects.create(
|
|
tenant_id=tenant.id,
|
|
uid="finding-uid",
|
|
scan=scan,
|
|
delta=Finding.DeltaChoices.NEW,
|
|
status=StatusChoices.FAIL,
|
|
status_extended="failed",
|
|
severity=Severity.high,
|
|
impact=Severity.high,
|
|
impact_extended="",
|
|
raw_result={},
|
|
check_id="check-1",
|
|
check_metadata={"checktitle": "Check title"},
|
|
first_seen_at=scan.inserted_at,
|
|
)
|
|
|
|
# Map finding to all 3 resources
|
|
for resource in resources:
|
|
ResourceFindingMapping.objects.create(
|
|
tenant_id=tenant.id,
|
|
resource=resource,
|
|
finding=finding,
|
|
)
|
|
|
|
finding_dict = {
|
|
"id": finding.id,
|
|
"uid": finding.uid,
|
|
"inserted_at": finding.inserted_at,
|
|
"updated_at": finding.updated_at,
|
|
"first_seen_at": finding.first_seen_at,
|
|
"scan_id": scan.id,
|
|
"delta": finding.delta,
|
|
"status": finding.status,
|
|
"status_extended": finding.status_extended,
|
|
"severity": finding.severity,
|
|
"check_id": finding.check_id,
|
|
"check_metadata__checktitle": finding.check_metadata["checktitle"],
|
|
"muted": finding.muted,
|
|
"muted_reason": finding.muted_reason,
|
|
}
|
|
|
|
# _enrich_batch_with_resources queries ResourceFindingMapping directly
|
|
# No RLS mock needed - test DB doesn't enforce RLS policies
|
|
with patch(
|
|
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
|
|
"default",
|
|
):
|
|
result = findings_module._enrich_batch_with_resources(
|
|
[finding_dict], str(tenant.id)
|
|
)
|
|
|
|
assert len(result) == 3
|
|
result_resource_uids = {r.resource_uid for r in result}
|
|
assert result_resource_uids == {r.uid for r in resources}
|
|
|
|
# All should have same finding data
|
|
for r in result:
|
|
assert r.id == str(finding.id)
|
|
assert r.status == "FAIL"
|
|
|
|
def test_enrich_batch_with_resources_no_resources_skips(
|
|
self,
|
|
tenants_fixture,
|
|
providers_fixture,
|
|
):
|
|
"""Finding without resources should be skipped"""
|
|
tenant = tenants_fixture[0]
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
|
|
scan = Scan.objects.create(
|
|
name="Test Scan",
|
|
provider=provider,
|
|
trigger=Scan.TriggerChoices.MANUAL,
|
|
state=StateChoices.COMPLETED,
|
|
tenant_id=tenant.id,
|
|
)
|
|
|
|
finding = Finding.objects.create(
|
|
tenant_id=tenant.id,
|
|
uid="orphan-finding",
|
|
scan=scan,
|
|
delta=Finding.DeltaChoices.NEW,
|
|
status=StatusChoices.FAIL,
|
|
status_extended="failed",
|
|
severity=Severity.high,
|
|
impact=Severity.high,
|
|
impact_extended="",
|
|
raw_result={},
|
|
check_id="check-1",
|
|
check_metadata={"checktitle": "Check title"},
|
|
first_seen_at=scan.inserted_at,
|
|
)
|
|
# Note: No ResourceFindingMapping created
|
|
|
|
finding_dict = {
|
|
"id": finding.id,
|
|
"uid": finding.uid,
|
|
"inserted_at": finding.inserted_at,
|
|
"updated_at": finding.updated_at,
|
|
"first_seen_at": finding.first_seen_at,
|
|
"scan_id": scan.id,
|
|
"delta": finding.delta,
|
|
"status": finding.status,
|
|
"status_extended": finding.status_extended,
|
|
"severity": finding.severity,
|
|
"check_id": finding.check_id,
|
|
"check_metadata__checktitle": finding.check_metadata["checktitle"],
|
|
"muted": finding.muted,
|
|
"muted_reason": finding.muted_reason,
|
|
}
|
|
|
|
# Mock logger to verify no warning is emitted
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
|
|
"default",
|
|
),
|
|
patch("tasks.jobs.attack_paths.findings.logger") as mock_logger,
|
|
):
|
|
result = findings_module._enrich_batch_with_resources(
|
|
[finding_dict], str(tenant.id)
|
|
)
|
|
|
|
assert len(result) == 0
|
|
mock_logger.warning.assert_not_called()
|
|
|
|
def test_generator_is_lazy(self, providers_fixture):
|
|
"""Generator should not execute queries until iterated"""
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
scan_id = "some-scan-id"
|
|
|
|
with (
|
|
patch("tasks.jobs.attack_paths.findings.rls_transaction") as mock_rls,
|
|
patch("tasks.jobs.attack_paths.findings.Finding") as mock_finding,
|
|
):
|
|
# Create generator but don't iterate
|
|
findings_module.stream_findings_with_resources(provider, scan_id)
|
|
|
|
# Nothing should be called yet
|
|
mock_rls.assert_not_called()
|
|
mock_finding.objects.filter.assert_not_called()
|
|
|
|
def test_load_findings_empty_generator(self, providers_fixture):
|
|
"""Empty generator should not call neo4j"""
|
|
provider = providers_fixture[0]
|
|
provider.provider = Provider.ProviderChoices.AWS
|
|
provider.save()
|
|
|
|
mock_session = MagicMock()
|
|
config = SimpleNamespace(update_tag=12345)
|
|
|
|
def empty_gen():
|
|
return
|
|
yield # Make it a generator
|
|
|
|
with (
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.get_root_node_label",
|
|
return_value="AWSAccount",
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.get_node_uid_field",
|
|
return_value="arn",
|
|
),
|
|
patch(
|
|
"tasks.jobs.attack_paths.findings.get_provider_resource_label",
|
|
return_value="AWSResource",
|
|
),
|
|
):
|
|
findings_module.load_findings(mock_session, empty_gen(), provider, config)
|
|
|
|
mock_session.run.assert_not_called()
|
|
|
|
|
|
class TestInternetAnalysis:
|
|
def _make_provider_and_config(self):
|
|
provider = MagicMock()
|
|
provider.provider = "aws"
|
|
provider.uid = "123456789012"
|
|
config = SimpleNamespace(update_tag=1234567890)
|
|
return provider, config
|
|
|
|
def test_analysis_creates_node_and_relationships(self):
|
|
"""Verify both Cypher statements are executed and relationship count returned."""
|
|
mock_session = MagicMock()
|
|
mock_result = MagicMock()
|
|
mock_result.single.return_value = {"relationships_merged": 3}
|
|
mock_session.run.side_effect = [None, mock_result]
|
|
provider, config = self._make_provider_and_config()
|
|
|
|
with patch(
|
|
"tasks.jobs.attack_paths.internet.get_root_node_label",
|
|
return_value="AWSAccount",
|
|
):
|
|
result = internet_module.analysis(mock_session, provider, config)
|
|
|
|
assert mock_session.run.call_count == 2
|
|
assert result == 3
|
|
|
|
def test_analysis_zero_exposed_resources(self):
|
|
"""When no resources are exposed, zero relationships are created."""
|
|
mock_session = MagicMock()
|
|
mock_result = MagicMock()
|
|
mock_result.single.return_value = {"relationships_merged": 0}
|
|
mock_session.run.side_effect = [None, mock_result]
|
|
provider, config = self._make_provider_and_config()
|
|
|
|
with patch(
|
|
"tasks.jobs.attack_paths.internet.get_root_node_label",
|
|
return_value="AWSAccount",
|
|
):
|
|
result = internet_module.analysis(mock_session, provider, config)
|
|
|
|
assert result == 0
|