regscale-cli 6.16.0.0__py3-none-any.whl → 6.16.2.0__py3-none-any.whl

regscale/integrations/commercial/wizv2/utils.py

@@ -1,3 +1,7 @@
+"""
+Wiz V2 Utils
+"""
+
 import codecs
 import csv
 import datetime
@@ -32,10 +36,10 @@ from regscale.integrations.commercial.wizv2.constants import (
     MAX_RETRIES,
     CHECK_INTERVAL_FOR_DOWNLOAD_REPORT,
 )
+from regscale.integrations.commercial.wizv2.models import ComplianceReport, ComplianceCheckStatus
 from regscale.integrations.commercial.wizv2.variables import WizVariables
 from regscale.integrations.commercial.wizv2.wiz_auth import wiz_authenticate
 from regscale.models import File, Sbom, SecurityPlan, Catalog, ControlImplementation, Assessment, regscale_models
-from regscale.integrations.commercial.wizv2.models import ComplianceReport, ComplianceCheckStatus
 from regscale.utils import PaginatedGraphQLClient
 from regscale.utils.decorators import deprecated
 
@@ -51,41 +55,29 @@ def get_notes_from_wiz_props(wiz_entity_properties: Dict, external_id: str) -> s
     :return: Notes
     :rtype: str
     """
+    # Define property mappings with display names and keys
+    property_mappings = [
+        ("External ID", lambda: external_id, lambda x: x),
+        ("Cloud Platform", "cloudPlatform", str),
+        ("Provider Unique ID", "providerUniqueId", str),
+        ("cloudProviderURL", "cloudProviderURL", lambda x: f'<a href="{x}" target="_blank">{x}</a>'),
+        ("Vertex ID", "_vertexID", str),
+        ("Severity Name", "severity_name", str),
+        ("Severity Description", "severity_description", str),
+    ]
+
     notes = []
-    notes.append(f"External ID: {external_id}") if external_id else None
-    (
-        notes.append(f"Cloud Platform: {wiz_entity_properties.get('cloudPlatform')}")
-        if wiz_entity_properties.get("cloudPlatform")
-        else None
-    )
-    (
-        notes.append(f"Provider Unique ID: {wiz_entity_properties.get('providerUniqueId')}")
-        if wiz_entity_properties.get("providerUniqueId")
-        else None
-    )
-    (
-        notes.append(
-            f"""cloudProviderURL:<a href="{wiz_entity_properties.get("cloudProviderURL")}"
-                            target="_blank">{wiz_entity_properties.get("cloudProviderURL")}</a>"""
-        )
-        if wiz_entity_properties.get("cloudProviderURL")
-        else None
-    )
-    (
-        notes.append(f"Vertex ID: {wiz_entity_properties.get('_vertexID')}")
-        if wiz_entity_properties.get("_vertexID")
-        else None
-    )
-    (
-        notes.append(f"Severity Name: {wiz_entity_properties.get('severity_name')}")
-        if wiz_entity_properties.get("severity_name")
-        else None
-    )
-    (
-        notes.append(f"Severity Description: {wiz_entity_properties.get('severity_description')}")
-        if wiz_entity_properties.get("severity_description")
-        else None
-    )
+    for display_name, key_or_func, formatter in property_mappings:
+        # Handle external_id special case
+        if callable(key_or_func):
+            value = key_or_func()
+        else:
+            value = wiz_entity_properties.get(key_or_func)
+
+        if value:
+            formatted_value = formatter(value)
+            notes.append(f"{display_name}: {formatted_value}")
+
     return "<br>".join(notes)
 
 
@@ -108,7 +100,6 @@ def create_asset_type(asset_type: str) -> str:
     :return: Asset type
     :rtype: str
     """
-    #
     asset_type = asset_type.title().replace("_", " ")
     meta_data_list = regscale_models.Metadata.get_metadata_by_module_field(module="assets", field="Asset Type")
     if not any(meta_data.value == asset_type for meta_data in meta_data_list):
@@ -360,8 +351,9 @@ def create_report_if_needed(
         )
         logger.info(f"Wiz compliance report created with ID {wiz_report_id}")
         return [wiz_report_id]
-
-    return [report["id"] for report in reports if any(frame in report["name"] for frame in frames)]
+    logger.debug(f"Returning report ids for these reports {(report['name'] for report in reports)}")
+    reports = [report["id"] for report in reports if any(frame in report["name"] for frame in frames)]
+    return reports
 
 
 def fetch_and_process_report_data(wiz_report_ids: List) -> List:
@@ -384,6 +376,95 @@ def fetch_and_process_report_data(wiz_report_ids: List) -> List:
     return report_data
 
 
