mirror of
https://github.com/prowler-cloud/prowler.git
synced 2026-10-19 09:54:16 +00:00
chore(api): Merge API repository
This commit is contained in:
commit
6fe8c81312
102 files changed
+30515
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,31 @@
|
||||
import json
|
||||
|
||||
from django.utils import timezone
|
||||
from django_celery_beat.models import PeriodicTask, IntervalSchedule
|
||||
|
||||
from api.models import Provider
|
||||
|
||||
|
||||
def schedule_provider_scan(provider_instance: Provider):
|
||||
schedule, _ = IntervalSchedule.objects.get_or_create(
|
||||
every=24,
|
||||
period=IntervalSchedule.HOURS,
|
||||
)
|
||||
|
||||
# Create a unique name for the periodic task
|
||||
task_name = f"scan-perform-scheduled-{provider_instance.id}"
|
||||
|
||||
# Schedule the task
|
||||
PeriodicTask.objects.create(
|
||||
interval=schedule,
|
||||
name=task_name,
|
||||
task="scan-perform-scheduled",
|
||||
kwargs=json.dumps(
|
||||
{
|
||||
"tenant_id": str(provider_instance.tenant_id),
|
||||
"provider_id": str(provider_instance.id),
|
||||
}
|
||||
),
|
||||
start_time=provider_instance.inserted_at + timezone.timedelta(hours=24),
|
||||
one_off=False,
|
||||
)
|
||||
@@ -0,0 +1,41 @@
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from celery.utils.log import get_task_logger
|
||||
|
||||
from api.models import Provider
|
||||
from api.utils import prowler_provider_connection_test
|
||||
|
||||
logger = get_task_logger(__name__)
|
||||
|
||||
|
||||
def check_provider_connection(provider_id: str):
|
||||
"""
|
||||
Business logic to check the connection status of a provider.
|
||||
|
||||
Args:
|
||||
provider_id (str): The primary key of the Provider instance to check.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- 'connected' (bool): Indicates whether the provider is successfully connected.
|
||||
- 'error' (str or None): The error message if the connection failed, otherwise `None`.
|
||||
|
||||
Raises:
|
||||
ValueError: If the provider type is not supported.
|
||||
Model.DoesNotExist: If the provider does not exist.
|
||||
"""
|
||||
provider_instance = Provider.objects.get(pk=provider_id)
|
||||
try:
|
||||
connection_result = prowler_provider_connection_test(provider_instance)
|
||||
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)
|
||||
provider_instance.save()
|
||||
|
||||
connection_error = f"{connection_result.error}" if connection_result.error else None
|
||||
return {"connected": connection_result.is_connected, "error": connection_error}
|
||||
@@ -0,0 +1,25 @@
|
||||
from celery.utils.log import get_task_logger
|
||||
|
||||
logger = get_task_logger(__name__)
|
||||
|
||||
|
||||
def delete_instance(model, pk: str):
|
||||
"""
|
||||
Deletes an instance of the specified model.
|
||||
|
||||
This function retrieves an instance of the provided model using its primary key
|
||||
and deletes it from the database.
|
||||
|
||||
Args:
|
||||
model (Model): The Django model class from which to delete an instance.
|
||||
pk (str): The primary key of the instance to delete.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the number of objects deleted and a dictionary
|
||||
with the count of deleted objects per model,
|
||||
including related models if applicable.
|
||||
|
||||
Raises:
|
||||
model.DoesNotExist: If no instance with the provided primary key exists.
|
||||
"""
|
||||
return model.objects.get(pk=pk).delete()
|
||||
@@ -0,0 +1,323 @@
|
||||
import time
|
||||
from copy import deepcopy
|
||||
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 api.compliance import (
|
||||
PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE,
|
||||
generate_scan_compliance,
|
||||
)
|
||||
from api.db_utils import tenant_transaction
|
||||
from api.models import (
|
||||
Provider,
|
||||
Scan,
|
||||
Finding,
|
||||
Resource,
|
||||
ResourceTag,
|
||||
StatusChoices as FindingStatus,
|
||||
StateChoices,
|
||||
ComplianceOverview,
|
||||
)
|
||||
from api.utils import initialize_prowler_provider
|
||||
from api.v1.serializers import ScanTaskSerializer
|
||||
|
||||
logger = get_task_logger(__name__)
|
||||
|
||||
|
||||
def _create_finding_delta(
|
||||
last_status: FindingStatus | None | str, new_status: FindingStatus | None
|
||||
) -> Finding.DeltaChoices:
|
||||
"""
|
||||
Determine the delta status of a finding based on its previous and current status.
|
||||
|
||||
Args:
|
||||
last_status (FindingStatus | None | str): The previous status of the finding. Can be None or a string representation.
|
||||
new_status (FindingStatus | None): The current status of the finding.
|
||||
|
||||
Returns:
|
||||
Finding.DeltaChoices: The delta status indicating if the finding is new, changed, or unchanged.
|
||||
- Returns `Finding.DeltaChoices.NEW` if `last_status` is None.
|
||||
- Returns `Finding.DeltaChoices.CHANGED` if `last_status` and `new_status` are different.
|
||||
- Returns `None` if the status hasn't changed.
|
||||
"""
|
||||
if last_status is None:
|
||||
return Finding.DeltaChoices.NEW
|
||||
return Finding.DeltaChoices.CHANGED if last_status != new_status else None
|
||||
|
||||
|
||||
def _store_resources(
|
||||
finding: ProwlerFinding, tenant_id: str, provider_instance: Provider
|
||||
) -> tuple[Resource, tuple[str, str]]:
|
||||
"""
|
||||
Store resource information from a finding, including tags, in the database.
|
||||
|
||||
Args:
|
||||
finding (ProwlerFinding): The finding object containing resource information.
|
||||
tenant_id (str): The ID of the tenant owning the resource.
|
||||
provider_instance (Provider): The provider instance associated with the resource.
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- Resource: The resource instance created or retrieved from the database.
|
||||
- tuple[str, str]: A tuple containing the resource UID and region.
|
||||
|
||||
"""
|
||||
with tenant_transaction(tenant_id):
|
||||
resource_instance, created = Resource.objects.get_or_create(
|
||||
tenant_id=tenant_id,
|
||||
provider=provider_instance,
|
||||
uid=finding.resource_uid,
|
||||
defaults={
|
||||
"region": finding.region,
|
||||
"service": finding.service_name,
|
||||
"type": finding.resource_type,
|
||||
},
|
||||
)
|
||||
|
||||
if not created:
|
||||
resource_instance.region = finding.region
|
||||
resource_instance.service = finding.service_name
|
||||
resource_instance.type = finding.resource_type
|
||||
resource_instance.save()
|
||||
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
|
||||
):
|
||||
"""
|
||||
Perform a scan using Prowler and store the findings and resources in the database.
|
||||
|
||||
Args:
|
||||
tenant_id (str): The ID of the tenant for which the scan is performed.
|
||||
scan_id (str): The ID of the scan instance.
|
||||
provider_id (str): The ID of the provider to scan.
|
||||
checks_to_execute (list[str], optional): A list of specific checks to execute. Defaults to None.
|
||||
|
||||
Returns:
|
||||
dict: Serialized data of the completed scan instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If the provider cannot be connected.
|
||||
|
||||
"""
|
||||
generate_compliance = False
|
||||
check_status_by_region = {}
|
||||
exception = None
|
||||
unique_resources = set()
|
||||
start_time = time.time()
|
||||
|
||||
with tenant_transaction(tenant_id):
|
||||
provider_instance = Provider.objects.get(pk=provider_id)
|
||||
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:
|
||||
with tenant_transaction(tenant_id):
|
||||
try:
|
||||
prowler_provider = initialize_prowler_provider(provider_instance)
|
||||
provider_instance.connected = True
|
||||
except Exception as e:
|
||||
provider_instance.connected = False
|
||||
raise ValueError(
|
||||
f"Provider {provider_instance.provider} is not connected: {e}"
|
||||
)
|
||||
finally:
|
||||
provider_instance.connection_last_checked_at = datetime.now(
|
||||
tz=timezone.utc
|
||||
)
|
||||
provider_instance.save()
|
||||
|
||||
generate_compliance = provider_instance.provider != Provider.ProviderChoices.GCP
|
||||
prowler_scan = ProwlerScan(provider=prowler_provider, checks=checks_to_execute)
|
||||
|
||||
resource_cache = {}
|
||||
tag_cache = {}
|
||||
last_status_cache = {}
|
||||
|
||||
for progress, findings in prowler_scan.scan():
|
||||
with tenant_transaction(tenant_id):
|
||||
for finding in findings:
|
||||
# Process resource
|
||||
resource_uid = finding.resource_uid
|
||||
if resource_uid not in resource_cache:
|
||||
# Get or create the resource
|
||||
resource_instance, _ = Resource.objects.get_or_create(
|
||||
tenant_id=tenant_id,
|
||||
provider=provider_instance,
|
||||
uid=resource_uid,
|
||||
defaults={
|
||||
"region": finding.region,
|
||||
"service": finding.service_name,
|
||||
"type": finding.resource_type,
|
||||
"name": finding.resource_name,
|
||||
},
|
||||
)
|
||||
resource_cache[resource_uid] = resource_instance
|
||||
else:
|
||||
resource_instance = resource_cache[resource_uid]
|
||||
|
||||
# Update resource fields if necessary
|
||||
updated_fields = []
|
||||
if resource_instance.region != finding.region:
|
||||
resource_instance.region = finding.region
|
||||
updated_fields.append("region")
|
||||
if resource_instance.service != finding.service_name:
|
||||
resource_instance.service = finding.service_name
|
||||
updated_fields.append("service")
|
||||
if resource_instance.type != finding.resource_type:
|
||||
resource_instance.type = finding.resource_type
|
||||
updated_fields.append("type")
|
||||
if updated_fields:
|
||||
resource_instance.save(update_fields=updated_fields)
|
||||
|
||||
# Update tags
|
||||
tags = []
|
||||
for key, value in finding.resource_tags.items():
|
||||
tag_key = (key, value)
|
||||
if tag_key not in tag_cache:
|
||||
tag_instance, _ = ResourceTag.objects.get_or_create(
|
||||
tenant_id=tenant_id, key=key, value=value
|
||||
)
|
||||
tag_cache[tag_key] = tag_instance
|
||||
else:
|
||||
tag_instance = tag_cache[tag_key]
|
||||
tags.append(tag_instance)
|
||||
resource_instance.upsert_or_delete_tags(tags=tags)
|
||||
|
||||
unique_resources.add(
|
||||
(resource_instance.uid, resource_instance.region)
|
||||
)
|
||||
|
||||
# Process finding
|
||||
finding_uid = finding.uid
|
||||
if finding_uid not in last_status_cache:
|
||||
most_recent_finding = (
|
||||
Finding.objects.filter(uid=finding_uid)
|
||||
.order_by("-id")
|
||||
.values("status")
|
||||
.first()
|
||||
)
|
||||
last_status = (
|
||||
most_recent_finding["status"]
|
||||
if most_recent_finding
|
||||
else None
|
||||
)
|
||||
last_status_cache[finding_uid] = last_status
|
||||
else:
|
||||
last_status = last_status_cache[finding_uid]
|
||||
|
||||
status = FindingStatus[finding.status]
|
||||
delta = _create_finding_delta(last_status, status)
|
||||
|
||||
# Create the finding
|
||||
finding_instance = Finding.objects.create(
|
||||
tenant_id=tenant_id,
|
||||
uid=finding_uid,
|
||||
delta=delta,
|
||||
check_metadata=finding.get_metadata(),
|
||||
status=status,
|
||||
status_extended=finding.status_extended,
|
||||
severity=finding.severity,
|
||||
impact=finding.severity,
|
||||
raw_result=finding.raw,
|
||||
check_id=finding.check_id,
|
||||
scan=scan_instance,
|
||||
)
|
||||
finding_instance.add_resources([resource_instance])
|
||||
|
||||
# Update compliance data if applicable
|
||||
if not generate_compliance or finding.status.value == "MUTED":
|
||||
continue
|
||||
|
||||
region_dict = check_status_by_region.setdefault(finding.region, {})
|
||||
current_status = region_dict.get(finding.check_id)
|
||||
if current_status == "FAIL":
|
||||
continue
|
||||
region_dict[finding.check_id] = finding.status.value
|
||||
|
||||
# Update scan progress
|
||||
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 generate_compliance:
|
||||
try:
|
||||
regions = prowler_provider.get_regions()
|
||||
except AttributeError:
|
||||
regions = set()
|
||||
|
||||
compliance_template = PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE[
|
||||
provider_instance.provider
|
||||
]
|
||||
compliance_overview_by_region = {
|
||||
region: deepcopy(compliance_template) for region in regions
|
||||
}
|
||||
|
||||
for region, check_status in check_status_by_region.items():
|
||||
compliance_data = compliance_overview_by_region.setdefault(
|
||||
region, deepcopy(compliance_template)
|
||||
)
|
||||
for check_name, status in check_status.items():
|
||||
generate_scan_compliance(
|
||||
compliance_data,
|
||||
provider_instance.provider,
|
||||
check_name,
|
||||
status,
|
||||
)
|
||||
|
||||
# Prepare compliance overview objects
|
||||
compliance_overview_objects = []
|
||||
for region, compliance_data in compliance_overview_by_region.items():
|
||||
for compliance_id, compliance in compliance_data.items():
|
||||
compliance_overview_objects.append(
|
||||
ComplianceOverview(
|
||||
tenant_id=tenant_id,
|
||||
scan=scan_instance,
|
||||
region=region,
|
||||
compliance_id=compliance_id,
|
||||
framework=compliance["framework"],
|
||||
version=compliance["version"],
|
||||
description=compliance["description"],
|
||||
requirements=compliance["requirements"],
|
||||
requirements_passed=compliance["requirements_status"]["passed"],
|
||||
requirements_failed=compliance["requirements_status"]["failed"],
|
||||
requirements_manual=compliance["requirements_status"]["manual"],
|
||||
total_requirements=compliance["total_requirements"],
|
||||
)
|
||||
)
|
||||
with tenant_transaction(tenant_id):
|
||||
ComplianceOverview.objects.bulk_create(compliance_overview_objects)
|
||||
|
||||
if exception is not None:
|
||||
raise exception
|
||||
|
||||
serializer = ScanTaskSerializer(instance=scan_instance)
|
||||
return serializer.data
|
||||
@@ -0,0 +1,112 @@
|
||||
from celery import shared_task
|
||||
|
||||
from api.db_utils import tenant_transaction
|
||||
from api.decorators import set_tenant
|
||||
from api.models import Provider, Scan
|
||||
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")
|
||||
@set_tenant
|
||||
def check_provider_connection_task(provider_id: str):
|
||||
"""
|
||||
Task to check the connection status of a provider.
|
||||
|
||||
Args:
|
||||
provider_id (str): The primary key of the Provider instance to check.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- 'connected' (bool): Indicates whether the provider is successfully connected.
|
||||
- 'error' (str or None): The error message if the connection failed, otherwise `None`.
|
||||
"""
|
||||
return check_provider_connection(provider_id=provider_id)
|
||||
|
||||
|
||||
@shared_task(base=RLSTask, name="provider-deletion")
|
||||
@set_tenant
|
||||
def delete_provider_task(provider_id: str):
|
||||
"""
|
||||
Task to delete a specific Provider instance.
|
||||
|
||||
Args:
|
||||
provider_id (str): The primary key of the `Provider` instance to be deleted.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- The number of instances deleted.
|
||||
- A dictionary with the count of deleted instances per model,
|
||||
including related models if cascading deletes were triggered.
|
||||
"""
|
||||
return delete_instance(model=Provider, pk=provider_id)
|
||||
|
||||
|
||||
@shared_task(base=RLSTask, name="scan-perform", queue="scans")
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@shared_task(base=RLSTask, bind=True, name="scan-perform-scheduled", queue="scans")
|
||||
def perform_scheduled_scan_task(self, tenant_id: str, provider_id: str):
|
||||
"""
|
||||
Task to perform a scheduled Prowler scan on a given provider.
|
||||
|
||||
This task creates and executes a Prowler scan for the provider identified by `provider_id`
|
||||
under the tenant identified by `tenant_id`. It initiates a new scan instance with the task ID
|
||||
for tracking purposes. This task is intended to be run on a schedule (e.g., daily) to
|
||||
automatically perform scans without manual intervention.
|
||||
|
||||
Args:
|
||||
self: The task instance (automatically passed when bind=True).
|
||||
tenant_id (str): The tenant ID under which the scan is being performed.
|
||||
provider_id (str): The primary key of the Provider instance to scan.
|
||||
|
||||
Returns:
|
||||
dict: The result of the scan execution, typically including the status and results
|
||||
of the performed checks.
|
||||
|
||||
"""
|
||||
task_id = self.request.id
|
||||
|
||||
with tenant_transaction(tenant_id):
|
||||
provider_instance = Provider.objects.get(pk=provider_id)
|
||||
|
||||
scan_instance = Scan.objects.create(
|
||||
tenant_id=tenant_id,
|
||||
name="Daily scheduled scan",
|
||||
provider=provider_instance,
|
||||
trigger=Scan.TriggerChoices.SCHEDULED,
|
||||
task_id=task_id,
|
||||
)
|
||||
|
||||
return perform_prowler_scan(
|
||||
tenant_id=tenant_id,
|
||||
scan_id=str(scan_instance.id),
|
||||
provider_id=provider_id,
|
||||
)
|
||||
Whitespace-only changes.
@@ -0,0 +1,72 @@
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from api.models import Provider
|
||||
from tasks.jobs.connection import check_provider_connection
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"provider_data",
|
||||
[
|
||||
{"provider": "aws", "uid": "123456789012", "alias": "aws"},
|
||||
],
|
||||
)
|
||||
@patch("tasks.jobs.connection.prowler_provider_connection_test")
|
||||
@pytest.mark.django_db
|
||||
def test_check_provider_connection(
|
||||
mock_provider_connection_test, tenants_fixture, provider_data
|
||||
):
|
||||
provider = Provider.objects.create(**provider_data, tenant_id=tenants_fixture[0].id)
|
||||
|
||||
mock_test_connection_result = MagicMock()
|
||||
mock_test_connection_result.is_connected = True
|
||||
|
||||
mock_provider_connection_test.return_value = mock_test_connection_result
|
||||
|
||||
check_provider_connection(
|
||||
provider_id=str(provider.id),
|
||||
)
|
||||
provider.refresh_from_db()
|
||||
|
||||
mock_provider_connection_test.assert_called_once()
|
||||
assert provider.connected is True
|
||||
assert provider.connection_last_checked_at is not None
|
||||
assert provider.connection_last_checked_at <= datetime.now(tz=timezone.utc)
|
||||
|
||||
|
||||
@patch("tasks.jobs.connection.Provider.objects.get")
|
||||
@pytest.mark.django_db
|
||||
def test_check_provider_connection_unsupported_provider(mock_provider_get):
|
||||
mock_provider_instance = MagicMock()
|
||||
mock_provider_instance.provider = "UNSUPPORTED_PROVIDER"
|
||||
mock_provider_get.return_value = mock_provider_instance
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Provider type UNSUPPORTED_PROVIDER not supported"
|
||||
):
|
||||
check_provider_connection("provider_id")
|
||||
|
||||
|
||||
@patch("tasks.jobs.connection.Provider.objects.get")
|
||||
@patch("tasks.jobs.connection.prowler_provider_connection_test")
|
||||
@pytest.mark.django_db
|
||||
def test_check_provider_connection_exception(
|
||||
mock_provider_connection_test, mock_provider_get
|
||||
):
|
||||
mock_provider_instance = MagicMock()
|
||||
mock_provider_instance.provider = Provider.ProviderChoices.AWS.value
|
||||
mock_provider_get.return_value = mock_provider_instance
|
||||
|
||||
mock_provider_connection_test.return_value = MagicMock()
|
||||
mock_provider_connection_test.return_value.is_connected = False
|
||||
mock_provider_connection_test.return_value.error = Exception()
|
||||
|
||||
result = check_provider_connection(provider_id="provider_id")
|
||||
|
||||
assert result["connected"] is False
|
||||
assert result["error"] is not None
|
||||
|
||||
mock_provider_instance.save.assert_called_once()
|
||||
assert mock_provider_instance.connected is False
|
||||
@@ -0,0 +1,22 @@
|
||||
import pytest
|
||||
from django.core.exceptions import ObjectDoesNotExist
|
||||
|
||||
from api.models import Provider
|
||||
from tasks.jobs.deletion import delete_instance
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestDeleteInstance:
|
||||
def test_delete_instance_success(self, providers_fixture):
|
||||
instance = providers_fixture[0]
|
||||
result = delete_instance(Provider, instance.id)
|
||||
|
||||
assert result
|
||||
with pytest.raises(ObjectDoesNotExist):
|
||||
Provider.objects.get(pk=instance.id)
|
||||
|
||||
def test_delete_instance_does_not_exist(self):
|
||||
non_existent_pk = "babf6796-cfcc-4fd3-9dcf-88d012247645"
|
||||
|
||||
with pytest.raises(ObjectDoesNotExist):
|
||||
delete_instance(Provider, non_existent_pk)
|
||||
@@ -0,0 +1,360 @@
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from api.models import (
|
||||
StateChoices,
|
||||
Severity,
|
||||
Finding,
|
||||
Resource,
|
||||
StatusChoices,
|
||||
Provider,
|
||||
)
|
||||
from tasks.jobs.scan import (
|
||||
perform_prowler_scan,
|
||||
_create_finding_delta,
|
||||
_store_resources,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestPerformScan:
|
||||
def test_perform_prowler_scan_success(
|
||||
self,
|
||||
tenants_fixture,
|
||||
scans_fixture,
|
||||
providers_fixture,
|
||||
):
|
||||
with (
|
||||
patch("api.db_utils.tenant_transaction"),
|
||||
patch(
|
||||
"tasks.jobs.scan.initialize_prowler_provider"
|
||||
) as mock_initialize_prowler_provider,
|
||||
patch("tasks.jobs.scan.ProwlerScan") as mock_prowler_scan_class,
|
||||
patch(
|
||||
"tasks.jobs.scan.PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE",
|
||||
new_callable=dict,
|
||||
) as mock_prowler_compliance_overview_template,
|
||||
patch(
|
||||
"api.compliance.PROWLER_CHECKS", new_callable=dict
|
||||
) as mock_prowler_checks,
|
||||
):
|
||||
# Set up the mock PROWLER_CHECKS
|
||||
mock_prowler_checks["aws"] = {
|
||||
"check1": {"compliance1"},
|
||||
"check2": {"compliance1", "compliance2"},
|
||||
}
|
||||
|
||||
# Set up the mock PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE
|
||||
mock_prowler_compliance_overview_template["aws"] = {
|
||||
"compliance1": {
|
||||
"framework": "Framework 1",
|
||||
"version": "1.0",
|
||||
"provider": "aws",
|
||||
"description": "Description of compliance1",
|
||||
"requirements": {
|
||||
"requirement1": {
|
||||
"name": "Requirement 1",
|
||||
"description": "Description of requirement 1",
|
||||
"attributes": [],
|
||||
"checks": {"check1": None, "check2": None},
|
||||
"checks_status": {
|
||||
"pass": 0,
|
||||
"fail": 0,
|
||||
"total": 2,
|
||||
},
|
||||
"status": "PASS",
|
||||
}
|
||||
},
|
||||
"requirements_status": {
|
||||
"passed": 1,
|
||||
"failed": 0,
|
||||
"manual": 0,
|
||||
},
|
||||
"total_requirements": 1,
|
||||
}
|
||||
}
|
||||
|
||||
# Ensure the database is empty
|
||||
assert Finding.objects.count() == 0
|
||||
assert Resource.objects.count() == 0
|
||||
|
||||
tenant = tenants_fixture[0]
|
||||
scan = scans_fixture[0]
|
||||
provider = providers_fixture[0]
|
||||
|
||||
# Ensure the provider type is 'aws' to match our mocks
|
||||
provider.provider = Provider.ProviderChoices.AWS
|
||||
provider.save()
|
||||
|
||||
tenant_id = str(tenant.id)
|
||||
scan_id = str(scan.id)
|
||||
provider_id = str(provider.id)
|
||||
checks_to_execute = ["check1", "check2"]
|
||||
|
||||
# Mock the findings returned by the prowler scan
|
||||
finding = MagicMock()
|
||||
finding.uid = "this_is_a_test_finding_id"
|
||||
finding.status = StatusChoices.PASS
|
||||
finding.status_extended = "test status extended"
|
||||
finding.severity = Severity.medium
|
||||
finding.check_id = "check1"
|
||||
finding.get_metadata.return_value = {"key": "value"}
|
||||
finding.resource_uid = "resource_uid"
|
||||
finding.resource_name = "resource_name"
|
||||
finding.region = "region"
|
||||
finding.service_name = "service_name"
|
||||
finding.resource_type = "resource_type"
|
||||
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
|
||||
finding.raw = {}
|
||||
|
||||
# Mock the ProwlerScan instance
|
||||
mock_prowler_scan_instance = MagicMock()
|
||||
mock_prowler_scan_instance.scan.return_value = [(100, [finding])]
|
||||
mock_prowler_scan_class.return_value = mock_prowler_scan_instance
|
||||
|
||||
# Mock prowler_provider.get_regions()
|
||||
mock_prowler_provider_instance = MagicMock()
|
||||
mock_prowler_provider_instance.get_regions.return_value = ["region"]
|
||||
mock_initialize_prowler_provider.return_value = (
|
||||
mock_prowler_provider_instance
|
||||
)
|
||||
|
||||
# Call the function under test
|
||||
perform_prowler_scan(tenant_id, scan_id, provider_id, checks_to_execute)
|
||||
|
||||
# Refresh instances from the database
|
||||
scan.refresh_from_db()
|
||||
scan_finding = Finding.objects.get(scan=scan)
|
||||
scan_resource = Resource.objects.get(provider=provider)
|
||||
|
||||
# Assertions
|
||||
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.uid == finding.uid
|
||||
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 == finding.raw
|
||||
|
||||
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
|
||||
assert scan_resource.name == finding.resource_name
|
||||
|
||||
# Assert that the resource tags have been created and associated
|
||||
tags = scan_resource.tags.all()
|
||||
assert tags.count() == 2
|
||||
tag_keys = {tag.key for tag in tags}
|
||||
tag_values = {tag.value for tag in tags}
|
||||
assert tag_keys == set(finding.resource_tags.keys())
|
||||
assert tag_values == set(finding.resource_tags.values())
|
||||
|
||||
@patch("tasks.jobs.scan.ProwlerScan")
|
||||
@patch(
|
||||
"tasks.jobs.scan.initialize_prowler_provider",
|
||||
side_effect=Exception("Connection error"),
|
||||
)
|
||||
@patch("api.db_utils.tenant_transaction")
|
||||
def test_perform_prowler_scan_no_connection(
|
||||
self,
|
||||
mock_tenant_transaction,
|
||||
mock_initialize_prowler_provider,
|
||||
mock_prowler_scan_class,
|
||||
tenants_fixture,
|
||||
scans_fixture,
|
||||
providers_fixture,
|
||||
):
|
||||
tenant = tenants_fixture[0]
|
||||
scan = scans_fixture[0]
|
||||
provider = providers_fixture[0]
|
||||
|
||||
tenant_id = str(tenant.id)
|
||||
scan_id = str(scan.id)
|
||||
provider_id = str(provider.id)
|
||||
checks_to_execute = ["check1", "check2"]
|
||||
|
||||
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
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"last_status, new_status, expected_delta",
|
||||
[
|
||||
(None, None, Finding.DeltaChoices.NEW),
|
||||
(None, StatusChoices.PASS, Finding.DeltaChoices.NEW),
|
||||
(StatusChoices.PASS, StatusChoices.PASS, None),
|
||||
(StatusChoices.PASS, StatusChoices.FAIL, Finding.DeltaChoices.CHANGED),
|
||||
(StatusChoices.FAIL, StatusChoices.PASS, Finding.DeltaChoices.CHANGED),
|
||||
],
|
||||
)
|
||||
def test_create_finding_delta(self, last_status, new_status, expected_delta):
|
||||
assert _create_finding_delta(last_status, new_status) == expected_delta
|
||||
|
||||
@patch("api.models.ResourceTag.objects.get_or_create")
|
||||
@patch("api.models.Resource.objects.get_or_create")
|
||||
@patch("api.db_utils.tenant_transaction")
|
||||
def test_store_resources_new_resource(
|
||||
self,
|
||||
mock_tenant_transaction,
|
||||
mock_get_or_create_resource,
|
||||
mock_get_or_create_tag,
|
||||
):
|
||||
tenant_id = "tenant123"
|
||||
provider_instance = MagicMock()
|
||||
provider_instance.id = "provider456"
|
||||
|
||||
finding = MagicMock()
|
||||
finding.resource_uid = "resource_uid_123"
|
||||
finding.resource_name = "resource_name"
|
||||
finding.region = "us-west-1"
|
||||
finding.service_name = "service_name"
|
||||
finding.resource_type = "resource_type"
|
||||
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
|
||||
|
||||
resource_instance = MagicMock()
|
||||
resource_instance.uid = finding.resource_uid
|
||||
resource_instance.region = finding.region
|
||||
|
||||
mock_get_or_create_resource.return_value = (resource_instance, True)
|
||||
tag_instance = MagicMock()
|
||||
mock_get_or_create_tag.return_value = (tag_instance, True)
|
||||
|
||||
resource, resource_uid_tuple = _store_resources(
|
||||
finding, tenant_id, provider_instance
|
||||
)
|
||||
|
||||
mock_get_or_create_resource.assert_called_once_with(
|
||||
tenant_id=tenant_id,
|
||||
provider=provider_instance,
|
||||
uid=finding.resource_uid,
|
||||
defaults={
|
||||
"region": finding.region,
|
||||
"service": finding.service_name,
|
||||
"type": finding.resource_type,
|
||||
},
|
||||
)
|
||||
|
||||
assert resource == resource_instance
|
||||
assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)
|
||||
resource_instance.upsert_or_delete_tags.assert_called_once()
|
||||
|
||||
@patch("api.models.ResourceTag.objects.get_or_create")
|
||||
@patch("api.models.Resource.objects.get_or_create")
|
||||
@patch("api.db_utils.tenant_transaction")
|
||||
def test_store_resources_existing_resource(
|
||||
self,
|
||||
mock_tenant_transaction,
|
||||
mock_get_or_create_resource,
|
||||
mock_get_or_create_tag,
|
||||
):
|
||||
tenant_id = "tenant123"
|
||||
provider_instance = MagicMock()
|
||||
provider_instance.id = "provider456"
|
||||
|
||||
finding = MagicMock()
|
||||
finding.resource_uid = "resource_uid_123"
|
||||
finding.resource_name = "resource_name"
|
||||
finding.region = "us-west-2"
|
||||
finding.service_name = "new_service"
|
||||
finding.resource_type = "new_type"
|
||||
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
|
||||
|
||||
resource_instance = MagicMock()
|
||||
resource_instance.uid = finding.resource_uid
|
||||
resource_instance.region = "us-west-1"
|
||||
resource_instance.service = "old_service"
|
||||
resource_instance.type = "old_type"
|
||||
|
||||
mock_get_or_create_resource.return_value = (resource_instance, False)
|
||||
|
||||
tag_instance = MagicMock()
|
||||
mock_get_or_create_tag.return_value = (tag_instance, True)
|
||||
|
||||
resource, resource_uid_tuple = _store_resources(
|
||||
finding, tenant_id, provider_instance
|
||||
)
|
||||
|
||||
mock_get_or_create_resource.assert_called_once_with(
|
||||
tenant_id=tenant_id,
|
||||
provider=provider_instance,
|
||||
uid=finding.resource_uid,
|
||||
defaults={
|
||||
"region": finding.region,
|
||||
"service": finding.service_name,
|
||||
"type": finding.resource_type,
|
||||
},
|
||||
)
|
||||
|
||||
# Check that resource fields were updated
|
||||
assert resource_instance.region == finding.region
|
||||
assert resource_instance.service == finding.service_name
|
||||
assert resource_instance.type == finding.resource_type
|
||||
resource_instance.save.assert_called_once()
|
||||
|
||||
assert resource == resource_instance
|
||||
assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)
|
||||
resource_instance.upsert_or_delete_tags.assert_called_once()
|
||||
|
||||
@patch("api.models.ResourceTag.objects.get_or_create")
|
||||
@patch("api.models.Resource.objects.get_or_create")
|
||||
@patch("api.db_utils.tenant_transaction")
|
||||
def test_store_resources_with_tags(
|
||||
self,
|
||||
mock_tenant_transaction,
|
||||
mock_get_or_create_resource,
|
||||
mock_get_or_create_tag,
|
||||
):
|
||||
tenant_id = "tenant123"
|
||||
provider_instance = MagicMock()
|
||||
provider_instance.id = "provider456"
|
||||
|
||||
finding = MagicMock()
|
||||
finding.resource_uid = "resource_uid_123"
|
||||
finding.resource_name = "resource_name"
|
||||
finding.region = "us-west-1"
|
||||
finding.service_name = "service_name"
|
||||
finding.resource_type = "resource_type"
|
||||
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
|
||||
|
||||
resource_instance = MagicMock()
|
||||
resource_instance.uid = finding.resource_uid
|
||||
resource_instance.region = finding.region
|
||||
|
||||
mock_get_or_create_resource.return_value = (resource_instance, True)
|
||||
tag_instance_1 = MagicMock()
|
||||
tag_instance_2 = MagicMock()
|
||||
mock_get_or_create_tag.side_effect = [
|
||||
(tag_instance_1, True),
|
||||
(tag_instance_2, True),
|
||||
]
|
||||
|
||||
resource, resource_uid_tuple = _store_resources(
|
||||
finding, tenant_id, provider_instance
|
||||
)
|
||||
|
||||
mock_get_or_create_tag.assert_any_call(
|
||||
tenant_id=tenant_id, key="tag1", value="value1"
|
||||
)
|
||||
mock_get_or_create_tag.assert_any_call(
|
||||
tenant_id=tenant_id, key="tag2", value="value2"
|
||||
)
|
||||
resource_instance.upsert_or_delete_tags.assert_called_once()
|
||||
tags_passed = resource_instance.upsert_or_delete_tags.call_args[1]["tags"]
|
||||
assert tag_instance_1 in tags_passed
|
||||
assert tag_instance_2 in tags_passed
|
||||
|
||||
assert resource == resource_instance
|
||||
assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)
|
||||
Reference in new issue
Block a user