From 102d0999471288263a2ca960ad091147c2b3f925 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Fern=C3=A1ndez=20Poyatos?= Date: Wed, 26 Feb 2025 16:15:37 +0100 Subject: [PATCH] feat(findings): Add Django management command to populate database with dummy data (#7049) --- api/README.md | 63 +++++ api/poetry.lock | 23 +- api/pyproject.toml | 1 + .../api/management/commands/findings.py | 237 ++++++++++++++++++ 4 files changed, 323 insertions(+), 1 deletion(-) create mode 100644 api/src/backend/api/management/commands/findings.py diff --git a/api/README.md b/api/README.md index cd38f532d5..063d8342f6 100644 --- a/api/README.md +++ b/api/README.md @@ -269,3 +269,66 @@ poetry shell cd src/backend pytest ``` + +# Custom commands + +Django provides a way to create custom commands that can be run from the command line. + +> These commands can be found in: ```prowler/api/src/backend/api/management/commands``` + +To run a custom command, you need to be in the `prowler/api/src/backend` directory and run: + +```console +poetry shell +python manage.py +``` + +## Generate dummy data + +```console +python manage.py findings --tenant + --findings --re +sources --batch --alias +``` + +This command creates, for a given tenant, a provider, scan and a set of findings and resources related altogether. + +> Scan progress and state are updated in real time. +> - 0-33%: Create resources. +> - 33-66%: Create findings. +> - 66%: Create resource-finding mapping. +> +> The last step is required to access the findings details, since the UI needs that to print all the information. + +### Example + +```console +~/backend $ poetry run python manage.py findings --tenant +fffb1893-3fc7-4623-a5d9-fae47da1c528 --findings 25000 --re +sources 1000 --batch 5000 --alias test-script + +Starting data population + Tenant: fffb1893-3fc7-4623-a5d9-fae47da1c528 + Alias: test-script + Resources: 1000 + Findings: 25000 + Batch size: 5000 + + +Creating resources... +100%|███████████████████████| 1/1 [00:00<00:00, 7.72it/s] +Resources created successfully. + + +Creating findings... +100%|███████████████████████| 5/5 [00:05<00:00, 1.09s/it] +Findings created successfully. + + +Creating resource-finding mappings... +100%|███████████████████████| 5/5 [00:02<00:00, 1.81it/s] +Resource-finding mappings created successfully. + + +Successfully populated test data. +``` diff --git a/api/poetry.lock b/api/poetry.lock index 5475c78724..42359bc639 100644 --- a/api/poetry.lock +++ b/api/poetry.lock @@ -4776,6 +4776,27 @@ files = [ {file = "tomlkit-0.13.2.tar.gz", hash = "sha256:fff5fe59a87295b278abd31bec92c15d9bc4a06885ab12bcea52c71119392e79"}, ] +[[package]] +name = "tqdm" +version = "4.67.1" +description = "Fast, Extensible Progress Meter" +optional = false +python-versions = ">=3.7" +files = [ + {file = "tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2"}, + {file = "tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "platform_system == \"Windows\""} + +[package.extras] +dev = ["nbval", "pytest (>=6)", "pytest-asyncio (>=0.24)", "pytest-cov", "pytest-timeout"] +discord = ["requests"] +notebook = ["ipywidgets (>=6)"] +slack = ["slack-sdk"] +telegram = ["requests"] + [[package]] name = "typer" version = "0.15.1" @@ -5154,4 +5175,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.0" python-versions = ">=3.11,<3.13" -content-hash = "a8927a21c9e3006f5378b4bdf72108e0d45c2d803d59733f9824b3eaf12ff5e1" +content-hash = "3360123df5b0ad1fb39f186b801c3554fbaeea06f32765b8a32a2bb7eb87b41b" diff --git a/api/pyproject.toml b/api/pyproject.toml index 4a62cada47..e02d146b31 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -52,6 +52,7 @@ pytest-randomly = "3.15.0" pytest-xdist = "3.6.1" ruff = "0.5.0" safety = "3.2.9" +tqdm = "4.67.1" vulture = "2.14" [tool.poetry.scripts] diff --git a/api/src/backend/api/management/commands/findings.py b/api/src/backend/api/management/commands/findings.py new file mode 100644 index 0000000000..f81f8012a3 --- /dev/null +++ b/api/src/backend/api/management/commands/findings.py @@ -0,0 +1,237 @@ +import random +from datetime import datetime, timezone +from math import ceil +from uuid import uuid4 + +from django.core.management.base import BaseCommand +from tqdm import tqdm + +from api.db_utils import rls_transaction +from api.models import ( + Finding, + Provider, + Resource, + ResourceFindingMapping, + Scan, + StatusChoices, +) +from prowler.lib.check.models import CheckMetadata + + +class Command(BaseCommand): + help = "Populates the database with test data for performance testing." + + def add_arguments(self, parser): + parser.add_argument( + "--tenant", + type=str, + required=True, + help="Tenant id for which the data will be populated.", + ) + parser.add_argument( + "--resources", + type=int, + required=True, + help="The number of resources to create.", + ) + parser.add_argument( + "--findings", + type=int, + required=True, + help="The number of findings to create.", + ) + parser.add_argument( + "--batch", type=int, required=True, help="The batch size for bulk creation." + ) + parser.add_argument( + "--alias", + type=str, + required=False, + help="Optional alias for the provider and scan", + ) + + def handle(self, *args, **options): + tenant_id = options["tenant"] + num_resources = options["resources"] + num_findings = options["findings"] + batch_size = options["batch"] + alias = options["alias"] or "Testing" + uid_token = str(uuid4()) + + self.stdout.write(self.style.NOTICE("Starting data population")) + self.stdout.write(self.style.NOTICE(f"\tTenant: {tenant_id}")) + self.stdout.write(self.style.NOTICE(f"\tAlias: {alias}")) + self.stdout.write(self.style.NOTICE(f"\tResources: {num_resources}")) + self.stdout.write(self.style.NOTICE(f"\tFindings: {num_findings}")) + self.stdout.write(self.style.NOTICE(f"\tBatch size: {batch_size}\n\n")) + + # Resource metadata + possible_regions = [ + "us-east-1", + "us-east-2", + "us-west-1", + "us-west-2", + "ca-central-1", + "eu-central-1", + "eu-west-1", + "eu-west-2", + "eu-west-3", + "ap-southeast-1", + "ap-southeast-2", + "ap-northeast-1", + "ap-northeast-2", + "ap-south-1", + "sa-east-1", + ] + possible_services = [] + possible_types = [] + + bulk_check_metadata = CheckMetadata.get_bulk(provider="aws") + for check_metadata in bulk_check_metadata.values(): + if check_metadata.ServiceName not in possible_services: + possible_services.append(check_metadata.ServiceName) + if ( + check_metadata.ResourceType + and check_metadata.ResourceType not in possible_types + ): + possible_types.append(check_metadata.ResourceType) + + with rls_transaction(tenant_id): + provider, _ = Provider.all_objects.get_or_create( + tenant_id=tenant_id, + provider="aws", + connected=True, + uid=str(random.randint(100000000000, 999999999999)), + defaults={ + "alias": alias, + }, + ) + + with rls_transaction(tenant_id): + scan = Scan.all_objects.create( + tenant_id=tenant_id, + provider=provider, + name=alias, + trigger="manual", + state="executing", + progress=0, + started_at=datetime.now(timezone.utc), + ) + scan_state = "completed" + + try: + # Create resources + resources = [] + + for i in range(num_resources): + resources.append( + Resource( + tenant_id=tenant_id, + provider_id=provider.id, + uid=f"testing-{uid_token}-{i}", + name=f"Testing {uid_token}-{i}", + region=random.choice(possible_regions), + service=random.choice(possible_services), + type=random.choice(possible_types), + ) + ) + + num_batches = ceil(len(resources) / batch_size) + self.stdout.write(self.style.WARNING("Creating resources...")) + for i in tqdm(range(0, len(resources), batch_size), total=num_batches): + with rls_transaction(tenant_id): + Resource.all_objects.bulk_create(resources[i : i + batch_size]) + self.stdout.write(self.style.SUCCESS("Resources created successfully.\n\n")) + + with rls_transaction(tenant_id): + scan.progress = 33 + scan.save() + + # Create Findings + findings = [] + possible_deltas = ["new", "changed", None] + possible_severities = ["critical", "high", "medium", "low"] + findings_resources_mapping = [] + + for i in range(num_findings): + severity = random.choice(possible_severities) + check_id = random.randint(1, 1000) + assigned_resource_num = random.randint(0, len(resources) - 1) + assigned_resource = resources[assigned_resource_num] + findings_resources_mapping.append(assigned_resource_num) + + findings.append( + Finding( + tenant_id=tenant_id, + scan=scan, + uid=f"testing-{uid_token}-{i}", + delta=random.choice(possible_deltas), + check_id=f"check-{check_id}", + status=random.choice(list(StatusChoices)), + severity=severity, + impact=severity, + raw_result={}, + check_metadata={ + "checktitle": f"Test title for check {check_id}", + "risk": f"Testing risk {uid_token}-{i}", + "provider": "aws", + "severity": severity, + "categories": ["category1", "category2", "category3"], + "description": "This is a random description that should not matter for testing purposes.", + "servicename": assigned_resource.service, + "resourcetype": assigned_resource.type, + }, + ) + ) + + num_batches = ceil(len(findings) / batch_size) + self.stdout.write(self.style.WARNING("Creating findings...")) + for i in tqdm(range(0, len(findings), batch_size), total=num_batches): + with rls_transaction(tenant_id): + Finding.all_objects.bulk_create(findings[i : i + batch_size]) + self.stdout.write(self.style.SUCCESS("Findings created successfully.\n\n")) + + with rls_transaction(tenant_id): + scan.progress = 66 + scan.save() + + # Create ResourceFindingMapping + mappings = [] + for index, f in enumerate(findings): + mappings.append( + ResourceFindingMapping( + tenant_id=tenant_id, + resource=resources[findings_resources_mapping[index]], + finding=f, + ) + ) + + num_batches = ceil(len(mappings) / batch_size) + self.stdout.write( + self.style.WARNING("Creating resource-finding mappings...") + ) + for i in tqdm(range(0, len(mappings), batch_size), total=num_batches): + with rls_transaction(tenant_id): + ResourceFindingMapping.objects.bulk_create( + mappings[i : i + batch_size] + ) + self.stdout.write( + self.style.SUCCESS( + "Resource-finding mappings created successfully.\n\n" + ) + ) + except Exception as e: + self.stdout.write(self.style.ERROR(f"Failed to populate test data: {e}")) + scan_state = "failed" + finally: + scan.completed_at = datetime.now(timezone.utc) + scan.duration = int( + (datetime.now(timezone.utc) - scan.started_at).total_seconds() + ) + scan.progress = 100 + scan.state = scan_state + scan.unique_resource_count = num_resources + with rls_transaction(tenant_id): + scan.save() + + self.stdout.write(self.style.NOTICE("Successfully populated test data."))