regscale-cli 6.21.2.0__py3-none-any.whl → 6.28.2.1__py3-none-any.whl

regscale/integrations/commercial/wizv2/{data_fetcher.py → fetchers/policy_assessment.py}

@@ -8,8 +8,8 @@ import os
 from datetime import datetime
 from typing import Dict, List, Optional, Any, Callable
 
-from regscale.integrations.commercial.wizv2.async_client import run_async_queries
-from regscale.integrations.commercial.wizv2.constants import WizVulnerabilityType, WIZ_POLICY_QUERY
+from regscale.integrations.commercial.wizv2.core.client import run_async_queries
+from regscale.integrations.commercial.wizv2.core.constants import WizVulnerabilityType, WIZ_POLICY_QUERY
 
 logger = logging.getLogger("regscale")
 
@@ -127,14 +127,10 @@ class WizApiClient:
             "Content-Type": "application/json",
         }
 
-    def fetch_policy_assessments_async(
-        self, wiz_project_id: str, progress_callback: Optional[Callable] = None
-    ) -> List[Dict[str, Any]]:
+    def fetch_policy_assessments_async(self) -> List[Dict[str, Any]]:
         """
         Fetch policy assessments using async client.
 
-        :param wiz_project_id: Wiz project ID
-        :param progress_callback: Optional progress callback
         :return: List of policy assessment nodes
         """
         try:
@@ -204,7 +200,7 @@ class WizApiClient:
                 logger.debug(f"Filter variant {filter_variant} failed: {e}")
                 continue
 
-        raise Exception(f"All filter variants failed. Last error: {last_error}")
+        raise RuntimeError(f"All filter variants failed. Last error: {last_error}")
 
     def _create_requests_session(self):
         """
@@ -333,6 +329,9 @@ class PolicyAssessmentFetcher:
             logger.debug("Async client failed, falling back to requests")
             nodes = self._fetch_with_requests()
 
+        # Clean data (trim whitespace from externalId values)
+        nodes = self._clean_node_data(nodes)
+
         # Filter to framework
         filtered_nodes = self._filter_nodes_to_framework(nodes)
 
@@ -347,7 +346,7 @@ class PolicyAssessmentFetcher:
 
         :return: List of policy assessment nodes
         """
-        return self.api_client.fetch_policy_assessments_async(self.wiz_project_id)
+        return self.api_client.fetch_policy_assessments_async()
 
     def _fetch_with_requests(self) -> List[Dict[str, Any]]:
         """
@@ -399,3 +398,61 @@ class PolicyAssessmentFetcher:
                 filtered_nodes.append(node)
 
         return filtered_nodes
+
+    def _clean_node_data(self, nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Clean node data by trimming whitespace from externalId values.
+
+        This fixes issues where Wiz API returns control IDs with trailing/leading
+        whitespace (e.g., "AC-14 " instead of "AC-14") which prevents proper matching
+        with RegScale control implementations.
+
+        :param nodes: Raw assessment nodes from Wiz API
+        :return: Cleaned nodes with trimmed externalId values
+        """
+        cleaned_nodes = []
+
+        for node in nodes:
+            try:
+                cleaned_node = self._clean_single_node(node)
+                cleaned_nodes.append(cleaned_node)
+            except Exception:
+                # On error, include the original node unchanged (defensive)
+                cleaned_nodes.append(node)
+
+        return cleaned_nodes
+
+    def _clean_single_node(self, node: Dict[str, Any]) -> Dict[str, Any]:
+        """Clean a single node's data."""
+        # Deep copy the node to avoid modifying the original
+        cleaned_node = dict(node)
+
+        policy = cleaned_node.get("policy")
+        if self._should_clean_policy(policy):
+            cleaned_node["policy"] = self._clean_policy_subcategories(policy)
+
+        return cleaned_node
+
+    def _should_clean_policy(self, policy: Dict[str, Any]) -> bool:
+        """Check if policy should be cleaned."""
+        return policy and "securitySubCategories" in policy
+
+    def _clean_policy_subcategories(self, policy: Dict[str, Any]) -> Dict[str, Any]:
+        """Clean security subcategories in policy."""
+        subcategories = policy["securitySubCategories"]
+        cleaned_subcategories = [self._clean_subcategory(subcat) for subcat in subcategories]
+
+        cleaned_policy = dict(policy)
+        cleaned_policy["securitySubCategories"] = cleaned_subcategories
+        return cleaned_policy
+
+    def _clean_subcategory(self, subcat: Dict[str, Any]) -> Dict[str, Any]:
+        """Clean a single subcategory's externalId."""
+        cleaned_subcat = dict(subcat)
+
+        if "externalId" in cleaned_subcat:
+            original_id = cleaned_subcat["externalId"]
+            cleaned_id = original_id.strip() if isinstance(original_id, str) else original_id
+            cleaned_subcat["externalId"] = cleaned_id
+
+        return cleaned_subcat

