feat(Scan, Resource, Finding): PRWLR-4015 Implement on demand scans and findings/resources in database (#48)

* feat(Users): PRWLR-4718 make user email the default login username

* feat(Token): PRWLR-4718 add serializers, views and urls for access and refresh tokens

* feat(Token): PRWLR-4718 add first membership tenant in token if not present in json body

* feat(Users): PRWLR-4718 add company_name to model

* feat(Users): PRWLR-4718 create tenant and membership when creating new user

* fix(BaseView): PRWLR-4718 add tenant_id to serializer context

* fix(Tests): PRWLR-4718 use authorization with unit tests

* fix(Views): PRWLR-4718 fix tenant retrieval from request

* fix(Tests): PRWLR-4718 fix tests

* fix(Fixtures): PRWLR-4718 fix tenant memberships ordering

* chore(Tokens): PRWLR-4718 update token url

* chore(Spec): PRWLR-4718 update API spec

* feat(Tokens): PRWLR-4718 enable token refresh blacklisting

* feat(Tokens): PRWLR-4718 implement RS256 algorithm and dev valid keys

* chore(env): PRWLR-4718 update .env.example

* chore(Deps): PRWLR-4015 update prowler dep

* fix(Resources, Findings): PRWLR-4015 fix permission issues on models and migrations

* feat(Scans, Tasks, Resources, Findings): PRWLR-4015 perform whole flow when executing scans

* fix(Settings): PRWLR-4015 fix devel JWT settings

* chore(Scans, Tasks): PRWLR-4015 improve docs and responses format

* test(Scan, Provider): PRWLR-4015 fix unit tests

* chore(Environment): PRWLR-4015 fix .env.example values

* test(Scan): PRWLR-4015 add unit tests for scan task

* chore(Task): PRWLR-4015 give proper format to task result on scan perform

* feat(Provider, Scan): PRWLR-4015 add all provider types to connection check and scan

* fix(Logging): PRWLR-4015 fix API logger middleware to include tenant and user IDs

src/backend/tasks/jobs/connection.py

@@ -2,6 +2,9 @@ from datetime import datetime, timezone
 
 from celery.utils.log import get_task_logger
 from prowler.providers.aws.aws_provider import AwsProvider
+from prowler.providers.azure.azure_provider import AzureProvider
+from prowler.providers.gcp.gcp_provider import GcpProvider
+from prowler.providers.kubernetes.kubernetes_provider import KubernetesProvider
 
 from api.models import Provider
 
@@ -26,20 +29,26 @@ def check_provider_connection(provider_id: str):
     """
     provider_instance = Provider.objects.get(pk=provider_id)
     match provider_instance.provider:
+        # TODO Refactor when proper credentials are implemented
         case Provider.ProviderChoices.AWS.value:
-            # TODO Refactor when proper credentials are implemented
-            try:
-                connection_result = AwsProvider.test_connection(
-                    raise_on_exception=False
-                )
-            # TODO: Improve this exception handling when SDK exceptions are implemented
-            except Exception as e:
-                logger.warning(str(e))
-                raise e
+            prowler_provider = AwsProvider
+        case Provider.ProviderChoices.GCP.value:
+            prowler_provider = GcpProvider
+        case Provider.ProviderChoices.AZURE.value:
+            prowler_provider = AzureProvider
+        case Provider.ProviderChoices.KUBERNETES.value:
+            prowler_provider = KubernetesProvider
         case _:
             raise ValueError(
                 f"Provider type {provider_instance.provider} not supported"
             )
+    try:
+        connection_result = prowler_provider.test_connection(raise_on_exception=False)
+    except Exception as e:
+        logger.warning(
+            f"Unexpected exception checking {provider_instance.provider} provider connection: {str(e)}"
+        )
+        raise e
 
     provider_instance.connected = connection_result.is_connected
     provider_instance.connection_last_checked_at = datetime.now(tz=timezone.utc)

src/backend/tasks/jobs/scan.py

@@ -0,0 +1,146 @@
+import json
+import time
+from datetime import datetime, timezone
+
+from celery.utils.log import get_task_logger
+from prowler.lib.outputs.finding import Finding as ProwlerFinding
+from prowler.lib.scan.scan import Scan as ProwlerScan
+from prowler.providers.aws.aws_provider import AwsProvider
+from prowler.providers.azure.azure_provider import AzureProvider
+from prowler.providers.gcp.gcp_provider import GcpProvider
+from prowler.providers.kubernetes.kubernetes_provider import KubernetesProvider
+
+from api.db_utils import tenant_transaction
+from api.models import (
+    Provider,
+    Scan,
+    Finding,
+    Resource,
+    ResourceTag,
+    StatusChoices as FindingStatus,
+    StateChoices,
+)
+from api.v1.serializers import ScanTaskSerializer
+from config.celery import TaskTimeoutError
+
+logger = get_task_logger(__name__)
+
+
+def store_finding(
+    finding: ProwlerFinding,
+    tenant_id: str,
+    scan_instance: Scan,
+    resource_instance: Resource,
+) -> Finding:
+    with tenant_transaction(tenant_id):
+        finding_instance = Finding.objects.create(
+            tenant_id=tenant_id,
+            delta=Finding.DeltaChoices.NEW,
+            status=FindingStatus[finding.status.value],
+            status_extended=finding.status_extended,
+            severity=finding.severity.value,
+            impact=finding.severity.value,
+            raw_result=json.loads(finding.json()),
+            check_id=finding.check_id,
+            scan=scan_instance,
+        )
+        finding_instance.add_resources([resource_instance])
+    return finding_instance
+
+
+def store_resources(
+    finding: ProwlerFinding, tenant_id: str, provider_instance: Provider
+) -> tuple[Resource, tuple[str, str]]:
+    with tenant_transaction(tenant_id):
+        resource_instance, _ = Resource.objects.get_or_create(
+            tenant_id=tenant_id,
+            provider=provider_instance,
+            uid=finding.resource_uid,
+            region=finding.region,
+            service=finding.service_name,
+            type=finding.resource_type,
+        )
+    with tenant_transaction(tenant_id):
+        tags = [
+            ResourceTag.objects.get_or_create(
+                tenant_id=tenant_id, key=key, value=value
+            )[0]
+            for key, value in finding.resource_tags.items()
+        ]
+        resource_instance.upsert_or_delete_tags(tags=tags)
+    return resource_instance, (resource_instance.uid, resource_instance.region)
+
+
+def perform_prowler_scan(
+    tenant_id: str, scan_id: str, provider_id: str, checks_to_execute: list[str] = None
+):
+    with tenant_transaction(tenant_id):
+        exception = None
+        provider_instance = Provider.objects.get(pk=provider_id)
+        start_time = time.time()
+        unique_resources = set()
+        # Prevent race conditions
+        while not Scan.objects.filter(id=scan_id).exists():
+            if time.time() - start_time > 10:
+                raise TaskTimeoutError(
+                    f"Could not find scan with given id {scan_id} within 10 seconds"
+                )
+            time.sleep(0.1)
+        scan_instance = Scan.objects.get(pk=scan_id)
+        scan_instance.state = StateChoices.EXECUTING
+        scan_instance.started_at = datetime.now(tz=timezone.utc)
+        scan_instance.save()
+    try:
+        match provider_instance.provider:
+            case Provider.ProviderChoices.AWS.value:
+                prowler_provider = AwsProvider
+            case Provider.ProviderChoices.GCP.value:
+                prowler_provider = GcpProvider
+            case Provider.ProviderChoices.AZURE.value:
+                prowler_provider = AzureProvider
+            case Provider.ProviderChoices.KUBERNETES.value:
+                prowler_provider = KubernetesProvider
+            case _:
+                raise ValueError(
+                    f"Provider type {provider_instance.provider} not supported"
+                )
+        with tenant_transaction(tenant_id):
+            connection_status = prowler_provider.test_connection(
+                raise_on_exception=False
+            )
+            provider_instance.connected = connection_status.is_connected
+            provider_instance.connection_last_checked_at = datetime.now(tz=timezone.utc)
+            provider_instance.save()
+
+        if connection_status.is_connected is False:
+            raise ValueError(f"Provider {provider_instance.provider} is not connected")
+
+        prowler_scan = ProwlerScan(
+            provider=prowler_provider(), checks_to_execute=checks_to_execute or []
+        )
+        for progress, findings in prowler_scan.scan():
+            for finding in findings:
+                resource_instance, resource_uid_tuple = store_resources(
+                    finding, tenant_id, provider_instance
+                )
+                store_finding(finding, tenant_id, scan_instance, resource_instance)
+                unique_resources.add(resource_uid_tuple)
+            with tenant_transaction(tenant_id):
+                scan_instance.progress = progress
+                scan_instance.save()
+
+        scan_instance.state = StateChoices.COMPLETED
+    except Exception as e:
+        logger.error(f"Error performing scan {scan_id}: {e}")
+        exception = e
+        scan_instance.state = StateChoices.FAILED
+    finally:
+        with tenant_transaction(tenant_id):
+            scan_instance.duration = time.time() - start_time
+            scan_instance.completed_at = datetime.now(tz=timezone.utc)
+            scan_instance.unique_resource_count = len(unique_resources)
+            scan_instance.save()
+        if exception is not None:
+            raise exception
+        serializer = ScanTaskSerializer(instance=scan_instance)
+        return serializer.data

src/backend/tasks/tasks.py

@@ -5,6 +5,7 @@ from api.models import Provider
 from config.celery import RLSTask
 from tasks.jobs.connection import check_provider_connection
 from tasks.jobs.deletion import delete_instance
+from tasks.jobs.scan import perform_prowler_scan
 
 
 @shared_task(base=RLSTask, name="provider-connection-check")
@@ -40,3 +41,31 @@ def delete_provider_task(provider_id: str):
               including related models if cascading deletes were triggered.
     """
     return delete_instance(model=Provider, pk=provider_id)
+
+
+@shared_task(base=RLSTask, name="scan-perform")
+def perform_scan_task(
+    tenant_id: str, scan_id: str, provider_id: str, checks_to_execute: list[str] = None
+):
+    """
+    Task to perform a Prowler scan on a given provider.
+
+    This task runs a Prowler scan on the provider identified by `provider_id`
+    under the tenant identified by `tenant_id`. The scan will use the `scan_id`
+    for tracking purposes.
+
+    Args:
+        tenant_id (str): The tenant ID under which the scan is being performed.
+        scan_id (str): The ID of the scan to be performed.
+        provider_id (str): The primary key of the Provider instance to scan.
+        checks_to_execute (list[str], optional): A list of specific checks to perform during the scan. Defaults to None.
+
+    Returns:
+        dict: The result of the scan execution, typically including the status and results of the performed checks.
+    """
+    return perform_prowler_scan(
+        tenant_id=tenant_id,
+        scan_id=scan_id,
+        provider_id=provider_id,
+        checks_to_execute=checks_to_execute,
+    )

src/backend/tasks/tests/test_scan.py

@@ -0,0 +1,111 @@
+import json
+from unittest.mock import patch, MagicMock, PropertyMock
+
+import pytest
+
+from api.models import StateChoices, StatusChoices, Severity, Finding, Resource
+from tasks.jobs.scan import perform_prowler_scan
+
+
+@pytest.mark.django_db
+class TestPerformScan:
+    @patch("tasks.jobs.scan.ProwlerScan")
+    @patch("tasks.jobs.scan.AwsProvider")
+    @patch("api.db_utils.tenant_transaction")
+    def test_perform_prowler_scan_success(
+        self,
+        mock_tenant_transaction,
+        mock_aws_provider,
+        mock_prowler_scan,
+        tenants_fixture,
+        scans_fixture,
+        providers_fixture,
+    ):
+        assert len(Finding.objects.all()) == 0
+        assert len(Resource.objects.all()) == 0
+
+        tenant, _ = tenants_fixture
+        scan, *_ = scans_fixture
+        provider, *_ = providers_fixture
+
+        tenant_id = str(tenant.id)
+        scan_id = str(scan.id)
+        provider_id = str(provider.id)
+        checks_to_execute = ["check1", "check2"]
+
+        connection_status_mock = MagicMock()
+        type(connection_status_mock).is_connected = PropertyMock(return_value=True)
+        mock_aws_provider.test_connection.return_value = connection_status_mock
+        mock_aws_provider_instance = mock_aws_provider.return_value
+        mock_aws_provider_instance.test_connection.return_value = connection_status_mock
+
+        finding = MagicMock()
+        finding.status = StatusChoices.PASS
+        finding.status_extended = "test status extended"
+        finding.severity = Severity.medium
+        finding.check_id = "check1"
+        finding.json.return_value = '{"key": "value"}'
+        finding.resource_uid = "resource_uid"
+        finding.region = "region"
+        finding.service_name = "service_name"
+        finding.resource_type = "resource_type"
+        mock_prowler_scan.return_value.scan.return_value = [(100, [finding])]
+
+        perform_prowler_scan(tenant_id, scan_id, provider_id, checks_to_execute)
+
+        scan.refresh_from_db()
+        scan_finding = Finding.objects.get(scan=scan)
+        scan_resource = Resource.objects.get(provider=provider)
+
+        assert scan.tenant == tenant
+        assert scan.provider == provider
+        assert scan.state == StateChoices.COMPLETED
+        assert scan.completed_at is not None
+        assert scan.duration is not None
+        assert scan.started_at is not None
+        assert scan.unique_resource_count == 1
+        assert scan.progress == 100
+        assert scan_finding.status == finding.status
+        assert scan_finding.status_extended == finding.status_extended
+        assert scan_finding.severity == finding.severity
+        assert scan_finding.check_id == finding.check_id
+        assert scan_finding.raw_result == json.loads(finding.json())
+
+        assert scan_resource.tenant == tenant
+        assert scan_resource.uid == finding.resource_uid
+        assert scan_resource.region == finding.region
+        assert scan_resource.service == finding.service_name
+        assert scan_resource.type == finding.resource_type
+
+    @patch("tasks.jobs.scan.ProwlerScan")
+    @patch("tasks.jobs.scan.AwsProvider")
+    @patch("api.db_utils.tenant_transaction")
+    def test_perform_prowler_scan_no_connection(
+        self,
+        mock_tenant_transaction,
+        mock_aws_provider,
+        mock_prowler_scan,
+        tenants_fixture,
+        scans_fixture,
+        providers_fixture,
+    ):
+        tenant, _ = tenants_fixture
+        scan, *_ = scans_fixture
+        provider, *_ = providers_fixture
+
+        tenant_id = str(tenant.id)
+        scan_id = str(scan.id)
+        provider_id = str(provider.id)
+        checks_to_execute = ["check1", "check2"]
+
+        connection_status_mock = MagicMock()
+        type(connection_status_mock).is_connected = PropertyMock(return_value=False)
+        mock_aws_provider.test_connection.return_value = connection_status_mock
+        mock_aws_provider_instance = mock_aws_provider.return_value
+        mock_aws_provider_instance.test_connection.return_value = connection_status_mock
+
+        with pytest.raises(ValueError):
+            perform_prowler_scan(tenant_id, scan_id, provider_id, checks_to_execute)
+
+        scan.refresh_from_db()
+        assert scan.state == StateChoices.FAILED