+def get_or_create_report_id(
+    project_id: str,
+    frameworks: List[str],
+    wiz_frameworks: List[Dict],
+    existing_reports: List[Dict],
+    target_framework: str,
+) -> str:
+    """
+    Get an existing report ID or create a new one for the target framework.
+
+    :param project_id: Project identifier
+    :param frameworks: List of framework names
+    :param wiz_frameworks: List of framework details with IDs
+    :param existing_reports: List of existing reports
+    :param target_framework: Target framework name with underscores
+    :return: Single report ID
+    """
+    report_name = f"{target_framework}_project_{project_id}"
+
+    # Check for existing report with exact name
+    for report in existing_reports:
+        if report.get("name") == report_name:
+            logger.info(f"Found existing report '{report_name}' with ID {report['id']}")
+            return report["id"]
+
+    # Create new report if no exact match found
+    try:
+        framework_index = frameworks.index(target_framework)
+        framework_id = wiz_frameworks[framework_index].get("id")
+
+        report_id = create_compliance_report(
+            wiz_project_id=project_id, report_name=report_name, framework_id=framework_id
+        )
+        logger.info(f"Created new report '{report_name}' with ID {report_id}")
+        return report_id
+    except ValueError:
+        logger.error(f"Framework '{target_framework}' not found in frameworks list")
+        raise
+
+
+def fetch_report_data(report_id: str) -> List[Dict]:
+    """
+    Fetch and process data for a single report ID.
+
+    :param report_id: Report identifier
+    :return: List of report data rows
+    """
+    try:
+        download_url = get_report_url_and_status(report_id)
+        logger.info(f"Fetching report {report_id} from: {download_url}")
+
+        with closing(requests.get(url=download_url, stream=True, timeout=10)) as response:
+            response.raise_for_status()
+            logger.info(f"Streaming and parsing report {report_id}")
+
+            reader = csv.DictReader(codecs.iterdecode(response.iter_lines(), encoding="utf-8"), delimiter=",")
+            return list(reader)
+    except requests.RequestException as e:
+        error_and_exit(f"Failed to fetch report {report_id}: {str(e)}")
+    except csv.Error as e:
+        error_and_exit(f"Failed to parse CSV for report {report_id}: {str(e)}")
+
+
+# Usage example
+def process_single_report(
+    project_id: str,
+    frameworks: List[str],
+    wiz_frameworks: List[Dict],
+    existing_reports: List[Dict],
+    target_framework: str,
+) -> List[Dict]:
+    """Process a single report and return its data.
+    :param project_id: Project identifier
+    :param frameworks: List of framework names
+    :param wiz_frameworks: List of framework details with IDs
+    :param existing_reports: List of existing reports
+    :param target_framework: Target framework name with underscores
+    :return: List of report data rows
+    """
+    report_id = get_or_create_report_id(
+        project_id=project_id,
+        frameworks=frameworks,
+        wiz_frameworks=wiz_frameworks,
+        existing_reports=existing_reports,
+        target_framework=target_framework,
+    )
+    return fetch_report_data(report_id)
+
+
 def fetch_framework_report(wiz_project_id: str, snake_framework: str) -> List[Any]:
     """
     Fetch Framework Report from Wiz.
@@ -395,10 +476,17 @@ def fetch_framework_report(wiz_project_id: str, snake_framework: str) -> List[An
     """
     wiz_frameworks = fetch_frameworks()
     frames = get_framework_names(wiz_frameworks)
-    reports = list(query_reports())
+    reports = list(query_reports(wiz_project_id))
+
+    report_data = process_single_report(
+        project_id=wiz_project_id,
+        frameworks=frames,
+        wiz_frameworks=wiz_frameworks,
+        existing_reports=reports,
+        target_framework=snake_framework,
+    )
 
-    wiz_report_ids = create_report_if_needed(wiz_project_id, frames, wiz_frameworks, reports, snake_framework)
-    return fetch_and_process_report_data(wiz_report_ids)
+    return report_data
 
 
 def fetch_frameworks() -> list:
@@ -425,23 +513,21 @@ def fetch_frameworks() -> list:
     """
     variables = {
         "policyTypes": "CLOUD",
-        "first": 500,
     }
     resp = send_request(
         query=query,
         variables=variables,
         api_endpoint_url=WizVariables.wizUrl,
     )
-
-    if resp.ok:
-        # ["data"]["securityFrameworks"]["nodes"]
+    logger.debug(f"Response: {resp}")
+    if resp and resp.ok:
         data = resp.json()
         return data.get("data", {}).get("securityFrameworks", {}).get("nodes")
     else:
         error_and_exit(f"Wiz Error: {resp.status_code if resp else None} - {resp.text if resp else 'No response'}")
 
 
-def query_reports() -> list:
+def query_reports(wiz_project_id: str) -> list:
     """
     Query Report table from Wiz
 
