diff --git a/prowler/lib/outputs/csv/csv.py b/prowler/lib/outputs/csv/csv.py index b6976c2a2f..aac2acdd08 100644 --- a/prowler/lib/outputs/csv/csv.py +++ b/prowler/lib/outputs/csv/csv.py @@ -1,25 +1,50 @@ from csv import DictWriter -from typing import Any + +from prowler.lib.logger import logger +from prowler.lib.outputs.finding import Finding +from prowler.lib.outputs.output import Output +from prowler.lib.outputs.utils import unroll_dict, unroll_list -# TODO: remove this once we always use the new CSV(Output) -def write_csv(file_descriptor, headers, row): - csv_writer = DictWriter( - file_descriptor, - fieldnames=headers, - delimiter=";", - ) - if isinstance(row, dict): - csv_writer.writerow(row) - else: - csv_writer.writerow(row.__dict__) +class CSV(Output): + def transform(self, findings: list[Finding]) -> None: + """Transforms the findings into the CSV format. + Args: + findings (list[Finding]): a list of Finding objects -# TODO: remove this once we always use the new CSV(Output) -def generate_csv_fields(format: Any) -> list[str]: - """Generates the CSV headers for the given class""" - csv_fields = [] - # __fields__ is always available in the Pydantic's BaseModel class - for field in format.__dict__.get("__fields__").keys(): - csv_fields.append(field) - return csv_fields + """ + try: + for finding in findings: + finding_dict = {k.upper(): v for k, v in finding.dict().items()} + finding_dict["COMPLIANCE"] = unroll_dict(finding.compliance) + finding_dict["ACCOUNT_TAGS"] = unroll_list(finding.account_tags) + finding_dict["STATUS"] = finding.status.value + finding_dict["SEVERITY"] = finding.severity.value + self._data.append(finding_dict) + except Exception as error: + logger.error( + f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" + ) + + def batch_write_data_to_file(self) -> None: + """Writes the findings to a file using the CSV format using the `Output._file_descriptor`.""" + try: + if ( + getattr(self, "_file_descriptor", None) + and not self._file_descriptor.closed + and self._data + ): + csv_writer = DictWriter( + self._file_descriptor, + fieldnames=self._data[0].keys(), + delimiter=";", + ) + csv_writer.writeheader() + for finding in self._data: + csv_writer.writerow(finding) + self._file_descriptor.close() + except Exception as error: + logger.error( + f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" + ) diff --git a/prowler/lib/outputs/csv/models.py b/prowler/lib/outputs/csv/models.py deleted file mode 100644 index aac2acdd08..0000000000 --- a/prowler/lib/outputs/csv/models.py +++ /dev/null @@ -1,50 +0,0 @@ -from csv import DictWriter - -from prowler.lib.logger import logger -from prowler.lib.outputs.finding import Finding -from prowler.lib.outputs.output import Output -from prowler.lib.outputs.utils import unroll_dict, unroll_list - - -class CSV(Output): - def transform(self, findings: list[Finding]) -> None: - """Transforms the findings into the CSV format. - - Args: - findings (list[Finding]): a list of Finding objects - - """ - try: - for finding in findings: - finding_dict = {k.upper(): v for k, v in finding.dict().items()} - finding_dict["COMPLIANCE"] = unroll_dict(finding.compliance) - finding_dict["ACCOUNT_TAGS"] = unroll_list(finding.account_tags) - finding_dict["STATUS"] = finding.status.value - finding_dict["SEVERITY"] = finding.severity.value - self._data.append(finding_dict) - except Exception as error: - logger.error( - f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" - ) - - def batch_write_data_to_file(self) -> None: - """Writes the findings to a file using the CSV format using the `Output._file_descriptor`.""" - try: - if ( - getattr(self, "_file_descriptor", None) - and not self._file_descriptor.closed - and self._data - ): - csv_writer = DictWriter( - self._file_descriptor, - fieldnames=self._data[0].keys(), - delimiter=";", - ) - csv_writer.writeheader() - for finding in self._data: - csv_writer.writerow(finding) - self._file_descriptor.close() - except Exception as error: - logger.error( - f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" - ) diff --git a/tests/lib/outputs/csv/csv_test.py b/tests/lib/outputs/csv/csv_test.py index 6817981c33..6667e73ef5 100644 --- a/tests/lib/outputs/csv/csv_test.py +++ b/tests/lib/outputs/csv/csv_test.py @@ -9,8 +9,7 @@ from freezegun import freeze_time from mock import patch from prowler.config.config import prowler_version -from prowler.lib.outputs.csv.csv import write_csv -from prowler.lib.outputs.csv.models import CSV +from prowler.lib.outputs.csv.csv import CSV from prowler.lib.outputs.finding import Finding from prowler.lib.outputs.output import Output from tests.lib.outputs.fixtures.fixtures import generate_finding_output @@ -166,36 +165,6 @@ class TestCSV: output_instance.file_descriptor ) - # TODO(PRWLR-4187): remove this once we always use the new CSV(Output) - def test_write_csv_with_dict(self): - headers = ["provider", "account", "check_id"] - row = {"provider": "aws", "account": "account_try", "check_id": "account_check"} - mock_file = StringIO() - - write_csv(mock_file, headers, row) - - mock_file.seek(0) - content = mock_file.read() - assert "aws;account_try;account_check" in content - - # TODO(PRWLR-4187): remove this once we always use the new CSV(Output) - def test_write_csv_with_object(self): - class Row: - def __init__(self, provider, account, check_id): - self.provider = provider - self.account = account - self.check_id = check_id - - headers = ["provider", "account", "check_id"] - row = Row("aws", "account_try", "account_check") - mock_file = StringIO() - - write_csv(mock_file, headers, row) - - mock_file.seek(0) - content = mock_file.read() - assert "aws;account_try;account_check" in content - def test_csv_with_file_path(self): file_name = "test" extension = ".csv" diff --git a/tests/lib/outputs/outputs_test.py b/tests/lib/outputs/outputs_test.py index 117e41646a..716f4b3013 100644 --- a/tests/lib/outputs/outputs_test.py +++ b/tests/lib/outputs/outputs_test.py @@ -3,8 +3,6 @@ from unittest import mock import pytest from colorama import Fore -from prowler.lib.outputs.csv.csv import generate_csv_fields -from prowler.lib.outputs.finding import Finding from prowler.lib.outputs.outputs import extract_findings_statistics, set_report_color from prowler.lib.outputs.utils import ( parse_html_string, @@ -34,53 +32,6 @@ class TestOutputs: assert "Invalid Report Status. Must be PASS, FAIL or MANUAL" in str(exc.value) assert exc.type == Exception - def test_generate_common_csv_fields(self): - expected = [ - "auth_method", - "timestamp", - "account_uid", - "account_name", - "account_email", - "account_organization_uid", - "account_organization_name", - "account_tags", - "finding_uid", - "provider", - "check_id", - "check_title", - "check_type", - "status", - "status_extended", - "muted", - "service_name", - "subservice_name", - "severity", - "resource_type", - "resource_uid", - "resource_name", - "resource_details", - "resource_tags", - "partition", - "region", - "description", - "risk", - "related_url", - "remediation_recommendation_text", - "remediation_recommendation_url", - "remediation_code_nativeiac", - "remediation_code_terraform", - "remediation_code_cli", - "remediation_code_other", - "compliance", - "categories", - "depends_on", - "related_to", - "notes", - "prowler_version", - ] - - assert generate_csv_fields(Finding) == expected - def test_unroll_list_no_separator(self): list = ["test", "test1", "test2"]