diff --git a/api/poetry.lock b/api/poetry.lock index 88558a76d4..a6e2cffa0b 100644 --- a/api/poetry.lock +++ b/api/poetry.lock @@ -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" diff --git a/api/pyproject.toml b/api/pyproject.toml index cdcf59834a..f96fd1f54b 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -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" diff --git a/api/src/backend/api/attack_paths/database.py b/api/src/backend/api/attack_paths/database.py index 7c06fe1f13..c08461059c 100644 --- a/api/src/backend/api/attack_paths/database.py +++ b/api/src/backend/api/attack_paths/database.py @@ -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 diff --git a/api/src/backend/api/attack_paths/views_helpers.py b/api/src/backend/api/attack_paths/views_helpers.py index 7418a0302e..2545cb10f3 100644 --- a/api/src/backend/api/attack_paths/views_helpers.py +++ b/api/src/backend/api/attack_paths/views_helpers.py @@ -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()) diff --git a/api/src/backend/api/fixtures/dev/8_dev_attack_paths_scans.json b/api/src/backend/api/fixtures/dev/8_dev_attack_paths_scans.json index fdf310458a..8474c542c9 100644 --- a/api/src/backend/api/fixtures/dev/8_dev_attack_paths_scans.json +++ b/api/src/backend/api/fixtures/dev/8_dev_attack_paths_scans.json @@ -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", diff --git a/api/src/backend/api/migrations/0060_attack_paths_scan.py b/api/src/backend/api/migrations/0060_attack_paths_scan.py index 8f57039723..70a4d3770d 100644 --- a/api/src/backend/api/migrations/0060_attack_paths_scan.py +++ b/api/src/backend/api/migrations/0060_attack_paths_scan.py @@ -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), - ), - ), ], }, ), diff --git a/api/src/backend/api/models.py b/api/src/backend/api/models.py index 5f00c578e2..d28e7a3672 100644 --- a/api/src/backend/api/models.py +++ b/api/src/backend/api/models.py @@ -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: diff --git a/api/src/backend/api/tests/test_attack_paths.py b/api/src/backend/api/tests/test_attack_paths.py index 2c4e1484f8..3ee9820a7f 100644 --- a/api/src/backend/api/tests/test_attack_paths.py +++ b/api/src/backend/api/tests/test_attack_paths.py @@ -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): diff --git a/api/src/backend/api/tests/test_views.py b/api/src/backend/api/tests/test_views.py index 7ba0d7094b..662dc7cc77 100644 --- a/api/src/backend/api/tests/test_views.py +++ b/api/src/backend/api/tests/test_views.py @@ -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, diff --git a/api/src/backend/api/v1/views.py b/api/src/backend/api/v1/views.py index 9028f83130..a3064bff5f 100644 --- a/api/src/backend/api/v1/views.py +++ b/api/src/backend/api/v1/views.py @@ -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( diff --git a/api/src/backend/conftest.py b/api/src/backend/conftest.py index 20221aca2a..b267019d8f 100644 --- a/api/src/backend/conftest.py +++ b/api/src/backend/conftest.py @@ -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) diff --git a/api/src/backend/tasks/jobs/attack_paths/db_utils.py b/api/src/backend/tasks/jobs/attack_paths/db_utils.py index 2565dfa872..e76a67c149 100644 --- a/api/src/backend/tasks/jobs/attack_paths/db_utils.py +++ b/api/src/backend/tasks/jobs/attack_paths/db_utils.py @@ -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) diff --git a/api/src/backend/tasks/jobs/attack_paths/scan.py b/api/src/backend/tasks/jobs/attack_paths/scan.py index aa7bb7b756..b33639f27b 100644 --- a/api/src/backend/tasks/jobs/attack_paths/scan.py +++ b/api/src/backend/tasks/jobs/attack_paths/scan.py @@ -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): """ diff --git a/api/src/backend/tasks/jobs/deletion.py b/api/src/backend/tasks/jobs/deletion.py index 6eee63de6a..39d0eac8a7 100644 --- a/api/src/backend/tasks/jobs/deletion.py +++ b/api/src/backend/tasks/jobs/deletion.py @@ -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 diff --git a/api/src/backend/tasks/tests/test_attack_paths_scan.py b/api/src/backend/tasks/tests/test_attack_paths_scan.py index 334c02df11..a0bc0df29b 100644 --- a/api/src/backend/tasks/tests/test_attack_paths_scan.py +++ b/api/src/backend/tasks/tests/test_attack_paths_scan.py @@ -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] diff --git a/api/src/backend/tasks/tests/test_deletion.py b/api/src/backend/tasks/tests/test_deletion.py index fc90bee0e3..1fc10e1785 100644 --- a/api/src/backend/tasks/tests/test_deletion.py +++ b/api/src/backend/tasks/tests/test_deletion.py @@ -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}")