regscale/integrations/commercial/wizv2/file_cleanup.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""File cleanup utilities for Wiz integrations."""
+
+import logging
+import os
+from typing import List, Tuple
+
+logger = logging.getLogger("regscale")
+
+
+class ReportFileCleanup:
+    """Utility class for cleaning up old report files."""
+
+    @staticmethod
+    def cleanup_old_files(directory: str, file_prefix: str, extensions: List[str] = None, keep_count: int = 5) -> None:
+        """
+        Keep the most recent files matching the pattern, delete older ones.
+
+        :param str directory: Directory containing files to clean
+        :param str file_prefix: File name prefix to match (e.g., 'compliance_report_')
+        :param List[str] extensions: List of extensions to clean (e.g., ['.csv', '.json']), defaults to None
+        :param int keep_count: Number of most recent files per extension to keep, defaults to 5
+        :return: None
+        :rtype: None
+        """
+        if extensions is None:
+            extensions = [".csv", ".json", ".jsonl"]
+
+        try:
+            if not os.path.exists(directory):
+                return
+
+            matching_entries = ReportFileCleanup._find_matching_files(directory, file_prefix, extensions)
+            files_by_extension = ReportFileCleanup._group_files_by_extension(matching_entries)
+            files_deleted = ReportFileCleanup._cleanup_files_by_extension(files_by_extension, keep_count)
+
+            if files_deleted > 0:
+                logger.info(f"Cleaned up {files_deleted} old report files from {directory}")
+
+        except Exception as e:
+            logger.warning(f"Error during file cleanup in {directory}: {e}")
+
+    @staticmethod
+    def _find_matching_files(directory: str, file_prefix: str, extensions: List[str]) -> List[Tuple[str, str, str]]:
+        """
+        Find all files matching the prefix and extensions.
+
+        :param str directory: Directory to search for files
+        :param str file_prefix: File name prefix to match
+        :param List[str] extensions: List of file extensions to match
+        :return: List of tuples containing (filename, file_path, extension)
+        :rtype: List[Tuple[str, str, str]]
+        """
+        entries = []
+        for filename in os.listdir(directory):
+            if filename.startswith(file_prefix):
+                for ext in extensions:
+                    if filename.endswith(ext):
+                        file_path = os.path.join(directory, filename)
+                        entries.append((filename, file_path, ext))
+                        break
+        return entries
+
+    @staticmethod
+    def _group_files_by_extension(entries: List[Tuple[str, str, str]]) -> dict:
+        """
+        Group files by their extensions.
+
+        :param List[Tuple[str, str, str]] entries: List of file entries (filename, file_path, extension)
+        :return: Dictionary mapping extensions to lists of (filename, file_path) tuples
+        :rtype: dict
+        """
+        by_extension = {}
+        for filename, file_path, ext in entries:
+            if ext not in by_extension:
+                by_extension[ext] = []
+            by_extension[ext].append((filename, file_path))
+        return by_extension
+
+    @staticmethod
+    def _cleanup_files_by_extension(files_by_extension: dict, keep_count: int) -> int:
+        """
+        Clean up files for each extension, keeping the most recent ones.
+
+        :param dict files_by_extension: Dictionary mapping extensions to lists of (filename, file_path) tuples
+        :param int keep_count: Number of most recent files to keep for each extension
+        :return: Total number of files deleted
+        :rtype: int
+        """
+        files_deleted = 0
+
+        for ext, files in files_by_extension.items():
+            files.sort(key=lambda x: os.path.getmtime(x[1]), reverse=True)
+
+            for filename, file_path in files[keep_count:]:
+                try:
+                    os.remove(file_path)
+                    files_deleted += 1
+                    logger.debug(f"Deleted old report file: {filename}")
+                except Exception as e:
+                    logger.warning(f"Failed to delete file {filename}: {e}")
+
+        return files_deleted