chore(api): Merge API repository

This commit is contained in:
Pedro De Castro committed 2024-11-25 13:12:54 +01:00
commit 6fe8c81312
102 files changed
+30515

No files matched your search

View File
Whitespace-only changes.
+31
View File
@@ -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,
)
+41
View File
@@ -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}
+25
View File
@@ -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()
+323
View File
@@ -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
+112
View File
@@ -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)
+360
View File
@@ -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)