From 97d55d7aa7d6176e2bc1fb6bdbf0a7cc88237298 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A1n=20Jes=C3=BAs=20Pe=C3=B1a=20Rodr=C3=ADguez?= Date: Tue, 18 Feb 2025 14:15:56 +0100 Subject: [PATCH] fix: fix the batch writing when launching the CLI --- api/src/backend/tasks/tasks.py | 1 + .../outputs/compliance/compliance_output.py | 1 + prowler/lib/outputs/csv/csv.py | 2 +- prowler/lib/outputs/html/html.py | 2 +- prowler/lib/outputs/ocsf/ocsf.py | 2 +- prowler/lib/outputs/output.py | 7 ++- tests/lib/outputs/asff/asff_test.py | 2 +- tests/lib/outputs/csv/csv_test.py | 2 +- tests/lib/outputs/finding_test.py | 45 ++++++++++++++----- tests/lib/outputs/html/html_test.py | 2 +- tests/lib/outputs/ocsf/ocsf_test.py | 2 +- 11 files changed, 47 insertions(+), 21 deletions(-) diff --git a/api/src/backend/tasks/tasks.py b/api/src/backend/tasks/tasks.py index 7711916064..1425699455 100644 --- a/api/src/backend/tasks/tasks.py +++ b/api/src/backend/tasks/tasks.py @@ -255,6 +255,7 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str): findings=finding_outputs, file_path=output_directory, file_extension=config["suffix"], + from_cli=False, ) writer.close_file = is_last_batch output_writers[writer_class] = writer diff --git a/prowler/lib/outputs/compliance/compliance_output.py b/prowler/lib/outputs/compliance/compliance_output.py index 09fb746f33..bb0e49d452 100644 --- a/prowler/lib/outputs/compliance/compliance_output.py +++ b/prowler/lib/outputs/compliance/compliance_output.py @@ -33,6 +33,7 @@ class ComplianceOutput(Output): file_extension: str = "", ) -> None: self._data = [] + self.file_descriptor = None if not file_extension and file_path: self._file_extension = "".join(Path(file_path).suffixes) diff --git a/prowler/lib/outputs/csv/csv.py b/prowler/lib/outputs/csv/csv.py index dec5ab86bc..927108f84a 100644 --- a/prowler/lib/outputs/csv/csv.py +++ b/prowler/lib/outputs/csv/csv.py @@ -102,7 +102,7 @@ class CSV(Output): csv_writer.writeheader() for finding in self._data: csv_writer.writerow(finding) - if self.close_file: + if self.close_file or self._from_cli: self._file_descriptor.close() except Exception as error: logger.error( diff --git a/prowler/lib/outputs/html/html.py b/prowler/lib/outputs/html/html.py index 54b5d21121..851dbeb566 100644 --- a/prowler/lib/outputs/html/html.py +++ b/prowler/lib/outputs/html/html.py @@ -78,7 +78,7 @@ class HTML(Output): HTML.write_header(self._file_descriptor, provider, stats) for finding in self._data: self._file_descriptor.write(finding) - if self.close_file: + if self.close_file or self._from_cli: HTML.write_footer(self._file_descriptor) self._file_descriptor.close() except Exception as error: diff --git a/prowler/lib/outputs/ocsf/ocsf.py b/prowler/lib/outputs/ocsf/ocsf.py index 2ddba0ddd8..a9162115be 100644 --- a/prowler/lib/outputs/ocsf/ocsf.py +++ b/prowler/lib/outputs/ocsf/ocsf.py @@ -212,7 +212,7 @@ class OCSF(Output): logger.error( f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" ) - if self.close_file: + if self.close_file or self._from_cli: if self._file_descriptor.tell() != 1: self._file_descriptor.seek( self._file_descriptor.tell() - 1, os.SEEK_SET diff --git a/prowler/lib/outputs/output.py b/prowler/lib/outputs/output.py index da05640f3a..a33fc74e37 100644 --- a/prowler/lib/outputs/output.py +++ b/prowler/lib/outputs/output.py @@ -34,11 +34,14 @@ class Output(ABC): findings: List[Finding], file_path: str = None, file_extension: str = "", + from_cli: bool = True, ) -> None: self._data = [] self.close_file = False self.file_path = file_path - self.file_descriptor = None + self._file_descriptor = None + # This parameter is to avoid refactoring more code, the CLI does not write in batches, the API does + self._from_cli = from_cli if not file_extension and file_path: self._file_extension = "".join(Path(file_path).suffixes) @@ -48,7 +51,7 @@ class Output(ABC): if findings: self.transform(findings) - if not self.file_descriptor and file_path: + if not self._file_descriptor and file_path: self.create_file_descriptor(self.file_path) @property diff --git a/tests/lib/outputs/asff/asff_test.py b/tests/lib/outputs/asff/asff_test.py index 61386124bd..2da25bc3d7 100644 --- a/tests/lib/outputs/asff/asff_test.py +++ b/tests/lib/outputs/asff/asff_test.py @@ -577,7 +577,7 @@ class TestASFF: assert loads(content) == expected_asff def test_batch_write_data_to_file_without_findings(self): - assert not hasattr(ASFF([]), "_file_descriptor") + assert not ASFF([])._file_descriptor def test_asff_generate_status(self): assert ASFF.generate_status("PASS") == "PASSED" diff --git a/tests/lib/outputs/csv/csv_test.py b/tests/lib/outputs/csv/csv_test.py index 8a96ca3017..656060179d 100644 --- a/tests/lib/outputs/csv/csv_test.py +++ b/tests/lib/outputs/csv/csv_test.py @@ -119,7 +119,7 @@ class TestCSV: assert content == expected_csv def test_batch_write_data_to_file_without_findings(self): - assert not hasattr(CSV([]), "_file_descriptor") + assert not CSV([])._file_descriptor @pytest.fixture def mock_output_class(self): diff --git a/tests/lib/outputs/finding_test.py b/tests/lib/outputs/finding_test.py index 917dfa1062..01e493bd74 100644 --- a/tests/lib/outputs/finding_test.py +++ b/tests/lib/outputs/finding_test.py @@ -61,6 +61,7 @@ class DummyTag: self.key = key self.value = value + class DummyTags: def __init__(self, tags): self._tags = tags @@ -68,6 +69,7 @@ class DummyTags: def all(self): return self._tags + class DummyResource: def __init__(self, uid, name, region, tags): self.uid = uid @@ -75,27 +77,33 @@ class DummyResource: self.region = region self.tags = DummyTags(tags) + class DummyResources: """Simulate a collection with a first() method.""" + def __init__(self, resource): self._resource = resource def first(self): return self._resource + class DummyProvider: def __init__(self, uid): self.uid = uid + class DummyScan: def __init__(self, provider): self.provider = provider + class DummyAPIFinding: """ A dummy API finding model to simulate the database model. Attributes will be added dynamically. """ + pass @@ -518,7 +526,9 @@ class TestFinding: # Create a dummy resource with one tag tag = DummyTag("env", "prod") - resource = DummyResource(uid="res-uid-1", name="ResourceName1", region="us-east-1", tags=[tag]) + resource = DummyResource( + uid="res-uid-1", name="ResourceName1", region="us-east-1", tags=[tag] + ) resources = DummyResources(resource) # Create a dummy check_metadata dict with all required fields @@ -535,10 +545,7 @@ class TestFinding: "risk": "High risk", "relatedurl": "http://example.com", "remediation": { - "recommendation": { - "text": "Fix it", - "url": "http://fix.com" - }, + "recommendation": {"text": "Fix it", "url": "http://fix.com"}, "code": { "nativeiac": "iac_code", "terraform": "terraform_code", @@ -618,10 +625,18 @@ class TestFinding: # Create a dummy scan object with a unique_resource_count dummy_scan = SimpleNamespace(unique_resource_count=10) # Build summaries covering each severity branch. - ss1 = SimpleNamespace(_pass=1, fail=2, total=3, muted=2, severity="critical", scan=dummy_scan) - ss2 = SimpleNamespace(_pass=2, fail=0, total=2, muted=0, severity="high", scan=dummy_scan) - ss3 = SimpleNamespace(_pass=2, fail=3, total=5, muted=3, severity="medium", scan=dummy_scan) - ss4 = SimpleNamespace(_pass=3, fail=0, total=3, muted=0, severity="low", scan=dummy_scan) + ss1 = SimpleNamespace( + _pass=1, fail=2, total=3, muted=2, severity="critical", scan=dummy_scan + ) + ss2 = SimpleNamespace( + _pass=2, fail=0, total=2, muted=0, severity="high", scan=dummy_scan + ) + ss3 = SimpleNamespace( + _pass=2, fail=3, total=5, muted=3, severity="medium", scan=dummy_scan + ) + ss4 = SimpleNamespace( + _pass=3, fail=0, total=3, muted=0, severity="low", scan=dummy_scan + ) summaries = [ss1, ss2, ss3, ss4] stats = Finding._transform_findings_stats(summaries) @@ -657,8 +672,12 @@ class TestFinding: """ dummy_scan = SimpleNamespace(unique_resource_count=5) # Build summaries: one summary has fail > 0 but muted == 0 - ss1 = SimpleNamespace(_pass=1, fail=2, total=3, muted=0, severity="critical", scan=dummy_scan) - ss2 = SimpleNamespace(_pass=2, fail=1, total=3, muted=1, severity="high", scan=dummy_scan) + ss1 = SimpleNamespace( + _pass=1, fail=2, total=3, muted=0, severity="critical", scan=dummy_scan + ) + ss2 = SimpleNamespace( + _pass=2, fail=1, total=3, muted=1, severity="high", scan=dummy_scan + ) summaries = [ss1, ss2] stats = Finding._transform_findings_stats(summaries) @@ -705,7 +724,9 @@ class TestFinding: dummy_finding.check_metadata = {} # Provide a dummy resources with a minimal resource tag = DummyTag("env", "prod") - resource = DummyResource(uid="res-uid-1", name="ResourceName1", region="us-east-1", tags=[tag]) + resource = DummyResource( + uid="res-uid-1", name="ResourceName1", region="us-east-1", tags=[tag] + ) dummy_finding.resources = DummyResources(resource) with pytest.raises(KeyError): diff --git a/tests/lib/outputs/html/html_test.py b/tests/lib/outputs/html/html_test.py index f38b1c65ab..c9fe3ddba9 100644 --- a/tests/lib/outputs/html/html_test.py +++ b/tests/lib/outputs/html/html_test.py @@ -492,7 +492,7 @@ class TestHTML: assert content == get_aws_html_header(args) + pass_html_finding + html_footer def test_batch_write_data_to_file_without_findings(self): - assert not hasattr(HTML([]), "_file_descriptor") + assert not HTML([])._file_descriptor def test_write_header(self): mock_file = StringIO() diff --git a/tests/lib/outputs/ocsf/ocsf_test.py b/tests/lib/outputs/ocsf/ocsf_test.py index d272849905..e394981ce3 100644 --- a/tests/lib/outputs/ocsf/ocsf_test.py +++ b/tests/lib/outputs/ocsf/ocsf_test.py @@ -258,7 +258,7 @@ class TestOCSF: assert json.loads(content) == expected_json_output def test_batch_write_data_to_file_without_findings(self): - assert not hasattr(OCSF([]), "_file_descriptor") + assert not OCSF([])._file_descriptor def test_finding_output_cloud_pass_low_muted(self): finding_output = generate_finding_output(