ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-23 20:42:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(attack-paths): moved creation and loading of neo4j graph dump to scan task

Browse Source
This commit is contained in:
Josema Camacho
2026-01-13 19:23:30 +01:00
parent b59a90b4e6
commit fdc3b18542
16 changed files with 325 additions and 390 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/poetry.lock
Generated
+7 -7
View File
@@ -1258,8 +1258,8 @@ xmltodict = "*"
[package.source]
type = "git"
url = "https://github.com/prowler-cloud/cartography"
reference = "azure-mgmt-sql-3"
resolved_reference = "7c7d6d5014bf079066ee1645e9517c67fb36fe67"
reference = "master"
resolved_reference = "c134846c0db64747340f880cf0b5085f5e473e03"
[[package]]
name = "celery"
@@ -5478,8 +5478,8 @@ tzlocal = "5.3.1"
[package.source]
type = "git"
url = "https://github.com/prowler-cloud/prowler.git"
reference = "PROWLER-510-update-cartography-dependency"
resolved_reference = "221f6afbcf625e4b6334919f165a1419a1e7d5d4"
reference = "attack-paths-demo"
resolved_reference = "95d9e9a59f8e52e4096a5c17bf839e515b918239"
[[package]]
name = "psutil"
@@ -5696,7 +5696,7 @@ description = "PycURL -- A Python Interface To The cURL library"
optional = false
python-versions = ">=3.5"
groups = ["main"]
markers = "sys_platform != \"win32\" and platform_python_implementation == \"CPython\""
markers = "platform_python_implementation == \"CPython\" and sys_platform != \"win32\""
files = [
{file = "pycurl-7.45.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c31b390f1e2cd4525828f1bb78c1f825c0aab5d1588228ed71b22c4784bdb593"},
{file = "pycurl-7.45.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:942b352b69184cb26920db48e0c5cb95af39874b57dbe27318e60f1e68564e37"},
@@ -5956,7 +5956,7 @@ description = "The MSALRuntime Python Interop Package"
optional = false
python-versions = ">=3.6"
groups = ["main"]
markers = "sys_platform == \"win32\" and (platform_system == \"Windows\" or platform_system == \"Darwin\" or platform_system == \"Linux\")"
markers = "(platform_system == \"Windows\" or platform_system == \"Darwin\" or platform_system == \"Linux\") and sys_platform == \"win32\""
files = [
{file = "pymsalruntime-0.18.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:0c22e2e83faa10de422bbfaacc1bb2887c9025ee8a53f0fc2e4f7db01c4a7b66"},
{file = "pymsalruntime-0.18.1-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:8ce2944a0f944833d047bb121396091e00287e2b6373716106da86ea99abf379"},
@@ -7806,4 +7806,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.11,<3.13"
content-hash = "1c2a75921edf35350cd01ff76bc4d4316ffd4f4aba8d02b4454013d8b70a730d"
content-hash = "de57b503d0f96c22ac3f9b1e9845aefac0a5b32089c9d90928a357f991384db3"
api/pyproject.toml
+2 -2
View File
@@ -24,7 +24,7 @@ dependencies = [
"drf-spectacular-jsonapi==0.5.1",
"gunicorn==23.0.0",
"lxml==5.3.2",
"prowler @ git+https://github.com/prowler-cloud/prowler.git@PROWLER-510-update-cartography-dependency",
"prowler @ git+https://github.com/prowler-cloud/prowler.git@attack-paths-demo",
"psycopg2-binary==2.9.9",
"pytest-celery[redis] (>=1.0.1,<2.0.0)",
"sentry-sdk[django] (>=2.20.0,<3.0.0)",
@@ -37,7 +37,7 @@ dependencies = [
"matplotlib (>=3.10.6,<4.0.0)",
"reportlab (>=4.4.4,<5.0.0)",
"neo4j (<6.0.0)",
"cartography @ git+https://github.com/prowler-cloud/cartography@azure-mgmt-sql-3",
"cartography @ git+https://github.com/prowler-cloud/cartography@master",
]
description = "Prowler's API (Django/DRF)"
license = "Apache-2.0"
api/src/backend/api/attack_paths/database.py
+113 -6
View File
@@ -1,16 +1,20 @@
import logging
import os
import tempfile
import threading
from contextlib import contextmanager
from typing import Iterator
from uuid import UUID
from itertools import islice
from typing import Any, Iterator
import neo4j
import neo4j.exceptions
from cartography.intel import create_indexes as cartography_create_indexes
from django.conf import settings
from api.attack_paths.retryable_session import RetryableSession
from tasks.jobs.attack_paths import prowler as attack_paths_prowler
# Without this Celery goes crazy with Neo4j logging
logging.getLogger("neo4j").setLevel(logging.ERROR)
@@ -119,13 +123,116 @@ def drop_subgraph(database: str, root_node_label: str, root_node_id: str) -> int
return 0 # As there are no nodes to delete, the result is empty
def create_database_dump(database: str, root_node_label: str, root_node_id: str) -> str:
file_descriptor, dump_filename_path = tempfile.mkstemp(suffix=".neo4jdump")
os.close(file_descriptor)
query = """
CALL {
MATCH (rn:__ROOT_NODE_LABEL__ {id: $root_node_id})
CALL apoc.path.subgraphAll(rn, {}) YIELD nodes, relationships
RETURN nodes, relationships
}
WITH nodes, relationships
CALL {
// Chunk nodes
WITH nodes
UNWIND range(0, size(nodes) - 1, 5000) AS offset
UNWIND nodes[offset..offset + 5000] AS n
RETURN apoc.convert.toJson({
type:'node',
id:id(n),
labels:labels(n),
properties:properties(n)
}) AS line
UNION ALL
// Chunk relationships
WITH relationships
UNWIND range(0, size(relationships) - 1, 5000) AS offset
UNWIND relationships[offset..offset + 5000] AS r
RETURN apoc.convert.toJson({
type:'relationship',
id:id(r),
label:type(r),
start:id(startNode(r)),
end:id(endNode(r)),
properties:properties(r)
}) AS line
}
RETURN line
""".replace("__ROOT_NODE_LABEL__", root_node_label)
parameters = {"root_node_id": root_node_id}
with get_session(database) as neo4j_session:
with open(dump_filename_path, "w", encoding="utf-8") as f:
for record in neo4j_session.run(query, parameters):
f.write(record["line"])
f.write("\n")
return dump_filename_path
def load_database_dump(dump_filename_path: str, database: str) -> str:
BATCH_SIZE = 1000
query = """
// Read the batch
WITH [line IN $lines | apoc.convert.fromJsonMap(line)] AS rows
// Create nodes from the batch
CALL {
WITH rows
UNWIND rows AS row
WITH row
WHERE row.type = 'node'
MERGE (n {piid: row.id})
SET n += COALESCE(row.properties, {})
FOREACH (l IN COALESCE(row.labels, []) | SET n:$(l))
}
// Create relationships from the batch
CALL {
WITH rows
UNWIND rows AS row
WITH row
WHERE row.type = 'relationship'
MATCH (s {piid: row.start}), (t {piid: row.end})
CREATE (s)-[r:$(row.label)]->(t)
SET r += COALESCE(row.properties, {})
};
"""
def chunks(iterable, size):
it = iter(iterable)
while True:
batch = list(islice(it, size))
if not batch:
break
yield batch
with get_session(database) as neo4j_session:
with open(dump_filename_path, "r", encoding="utf-8") as f:
for batch in chunks(f, BATCH_SIZE):
neo4j_session.run(query, {"lines": batch}).consume()
cartography_create_indexes.run(neo4j_session, None)
attack_paths_prowler.create_indexes(neo4j_session)
os.remove(dump_filename_path)
# Neo4j functions related to Prowler + Cartography
DATABASE_NAME_TEMPLATE = "db-{attack_paths_scan_id}"
DATABASE_NAME_TEMPLATE = "db-{reference_id}" # For tenants' main databases, so it uses `tenant_id`
TEMPORAL_DATABASE_NAME_TEMPLATE = "db-tmp-{reference_id}" # For temporal databases, so it uses `attack_paths_scan_id`
def get_database_name(attack_paths_scan_id: UUID) -> str:
attack_paths_scan_id_str = str(attack_paths_scan_id).lower()
return DATABASE_NAME_TEMPLATE.format(attack_paths_scan_id=attack_paths_scan_id_str)
def get_database_name(reference_id: Any, temporal: bool = False) -> str:
lower_reference_id = str(reference_id).lower()
template = TEMPORAL_DATABASE_NAME_TEMPLATE if temporal else DATABASE_NAME_TEMPLATE
return template.format(reference_id=lower_reference_id)
# Exceptions
api/src/backend/api/attack_paths/views_helpers.py
+3 -1
View File
@@ -83,8 +83,10 @@ def execute_attack_paths_query(
definition: AttackPathsQueryDefinition,
parameters: dict[str, Any],
) -> dict[str, Any]:
tenant_database_name = graph_database.get_database_name(attack_paths_scan.tenant_id)
try:
with graph_database.get_session(attack_paths_scan.graph_database) as session:
with graph_database.get_session(tenant_database_name) as session:
result = session.run(definition.cypher, parameters)
return _serialize_graph(result.graph())
api/src/backend/api/fixtures/dev/8_dev_attack_paths_scans.json
-4
View File
@@ -9,8 +9,6 @@
"state": "completed",
"progress": 100,
"update_tag": 1693586667,
"graph_database": "db-a7f0f6de-6f8e-4b3a-8cbe-3f6dd9012345",
"is_graph_database_deleted": false,
"task": null,
"inserted_at": "2024-09-01T17:24:37Z",
"updated_at": "2024-09-01T17:44:37Z",
@@ -30,8 +28,6 @@
"state": "executing",
"progress": 48,
"update_tag": 1697625000,
"graph_database": "db-4a2fb2af-8a60-4d7d-9cae-4ca65e098765",
"is_graph_database_deleted": false,
"task": null,
"inserted_at": "2024-10-18T10:55:57Z",
"updated_at": "2024-10-18T10:56:15Z",
api/src/backend/api/migrations/0060_attack_paths_scan.py
-23
View File
@@ -59,14 +59,6 @@ class Migration(migrations.Migration):
null=True,
),
),
(
"graph_database",
models.CharField(blank=True, max_length=63, null=True),
),
(
"is_graph_database_deleted",
models.BooleanField(default=False),
),
(
"ingestion_exceptions",
models.JSONField(blank=True, default=dict, null=True),
@@ -125,21 +117,6 @@ class Migration(migrations.Migration):
fields=["tenant_id", "scan_id"],
name="aps_scan_lookup_idx",
),
models.Index(
fields=["tenant_id", "provider_id"],
name="aps_active_graph_idx",
include=["graph_database", "id"],
condition=models.Q(("is_graph_database_deleted", False)),
),
models.Index(
fields=["tenant_id", "provider_id", "-completed_at"],
name="aps_completed_graph_idx",
include=["graph_database", "id"],
condition=models.Q(
("state", "completed"),
("is_graph_database_deleted", False),
),
),
],
},
),
api/src/backend/api/models.py
-17
View File
@@ -662,8 +662,6 @@ class AttackPathsScan(RowLevelSecurityProtectedModel):
update_tag = models.BigIntegerField(
null=True, blank=True, help_text="Cartography update tag (epoch)"
)
graph_database = models.CharField(max_length=63, null=True, blank=True)
is_graph_database_deleted = models.BooleanField(default=False)
ingestion_exceptions = models.JSONField(default=dict, null=True, blank=True)
class Meta(RowLevelSecurityProtectedModel.Meta):
@@ -690,21 +688,6 @@ class AttackPathsScan(RowLevelSecurityProtectedModel):
fields=["tenant_id", "scan_id"],
name="aps_scan_lookup_idx",
),
models.Index(
fields=["tenant_id", "provider_id"],
name="aps_active_graph_idx",
include=["graph_database", "id"],
condition=Q(is_graph_database_deleted=False),
),
models.Index(
fields=["tenant_id", "provider_id", "-completed_at"],
name="aps_completed_graph_idx",
include=["graph_database", "id"],
condition=Q(
state=StateChoices.COMPLETED,
is_graph_database_deleted=False,
),
),
]
class JSONAPIMeta:
api/src/backend/api/tests/test_attack_paths.py
+44 -22
View File
@@ -10,19 +10,21 @@ from api.attack_paths import views_helpers
def test_normalize_run_payload_extracts_attributes_section():
attributes = {
"id": "aws-rds",
"parameters": {"ip": "192.0.2.0"},
}
payload = {
"data": {
"id": "ignored",
"attributes": {
"id": "aws-rds",
"parameters": {"ip": "192.0.2.0"},
},
"attributes": attributes,
}
}
result = views_helpers.normalize_run_payload(payload)
assert result == {"id": "aws-rds", "parameters": {"ip": "192.0.2.0"}}
assert result == attributes
def test_normalize_run_payload_passthrough_for_non_dict():
@@ -33,15 +35,18 @@ def test_normalize_run_payload_passthrough_for_non_dict():
def test_prepare_query_parameters_includes_provider_and_casts(
attack_paths_query_definition_factory,
):
limit = 5
uid = "123456789012"
definition = attack_paths_query_definition_factory(cast_type=int)
result = views_helpers.prepare_query_parameters(
definition,
{"limit": "5"},
provider_uid="123456789012",
{"limit": limit},
provider_uid=uid,
)
assert result["provider_uid"] == "123456789012"
assert result["limit"] == 5
assert result["provider_uid"] == uid
assert result["limit"] == limit
@pytest.mark.parametrize(
@@ -80,6 +85,12 @@ def test_prepare_query_parameters_validates_cast(
def test_execute_attack_paths_query_serializes_graph(
attack_paths_query_definition_factory, attack_paths_graph_stub_classes
):
tenant_1 = "tenant-1"
tenant_2 = "tenant-2"
node_1 = "node-1"
node_1_property_value = "value-1"
rel_1_type = "OWNS"
definition = attack_paths_query_definition_factory(
id="aws-rds",
name="RDS",
@@ -88,16 +99,16 @@ def test_execute_attack_paths_query_serializes_graph(
parameters=[],
)
parameters = {"provider_uid": "123"}
attack_paths_scan = SimpleNamespace(graph_database="tenant-db")
attack_paths_scan = SimpleNamespace(tenant_id=tenant_1)
node = attack_paths_graph_stub_classes.Node(
element_id="node-1",
element_id=node_1,
labels=["AWSAccount"],
properties={
"name": "account",
"complex": {
"items": [
attack_paths_graph_stub_classes.NativeValue("value"),
attack_paths_graph_stub_classes.NativeValue(node_1_property_value),
{"nested": 1},
]
},
@@ -105,7 +116,7 @@ def test_execute_attack_paths_query_serializes_graph(
)
relationship = attack_paths_graph_stub_classes.Relationship(
element_id="rel-1",
rel_type="OWNS",
rel_type=rel_1_type,
start_node=node,
end_node=attack_paths_graph_stub_classes.Node("node-2", ["RDSInstance"], {}),
properties={"weight": 1},
@@ -122,19 +133,26 @@ def test_execute_attack_paths_query_serializes_graph(
session_ctx.__enter__.return_value = session
session_ctx.__exit__.return_value = False
with patch(
"api.attack_paths.views_helpers.graph_database.get_session",
return_value=session_ctx,
) as mock_get_session:
with (
patch(
"api.attack_paths.views_helpers.graph_database.get_session",
return_value=session_ctx,
) as mock_get_session,
patch(
"api.attack_paths.views_helpers.graph_database.get_database_name",
return_value=tenant_2,
) as mock_get_db_name,
):
result = views_helpers.execute_attack_paths_query(
attack_paths_scan, definition, parameters
)
mock_get_session.assert_called_once_with("tenant-db")
mock_get_db_name.assert_called_once_with(tenant_1)
mock_get_session.assert_called_once_with(tenant_2)
session.run.assert_called_once_with(definition.cypher, parameters)
assert result["nodes"][0]["id"] == "node-1"
assert result["nodes"][0]["properties"]["complex"]["items"][0] == "value"
assert result["relationships"][0]["label"] == "OWNS"
assert result["nodes"][0]["id"] == node_1
assert result["nodes"][0]["properties"]["complex"]["items"][0] == node_1_property_value
assert result["relationships"][0]["label"] == rel_1_type
def test_execute_attack_paths_query_wraps_graph_errors(
@@ -147,7 +165,7 @@ def test_execute_attack_paths_query_wraps_graph_errors(
cypher="MATCH (n) RETURN n",
parameters=[],
)
attack_paths_scan = SimpleNamespace(graph_database="tenant-db")
attack_paths_scan = SimpleNamespace(tenant_id="tenant-1")
parameters = {"provider_uid": "123"}
class ExplodingContext:
@@ -162,6 +180,10 @@ def test_execute_attack_paths_query_wraps_graph_errors(
"api.attack_paths.views_helpers.graph_database.get_session",
return_value=ExplodingContext(),
),
patch(
"api.attack_paths.views_helpers.graph_database.get_database_name",
return_value="tenant-db",
),
patch("api.attack_paths.views_helpers.logger") as mock_logger,
):
with pytest.raises(APIException):
api/src/backend/api/tests/test_views.py
+30 -9
View File
@@ -3850,19 +3850,40 @@ class TestAttackPathsScanViewSet:
attack_paths_scan = create_attack_paths_scan(
provider,
scan=scans_fixture[0],
graph_database=None,
)
response = authenticated_client.post(
reverse(
"attack-paths-scans-queries-run", kwargs={"pk": attack_paths_scan.id}
query_definition = AttackPathsQueryDefinition(
id="aws-rds",
name="RDS inventory",
description="List account RDS assets",
provider=provider.provider,
cypher="MATCH (n) RETURN n",
parameters=[],
)
prepared_parameters = {"provider_uid": provider.uid}
empty_graph_payload = {"nodes": [], "relationships": []}
with (
patch("api.v1.views.get_query_by_id", return_value=query_definition),
patch(
"api.v1.views.attack_paths_views_helpers.prepare_query_parameters",
return_value=prepared_parameters,
),
data=self._run_payload(),
content_type=API_JSON_CONTENT_TYPE,
)
patch(
"api.v1.views.attack_paths_views_helpers.execute_attack_paths_query",
return_value=empty_graph_payload,
),
):
response = authenticated_client.post(
reverse(
"attack-paths-scans-queries-run",
kwargs={"pk": attack_paths_scan.id},
),
data=self._run_payload(),
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
assert "does not reference a graph database" in str(response.json())
assert response.status_code == status.HTTP_404_NOT_FOUND
def test_run_attack_paths_query_unknown_query(
self,
api/src/backend/api/v1/views.py
-147
View File
@@ -2386,15 +2386,6 @@ class AttackPathsScanViewSet(BaseRLSViewSet):
}
)
if not attack_paths_scan.graph_database:
logger.error(
f"The Attack Paths Scan {attack_paths_scan.id} does not reference a graph database"
)
return Response(
{"detail": "The Attack Paths scan does not reference a graph database"},
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
payload = attack_paths_views_helpers.normalize_run_payload(request.data)
serializer = AttackPathsQueryRunRequestSerializer(data=payload)
serializer.is_valid(raise_exception=True)
@@ -2425,144 +2416,6 @@ class AttackPathsScanViewSet(BaseRLSViewSet):
response_serializer = AttackPathsQueryResultSerializer(graph)
return Response(response_serializer.data, status=status_code)
@action(
detail=True,
methods=["get"],
url_path="export_dump",
url_name="export_dump",
)
def export_dump(self, request, pk=None):
from api.attack_paths import database as graph_database
attack_paths_scan = self.get_object()
# The right query, but it can fail with large graphs
# query = """
# MATCH (rn:AWSAccount {id: '552455647653'})
# CALL apoc.path.subgraphAll(rn, {})
# YIELD nodes, relationships
# CALL apoc.export.json.data(nodes, relationships, null, {stream: true, batchSize: 20})
# YIELD data
# RETURN data
# """
# So, let's do manual chungs, as `batchSize` is "ignored"
query = """
CALL {
MATCH (rn:AWSAccount {id: '552455647653'})
CALL apoc.path.subgraphAll(rn, {}) YIELD nodes, relationships
RETURN nodes, relationships
}
WITH nodes, relationships
CALL {
// Chunk nodes
WITH nodes
UNWIND range(0, size(nodes) - 1, 5000) AS offset
UNWIND nodes[offset..offset + 5000] AS n
RETURN apoc.convert.toJson({
type:'node',
id:id(n),
labels:labels(n),
properties:properties(n)
}) AS line
UNION ALL
// Chunk relationships
WITH relationships
UNWIND range(0, size(relationships) - 1, 5000) AS offset
UNWIND relationships[offset..offset + 5000] AS r
RETURN apoc.convert.toJson({
type:'relationship',
id:id(r),
label:type(r),
start:id(startNode(r)),
end:id(endNode(r)),
properties:properties(r)
}) AS line
}
RETURN line
"""
filepath = "/home/prowler/backend/dump.json"
with graph_database.get_session(attack_paths_scan.graph_database) as neo4j_session:
with open(filepath, "w", encoding="utf-8") as f:
for record in neo4j_session.run(query):
f.write(record["line"])
f.write("\n")
return Response(status=status.HTTP_204_NO_CONTENT)
@action(
detail=True,
methods=["get"],
url_path="import_dump",
url_name="import_dump",
)
def import_dump(self, request, pk=None):
from itertools import islice
from cartography.intel import create_indexes as cartography_create_indexes
from api.attack_paths import database as graph_database
from tasks.jobs.attack_paths import prowler as attack_paths_prowler
BATCH_SIZE = 1000
database_name = "db-tmp-import-dump"
graph_database.drop_database(database_name)
graph_database.create_database(database_name)
filepath = "/home/prowler/backend/dump.json"
query = """
// Read the batch
WITH [line IN $lines | apoc.convert.fromJsonMap(line)] AS rows
// Create nodes from the batch
CALL {
WITH rows
UNWIND rows AS row
WITH row
WHERE row.type = 'node'
MERGE (n {piid: row.id})
SET n += COALESCE(row.properties, {})
FOREACH (l IN COALESCE(row.labels, []) | SET n:$(l))
}
// Create relationships from the batch
CALL {
WITH rows
UNWIND rows AS row
WITH row
WHERE row.type = 'relationship'
MATCH (s {piid: row.start}), (t {piid: row.end})
CREATE (s)-[r:$(row.label)]->(t)
SET r += COALESCE(row.properties, {})
};
"""
def chunks(iterable, size):
it = iter(iterable)
while True:
batch = list(islice(it, size))
if not batch:
break
yield batch
with graph_database.get_session(database_name) as neo4j_session:
with open(filepath, "r", encoding="utf-8") as f:
for batch in chunks(f, BATCH_SIZE):
neo4j_session.run(query, {"lines": batch}).consume()
cartography_create_indexes.run(neo4j_session, None)
attack_paths_prowler.create_indexes(neo4j_session)
return Response(status=status.HTTP_204_NO_CONTENT)
@extend_schema_view(
list=extend_schema(
api/src/backend/conftest.py
+3 -1
View File
@@ -1502,8 +1502,10 @@ def create_attack_paths_scan():
"scan": scan_instance,
"state": state,
"progress": progress,
"graph_database": graph_database,
}
# `graph_database` is no longer a model field; keep accepting it for backward compatibility
# in tests but avoid passing it to the model constructor.
extra_fields.pop("graph_database", None)
payload.update(extra_fields)
return AttackPathsScan.objects.create(**payload)
api/src/backend/tasks/jobs/attack_paths/db_utils.py
-50
View File
@@ -1,6 +1,5 @@
from datetime import datetime, timezone
from typing import Any
from uuid import UUID
from cartography.config import Config as CartographyConfig
@@ -63,7 +62,6 @@ def starting_attack_paths_scan(
attack_paths_scan.state = StateChoices.EXECUTING
attack_paths_scan.started_at = datetime.now(tz=timezone.utc)
attack_paths_scan.update_tag = cartography_config.update_tag
attack_paths_scan.graph_database = cartography_config.neo4j_database
attack_paths_scan.save(
update_fields=[
@@ -71,7 +69,6 @@ def starting_attack_paths_scan(
"state",
"started_at",
"update_tag",
"graph_database",
]
)
@@ -109,50 +106,3 @@ def update_attack_paths_scan_progress(
with rls_transaction(attack_paths_scan.tenant_id):
attack_paths_scan.progress = progress
attack_paths_scan.save(update_fields=["progress"])
def get_old_attack_paths_scans(
tenant_id: str,
provider_id: str,
attack_paths_scan_id: str,
) -> list[ProwlerAPIAttackPathsScan]:
"""
An `old_attack_paths_scan` is any `completed` Attack Paths scan for the same provider,
with its graph database not deleted, excluding the current Attack Paths scan.
"""
with rls_transaction(tenant_id):
completed_scans_qs = (
ProwlerAPIAttackPathsScan.objects.filter(
provider_id=provider_id,
state=StateChoices.COMPLETED,
is_graph_database_deleted=False,
)
.exclude(id=attack_paths_scan_id)
.all()
)
return list(completed_scans_qs)
def update_old_attack_paths_scan(
old_attack_paths_scan: ProwlerAPIAttackPathsScan,
) -> None:
with rls_transaction(old_attack_paths_scan.tenant_id):
old_attack_paths_scan.is_graph_database_deleted = True
old_attack_paths_scan.save(update_fields=["is_graph_database_deleted"])
def get_provider_graph_database_names(tenant_id: str, provider_id: str) -> list[str]:
"""
Return existing graph database names for a tenant/provider.
Note: For accesing the `AttackPathsScan` we need to use `all_objects` manager because the provider is soft-deleted.
"""
with rls_transaction(tenant_id):
graph_databases_names_qs = ProwlerAPIAttackPathsScan.all_objects.filter(
provider_id=provider_id,
is_graph_database_deleted=False,
).values_list("graph_database", flat=True)
return list(graph_databases_names_qs)
api/src/backend/tasks/jobs/attack_paths/scan.py
+36 -20
View File
@@ -17,7 +17,7 @@ from api.models import (
StateChoices,
)
from api.utils import initialize_prowler_provider
from tasks.jobs.attack_paths import aws, db_utils, prowler, utils
from tasks.jobs.attack_paths import aws, db_utils, providers, prowler, utils
# Without this Celery goes crazy with Cartography logging
logging.getLogger("cartography").setLevel(logging.ERROR)
@@ -63,7 +63,7 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
# While creating the Cartography configuration, attributes `neo4j_user` and `neo4j_password` are not really needed in this config object
cartography_config = CartographyConfig(
neo4j_uri=graph_database.get_uri(),
neo4j_database=graph_database.get_database_name(attack_paths_scan.id),
neo4j_database=graph_database.get_database_name(attack_paths_scan.id, temporal=True),
update_tag=int(time.time()),
)
@@ -102,36 +102,20 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
# Post-processing: Just keeping it to be more Cartography compliant
cartography_ontology.run(neo4j_session, cartography_config)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 95)
cartography_analysis.run(neo4j_session, cartography_config)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 96)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 95)
# Adding Prowler nodes and relationships
prowler.analysis(
neo4j_session, prowler_api_provider, scan_id, cartography_config
)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 96)
logger.info(
f"Completed Cartography ({attack_paths_scan.id}) for "
f"{prowler_api_provider.provider.upper()} provider {prowler_api_provider.id}"
)
# Handling databases changes
old_attack_paths_scans = db_utils.get_old_attack_paths_scans(
prowler_api_provider.tenant_id,
prowler_api_provider.id,
attack_paths_scan.id,
)
for old_attack_paths_scan in old_attack_paths_scans:
graph_database.drop_database(old_attack_paths_scan.graph_database)
db_utils.update_old_attack_paths_scan(old_attack_paths_scan)
db_utils.finish_attack_paths_scan(
attack_paths_scan, StateChoices.COMPLETED, ingestion_exceptions
)
return ingestion_exceptions
except Exception as e:
exception_message = utils.stringify_exception(e, "Cartography failed")
logger.error(exception_message)
@@ -144,6 +128,38 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
)
raise
# Variables for database transfer
tenant_database = graph_database.get_database_name(tenant_id)
root_node_label = providers.get_root_node_label(prowler_api_provider.provider)
root_node_id = str(prowler_api_provider.uid)
# Create dump from temporal database
dump_filename_path = graph_database.create_database_dump(
cartography_config.neo4j_database, root_node_label, root_node_id
)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 97)
graph_database.drop_database(cartography_config.neo4j_database)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 98)
logger.info(f"Dump created at {dump_filename_path} for Attack Paths scan {attack_paths_scan.id}")
# Load dump into tenant's main database
graph_database.create_database(tenant_database)
graph_database.drop_subgraph(tenant_database, root_node_label, root_node_id)
db_utils.update_attack_paths_scan_progress(attack_paths_scan, 99)
logger.info(f"Tenant database {tenant_database} ready, loading dump now")
graph_database.load_database_dump(dump_filename_path, tenant_database)
logger.info(f"Dump loaded into tenant database {tenant_database} for Attack Paths scan {attack_paths_scan.id}")
db_utils.finish_attack_paths_scan(
attack_paths_scan, StateChoices.COMPLETED, ingestion_exceptions
)
return ingestion_exceptions
def _call_within_event_loop(fn, *args, **kwargs):
"""
api/src/backend/tasks/jobs/deletion.py
+13 -11
View File
@@ -13,7 +13,7 @@ from api.models import (
ScanSummary,
Tenant,
)
from tasks.jobs.attack_paths.db_utils import get_provider_graph_database_names
from tasks.jobs.attack_paths import providers
logger = get_task_logger(__name__)
@@ -33,16 +33,7 @@ def delete_provider(tenant_id: str, pk: str):
Raises:
Provider.DoesNotExist: If no instance with the provided primary key exists.
"""
# Delete the Attack Paths' graph databases related to the provider
graph_database_names = get_provider_graph_database_names(tenant_id, pk)
try:
for graph_database_name in graph_database_names:
graph_database.drop_database(graph_database_name)
except graph_database.GraphDatabaseQueryException as gdb_error:
logger.error(f"Error deleting Provider databases: {gdb_error}")
raise
# Get all provider related data and delete them in batches
# Get all provider related data
with rls_transaction(tenant_id):
instance = Provider.all_objects.get(pk=pk)
deletion_steps = [
@@ -53,6 +44,7 @@ def delete_provider(tenant_id: str, pk: str):
("AttackPathsScans", AttackPathsScan.all_objects.filter(provider=instance)),
]
# Delete related data in batches
deletion_summary = {}
for step_name, queryset in deletion_steps:
try:
@@ -62,6 +54,13 @@ def delete_provider(tenant_id: str, pk: str):
logger.error(f"Error deleting {step_name}: {db_error}")
raise
# Delete the Attack Paths' graph from the tenant graph database
tenant_graph_database = graph_database.get_database_name(tenant_id)
root_node_label = providers.get_root_node_label(instance.provider)
root_node_id = str(instance.uid)
graph_database.drop_subgraph(tenant_graph_database, root_node_label, root_node_id)
# Finally, delete the provider instance itself
try:
with rls_transaction(tenant_id):
_, provider_summary = instance.delete()
@@ -90,6 +89,9 @@ def delete_tenant(pk: str):
summary = delete_provider(pk, provider.id)
deletion_summary.update(summary)
tenant_graph_database = graph_database.get_database_name(pk)
graph_database.drop_database(tenant_graph_database)
Tenant.objects.using(MainRouter.admin_db).filter(id=pk).delete()
return deletion_summary
api/src/backend/tasks/tests/test_attack_paths_scan.py
+38 -9
View File
@@ -59,7 +59,7 @@ class TestAttackPathsRun:
),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
side_effect=["temp-db", "tenant-db"],
) as mock_get_db_name,
patch(
"tasks.jobs.attack_paths.scan.graph_database.create_database"
@@ -68,12 +68,28 @@ class TestAttackPathsRun:
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
) as mock_get_session,
patch(
"tasks.jobs.attack_paths.scan.graph_database.create_database_dump",
return_value="/tmp/dump",
) as mock_create_dump,
patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_database"
) as mock_drop_db,
patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_subgraph"
) as mock_drop_subgraph,
patch(
"tasks.jobs.attack_paths.scan.graph_database.load_database_dump"
) as mock_load_dump,
patch(
"tasks.jobs.attack_paths.scan.cartography_create_indexes.run"
) as mock_cartography_indexes,
patch(
"tasks.jobs.attack_paths.scan.cartography_analysis.run"
) as mock_cartography_analysis,
patch(
"tasks.jobs.attack_paths.scan.cartography_ontology.run"
) as mock_cartography_ontology,
patch(
"tasks.jobs.attack_paths.scan.prowler.create_indexes"
) as mock_prowler_indexes,
@@ -108,13 +124,14 @@ class TestAttackPathsRun:
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 == "db-scan-id"
assert config.neo4j_database == "temp-db"
mock_create_db.assert_called_once_with("db-scan-id")
mock_get_session.assert_called_once_with("db-scan-id")
mock_create_db.assert_has_calls([call("temp-db"), call("tenant-db")])
mock_get_session.assert_called_once_with("temp-db")
mock_cartography_indexes.assert_called_once_with(mock_session, config)
mock_prowler_indexes.assert_called_once_with(mock_session)
mock_cartography_analysis.assert_called_once_with(mock_session, config)
mock_cartography_ontology.assert_called_once_with(mock_session, config)
mock_prowler_analysis.assert_called_once_with(
mock_session,
provider,
@@ -129,10 +146,19 @@ class TestAttackPathsRun:
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_finish.assert_called_once_with(
attack_paths_scan, StateChoices.COMPLETED, ingestion_result
)
mock_get_db_name.assert_called_once_with(attack_paths_scan.id)
mock_update_progress.assert_any_call(attack_paths_scan, 96)
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)
assert mock_get_db_name.call_args_list == [
call(attack_paths_scan.id, temporal=True),
call(str(tenant.id)),
]
mock_create_dump.assert_called_once_with("temp-db", "AWSAccount", str(provider.uid))
mock_drop_db.assert_called_once_with("temp-db")
mock_drop_subgraph.assert_called_once_with("tenant-db", "AWSAccount", str(provider.uid))
mock_load_dump.assert_called_once_with("/tmp/dump", "tenant-db")
def test_run_failure_marks_scan_failed(
self, tenants_fixture, providers_fixture, scans_fixture
@@ -170,15 +196,17 @@ class TestAttackPathsRun:
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",
return_value="temp-db",
),
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.graph_database.drop_database") as mock_drop_db,
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.cartography_ontology.run"),
patch("tasks.jobs.attack_paths.scan.prowler.create_indexes"),
patch("tasks.jobs.attack_paths.scan.prowler.analysis"),
patch(
@@ -214,6 +242,7 @@ class TestAttackPathsRun:
assert failure_args[2] == {
"global_cartography_error": "Cartography failed: ingestion boom"
}
mock_drop_db.assert_called_once_with("temp-db")
def test_run_returns_early_for_unsupported_provider(self, tenants_fixture):
tenant = tenants_fixture[0]
api/src/backend/tasks/tests/test_deletion.py
+36 -61
View File
@@ -5,20 +5,17 @@ import pytest
from django.core.exceptions import ObjectDoesNotExist
from api.models import Provider, Tenant
from tasks.jobs.attack_paths import providers
from tasks.jobs.deletion import delete_provider, delete_tenant
@pytest.mark.django_db
class TestDeleteProvider:
def test_delete_provider_success(self, providers_fixture):
with patch(
"tasks.jobs.deletion.get_provider_graph_database_names"
) as mock_get_provider_graph_database_names, patch(
"tasks.jobs.deletion.graph_database.drop_database"
) as mock_drop_database:
graph_db_names = ["graph-db-1", "graph-db-2"]
mock_get_provider_graph_database_names.return_value = graph_db_names
with (
patch("tasks.jobs.deletion.graph_database.get_database_name", return_value="tenant-db"),
patch("tasks.jobs.deletion.graph_database.drop_subgraph") as mock_drop_subgraph,
):
instance = providers_fixture[0]
tenant_id = str(instance.tenant_id)
result = delete_provider(tenant_id, instance.id)
@@ -27,34 +24,21 @@ class TestDeleteProvider:
with pytest.raises(ObjectDoesNotExist):
Provider.objects.get(pk=instance.id)
mock_get_provider_graph_database_names.assert_called_once_with(
tenant_id, instance.id
)
mock_drop_database.assert_has_calls(
[call(graph_db_name) for graph_db_name in graph_db_names]
mock_drop_subgraph.assert_called_once_with(
"tenant-db",
providers.get_root_node_label(instance.provider),
str(instance.uid),
)
def test_delete_provider_does_not_exist(self, tenants_fixture):
with patch(
"tasks.jobs.deletion.get_provider_graph_database_names"
) as mock_get_provider_graph_database_names, patch(
"tasks.jobs.deletion.graph_database.drop_database"
) as mock_drop_database:
graph_db_names = ["graph-db-1"]
mock_get_provider_graph_database_names.return_value = graph_db_names
with patch("tasks.jobs.deletion.graph_database.drop_subgraph") as mock_drop_subgraph:
tenant_id = str(tenants_fixture[0].id)
non_existent_pk = "babf6796-cfcc-4fd3-9dcf-88d012247645"
with pytest.raises(ObjectDoesNotExist):
delete_provider(tenant_id, non_existent_pk)
mock_get_provider_graph_database_names.assert_called_once_with(
tenant_id, non_existent_pk
)
mock_drop_database.assert_has_calls(
[call(graph_db_name) for graph_db_name in graph_db_names]
)
mock_drop_subgraph.assert_not_called()
@pytest.mark.django_db
@@ -63,24 +47,17 @@ class TestDeleteTenant:
"""
Test successful deletion of a tenant and its related data.
"""
with patch(
"tasks.jobs.deletion.get_provider_graph_database_names"
) as mock_get_provider_graph_database_names, patch(
"tasks.jobs.deletion.graph_database.drop_database"
) as mock_drop_database:
with (
patch("tasks.jobs.deletion.graph_database.get_database_name", side_effect=lambda tenant_id: f"db-{tenant_id}"),
patch("tasks.jobs.deletion.graph_database.drop_subgraph") as mock_drop_subgraph,
patch("tasks.jobs.deletion.graph_database.drop_database") as mock_drop_database,
):
tenant = tenants_fixture[0]
providers = list(Provider.objects.filter(tenant_id=tenant.id))
graph_db_names_per_provider = [
[f"graph-db-{provider.id}"] for provider in providers
]
mock_get_provider_graph_database_names.side_effect = (
graph_db_names_per_provider
)
provider_list = list(Provider.objects.filter(tenant_id=tenant.id))
# Ensure the tenant and related providers exist before deletion
assert Tenant.objects.filter(id=tenant.id).exists()
assert providers
assert provider_list
# Call the function and validate the result
deletion_summary = delete_tenant(tenant.id)
@@ -89,30 +66,28 @@ class TestDeleteTenant:
assert not Tenant.objects.filter(id=tenant.id).exists()
assert not Provider.objects.filter(tenant_id=tenant.id).exists()
expected_calls = [
call(provider.tenant_id, provider.id) for provider in providers
expected_subgraph_calls = [
call(
f"db-{tenant.id}",
providers.get_root_node_label(provider.provider),
provider.uid,
)
for provider in provider_list
]
mock_get_provider_graph_database_names.assert_has_calls(
expected_calls, any_order=True
)
assert mock_get_provider_graph_database_names.call_count == len(
expected_calls
)
expected_drop_calls = [
call(graph_db_name[0]) for graph_db_name in graph_db_names_per_provider
]
mock_drop_database.assert_has_calls(expected_drop_calls, any_order=True)
assert mock_drop_database.call_count == len(expected_drop_calls)
mock_drop_subgraph.assert_has_calls(expected_subgraph_calls, any_order=True)
assert mock_drop_subgraph.call_count == len(expected_subgraph_calls)
mock_drop_database.assert_called_once_with(f"db-{tenant.id}")
def test_delete_tenant_with_no_providers(self, tenants_fixture):
"""
Test deletion of a tenant with no related providers.
"""
with patch(
"tasks.jobs.deletion.get_provider_graph_database_names"
) as mock_get_provider_graph_database_names, patch(
"tasks.jobs.deletion.graph_database.drop_database"
) as mock_drop_database:
with (
patch("tasks.jobs.deletion.graph_database.drop_subgraph") as mock_drop_subgraph,
patch("tasks.jobs.deletion.graph_database.get_database_name", side_effect=lambda tenant_id: f"db-{tenant_id}"),
patch("tasks.jobs.deletion.graph_database.drop_database") as mock_drop_database,
):
tenant = tenants_fixture[1] # Assume this tenant has no providers
providers = Provider.objects.filter(tenant_id=tenant.id)
@@ -126,5 +101,5 @@ class TestDeleteTenant:
assert deletion_summary == {} # No providers, so empty summary
assert not Tenant.objects.filter(id=tenant.id).exists()
mock_get_provider_graph_database_names.assert_not_called()
mock_drop_database.assert_not_called()
mock_drop_subgraph.assert_not_called()
mock_drop_database.assert_called_once_with(f"db-{tenant.id}")