@@ -450,7 +536,7 @@ def query_reports() -> list:
     """
 
     # The variables sent along with the above query
-    variables = {"first": 100, "filterBy": {}}
+    variables = {"first": 100, "filterBy": {"projectId": f"{wiz_project_id}"}}
 
     res = send_request(
         query=REPORTS_QUERY,
@@ -462,6 +548,7 @@ def query_reports() -> list:
         if "errors" in res.json().keys():
             error_and_exit(f'Wiz Error: {res.json()["errors"]}')
         json_result = res.json()
+        logger.debug("JSON Result: %s", json_result)
         result = json_result.get("data", {}).get("reports", {}).get("nodes")
     except requests.JSONDecodeError:
         error_and_exit(f"Unable to fetch reports from Wiz: {res.status_code}, {res.reason}")
@@ -583,7 +670,6 @@ def _sync_compliance(
     wiz_project_id: str,
     regscale_id: int,
     regscale_module: str,
-    include_not_implemented: bool,
     client_id: str,
     client_secret: str,
     catalog_id: int,
@@ -595,7 +681,6 @@ def _sync_compliance(
     :param str wiz_project_id: Wiz Project ID
     :param int regscale_id: RegScale ID
     :param str regscale_module: RegScale module
-    :param bool include_not_implemented: Include not implemented controls
     :param str client_id: Wiz Client ID
     :param str client_secret: Wiz Client Secret
     :param int catalog_id: Catalog ID, defaults to None
@@ -623,7 +708,7 @@ def _sync_compliance(
     }
     sync_framework = framework_mapping.get(framework)
     snake_framework = sync_framework.replace(" ", "_")
-    logger.info(snake_framework)
+    logger.debug(f"{snake_framework=}")
     logger.info("Fetching Wiz compliance report for project ID %s", wiz_project_id)
     report_data = fetch_framework_report(wiz_project_id, snake_framework)
     report_models = []
@@ -631,12 +716,12 @@ def _sync_compliance(
 
     catalog = Catalog.get_with_all_details(catalog_id=catalog_id)
     controls = catalog.get("controls") if catalog else []
-    passing_controls = dict()
-    failing_controls = dict()
-    controls_to_reports = dict()
-    existing_implementations = ControlImplementation.get_existing_control_implementations(parent_id=regscale_id)
+    passing_controls = {}
+    failing_controls = {}
+    controls_to_reports = {}
+
     compliance_job_progress.update(fetch_regscale_data_job, completed=True, advance=1)
-    logger.info(f"Analyzing ComplianceReport for framework {sync_framework} from Wiz")
+    logger.info("Analyzing ComplianceReport for framework %s from Wiz" % sync_framework)
     running_compliance_job = compliance_job_progress.add_task(
         "[#f68d1f]Building compliance posture from wiz report...",
         total=len(report_data),
@@ -654,20 +739,11 @@ def _sync_compliance(
                 )
                 report_models.append(cr)
                 compliance_job_progress.update(running_compliance_job, advance=1)
-        except ValidationError as e:
-            logger.error(f"Error creating ComplianceReport: {e}")
+        except ValidationError:
+            error_message = traceback.format_exc()
+            logger.error(f"Error creating ComplianceReport: {error_message}")
     try:
         saving_regscale_data_job = compliance_job_progress.add_task("[#f68d1f]Saving RegScale data...", total=1)
-        ControlImplementation.create_control_implementations(
-            controls=controls,
-            parent_id=regscale_id,
-            parent_module=regscale_module,
-            existing_implementation_dict=existing_implementations,
-            full_controls=passing_controls,
-            partial_controls={},
-            failing_controls=failing_controls,
-            include_not_implemented=include_not_implemented,
-        )
         create_assessment_from_compliance_report(
             controls_to_reports=controls_to_reports,
             regscale_id=regscale_id,
@@ -676,9 +752,9 @@ def _sync_compliance(
         )
         compliance_job_progress.update(saving_regscale_data_job, completed=True, advance=1)
 
-    except Exception as e:
-        logger.error(f"Error creating ControlImplementations from compliance report: {e}")
-        traceback.print_exc()
+    except Exception:
+        error_message = traceback.format_exc()
+        logger.error(f"Error creating ControlImplementations from compliance report: {error_message}")
     return report_models