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

regscale/integrations/commercial/wizv2/data_fetcher.py

@@ -0,0 +1,401 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Data fetching and caching logic for Wiz Policy Compliance."""
+
+import json
+import logging
+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
+
+logger = logging.getLogger("regscale")
+
+
+class WizDataCache:
+    """Manages caching of Wiz API responses."""
+
+    def __init__(self, cache_dir: str, cache_duration_minutes: int = 0) -> None:
+        """
+        Initialize the Wiz data cache.
+
+        :param cache_dir: Directory to store cache files
+        :param cache_duration_minutes: Cache TTL in minutes (0 = disabled)
+        """
+        self.cache_dir = cache_dir
+        self.cache_duration_minutes = cache_duration_minutes
+        self.force_refresh = False
+
+    def get_cache_file_path(self, wiz_project_id: str, framework_id: str) -> str:
+        """
+        Get cache file path for given project and framework.
+
+        :param wiz_project_id: Wiz project ID
+        :param framework_id: Framework ID
+        :return: Full path to cache file
+        """
+        os.makedirs(self.cache_dir, exist_ok=True)
+        return os.path.join(self.cache_dir, f"policy_assessments_{wiz_project_id}_{framework_id}.json")
+
+    def is_cache_valid(self, cache_file: str) -> bool:
+        """
+        Check if cache file exists and is within TTL.
+
+        :param cache_file: Path to cache file to check
+        :return: True if cache is valid, False otherwise
+        """
+        if self.force_refresh or self.cache_duration_minutes <= 0:
+            return False
+
+        if not os.path.exists(cache_file):
+            return False
+
+        try:
+            max_age_seconds = self.cache_duration_minutes * 60
+            file_age = datetime.now().timestamp() - os.path.getmtime(cache_file)
+            return file_age <= max_age_seconds
+        except Exception:
+            return False
+
+    def load_from_cache(self, cache_file: str) -> Optional[List[Dict[str, Any]]]:
+        """
+        Load data from cache file.
+
+        :param cache_file: Path to cache file to load
+        :return: Cached assessment nodes if valid, None otherwise
+        """
+        try:
+            with open(cache_file, "r", encoding="utf-8") as f:
+                data = json.load(f)
+
+            nodes = data.get("nodes") or data.get("assessments") or []
+            return nodes if isinstance(nodes, list) else None
+        except Exception as e:
+            logger.debug(f"Error loading cache: {e}")
+            return None
+
+    def save_to_cache(
+        self, cache_file: str, nodes: List[Dict[str, Any]], wiz_project_id: str, framework_id: str
+    ) -> None:
+        """
+        Save data to cache file.
+
+        :param cache_file: Path to cache file
+        :param nodes: Assessment nodes to cache
+        :param wiz_project_id: Wiz project ID for metadata
+        :param framework_id: Framework ID for metadata
+        """
+        if self.cache_duration_minutes <= 0:
+            return
+
+        try:
+            payload = {
+                "timestamp": datetime.now().isoformat(),
+                "wiz_project_id": wiz_project_id,
+                "framework_id": framework_id,
+                "nodes": nodes,
+            }
+            with open(cache_file, "w", encoding="utf-8") as f:
+                json.dump(payload, f, ensure_ascii=False)
+        except Exception as e:
+            logger.debug(f"Error writing cache: {e}")
+
+
+class WizApiClient:
+    """Handles Wiz API interactions."""
+
+    def __init__(self, endpoint: str, access_token: str) -> None:
+        """
+        Initialize the Wiz API client.
+
+        :param endpoint: Wiz GraphQL API endpoint URL
+        :param access_token: Wiz API access token
+        """
+        self.endpoint = endpoint
+        self.access_token = access_token
+
+    def get_headers(self) -> Dict[str, str]:
+        """
+        Get HTTP headers for API requests.
+
+        :return: Dictionary of HTTP headers including authorization
+        """
+        return {
+            "Authorization": f"Bearer {self.access_token}",
+            "Content-Type": "application/json",
+        }
+
+    def fetch_policy_assessments_async(
+        self, wiz_project_id: str, progress_callback: Optional[Callable] = None
+    ) -> 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:
+            page_size = 100
+            query_config = {
+                "type": WizVulnerabilityType.CONFIGURATION,
+                "query": WIZ_POLICY_QUERY,
+                "topic_key": "policyAssessments",
+                "variables": {"first": page_size},
+            }
+
+            # Import here to avoid circular imports during testing
+            from regscale.integrations.commercial.wizv2.utils import compliance_job_progress
+
+            with compliance_job_progress:
+                task = compliance_job_progress.add_task(
+                    f"[#f68d1f]Fetching Wiz policy assessments (async, page size: {page_size})...",
+                    total=1,
+                )
+
+                results = run_async_queries(
+                    endpoint=self.endpoint,
+                    headers=self.get_headers(),
+                    query_configs=[query_config],
+                    progress_tracker=compliance_job_progress,
+                    max_concurrent=1,
+                )
+
+                compliance_job_progress.update(task, completed=1, advance=1)
+
+            if results and len(results) == 1 and not results[0][2]:
+                return results[0][1] or []
+
+            return []
+        except Exception as e:
+            logger.debug(f"Async fetch failed: {e}")
+            raise
+
+    def fetch_policy_assessments_requests(
+        self,
+        base_variables: Dict[str, Any],
+        filter_variants: List[Optional[Dict[str, Any]]],
+        progress_callback: Optional[Callable] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Fetch policy assessments using requests library with filter variants.
+
+        :param base_variables: Base GraphQL variables
+        :param filter_variants: List of filter variants to try
+        :param progress_callback: Optional progress callback
+        :return: List of policy assessment nodes
+        """
+
+        session = self._create_requests_session()
+        last_error = None
+
+        for filter_variant in filter_variants:
+            try:
+                variables = base_variables.copy()
+                if filter_variant is not None:
+                    variables["filterBy"] = filter_variant
+
+                nodes = self._execute_paginated_query(session, variables, progress_callback)
+                return nodes
+            except Exception as e:
+                last_error = e
+                logger.debug(f"Filter variant {filter_variant} failed: {e}")
+                continue
+
+        raise Exception(f"All filter variants failed. Last error: {last_error}")
+
+    def _create_requests_session(self):
+        """
+        Create requests session with retry logic.
+
+        :return: Configured requests session with retry adapter
+        """
+        import requests
+        from requests.adapters import HTTPAdapter
+        from urllib3.util.retry import Retry
+
+        session = requests.Session()
+        retry = Retry(
+            total=5,
+            connect=5,
+            read=5,
+            backoff_factor=0.5,
+            status_forcelist=[429, 500, 502, 503, 504],
+            allowed_methods=["POST"],
+        )
+        adapter = HTTPAdapter(max_retries=retry)
+        session.mount("https://", adapter)
+        return session
+
+    def _execute_paginated_query(
+        self, session, variables: Dict[str, Any], progress_callback: Optional[Callable] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Execute paginated GraphQL query.
+
+        :param session: Requests session object
+        :param variables: GraphQL query variables
+        :param progress_callback: Optional callback for progress updates
+        :return: List of assessment nodes from all pages
+        """
+        import requests
+
+        nodes = []
+        after_cursor = variables.get("after")
+        page_index = 0
+
+        while True:
+            payload_vars = variables.copy()
+            payload_vars["after"] = after_cursor
+            payload = {"query": WIZ_POLICY_QUERY, "variables": payload_vars}
+
+            response = session.post(self.endpoint, json=payload, headers=self.get_headers(), timeout=300)
+
+            if response.status_code >= 400:
+                raise requests.HTTPError(f"{response.status_code} {response.text[:500]}")
+
+            data = response.json()
+            if "errors" in data:
+                raise RuntimeError(str(data["errors"]))
+
+            topic = data.get("data", {}).get("policyAssessments", {})
+            page_nodes = topic.get("nodes", [])
+            page_info = topic.get("pageInfo", {})
+
+            nodes.extend(page_nodes)
+            page_index += 1
+
+            if progress_callback:
+                try:
+                    progress_callback(page_index, len(page_nodes), len(nodes))
+                except Exception:
+                    pass
+
+            has_next = page_info.get("hasNextPage", False)
+            after_cursor = page_info.get("endCursor")
+
+            if not has_next:
+                break
+
+        return nodes
+
+
+class PolicyAssessmentFetcher:
+    """Main class for fetching Wiz policy assessments."""
+
+    def __init__(
+        self,
+        wiz_endpoint: str,
+        access_token: str,
+        wiz_project_id: str,
+        framework_id: str,
+        cache_duration_minutes: int = 0,
+    ) -> None:
+        """
+        Initialize the policy assessment fetcher.
+
+        :param wiz_endpoint: Wiz GraphQL API endpoint URL
+        :param access_token: Wiz API access token
+        :param wiz_project_id: Wiz project ID to query
+        :param framework_id: Framework ID to filter by
+        :param cache_duration_minutes: Cache TTL in minutes (0 = disabled)
+        """
+        self.api_client = WizApiClient(wiz_endpoint, access_token)
+        self.wiz_project_id = wiz_project_id
+        self.framework_id = framework_id
+        self.cache = WizDataCache("artifacts/wiz", cache_duration_minutes)
+
+    def fetch_policy_assessments(self) -> List[Dict[str, Any]]:
+        """
+        Fetch policy assessments from Wiz API with caching.
+
+        :return: List of filtered policy assessment nodes
+        """
+        logger.info("Fetching policy assessments from Wiz...")
+
+        # Try cache first
+        cache_file = self.cache.get_cache_file_path(self.wiz_project_id, self.framework_id)
+
+        if self.cache.is_cache_valid(cache_file):
+            cached_nodes = self.cache.load_from_cache(cache_file)
+            if cached_nodes is not None:
+                logger.info("Using cached Wiz policy assessments")
+                return cached_nodes
+
+        # Fetch from API
+        try:
+            # Try async client first
+            nodes = self._fetch_with_async_client()
+        except Exception:
+            # Fall back to requests
+            logger.debug("Async client failed, falling back to requests")
+            nodes = self._fetch_with_requests()
+
+        # Filter to framework
+        filtered_nodes = self._filter_nodes_to_framework(nodes)
+
+        # Save to cache
+        self.cache.save_to_cache(cache_file, filtered_nodes, self.wiz_project_id, self.framework_id)
+
+        return filtered_nodes
+
+    def _fetch_with_async_client(self) -> List[Dict[str, Any]]:
+        """
+        Fetch using async client.
+
+        :return: List of policy assessment nodes
+        """
+        return self.api_client.fetch_policy_assessments_async(self.wiz_project_id)
+
+    def _fetch_with_requests(self) -> List[Dict[str, Any]]:
+        """
+        Fetch using requests with filter variants.
+
+        :return: List of policy assessment nodes
+        """
+        page_size = 100
+        base_variables = {"first": page_size}
+
+        # Try multiple filter variants
+        filter_variants = [
+            {"project": [self.wiz_project_id]},
+            {"projectId": [self.wiz_project_id]},
+            {"projects": [self.wiz_project_id]},
+            {},  # Empty filterBy
+            None,  # Omit filterBy entirely
+        ]
+
+        return self.api_client.fetch_policy_assessments_requests(base_variables, filter_variants)
+
+    def _filter_nodes_to_framework(self, nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Filter nodes to only include items from the target framework.
+
+        :param nodes: Raw assessment nodes from Wiz API
+        :return: Filtered nodes belonging to the target framework
+        """
+        filtered_nodes = []
+
+        for node in nodes:
+            try:
+                subcats = ((node or {}).get("policy") or {}).get("securitySubCategories", [])
+
+                # Include if no subcategories (can't evaluate framework)
+                if not subcats:
+                    filtered_nodes.append(node)
+                    continue
+
+                # Include if any subcategory matches our framework
+                for subcat in subcats:
+                    framework_id = subcat.get("category", {}).get("framework", {}).get("id")
+                    if framework_id == self.framework_id:
+                        filtered_nodes.append(node)
+                        break
+
+            except Exception:
+                # Include on error (defensive)
+                filtered_nodes.append(node)
+
+        return filtered_nodes

regscale/integrations/commercial/wizv2/finding_processor.py

@@ -0,0 +1,295 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Finding processing and consolidation logic for Wiz Policy Compliance."""
+
+import logging
+from typing import Dict, List, Optional, Iterator, Any, Set, Union
+from collections import defaultdict
+
+from regscale.integrations.scanner_integration import IntegrationFinding
+from regscale.integrations.commercial.wizv2.policy_compliance_helpers import AssetConsolidator
+from regscale.models import regscale_models
+
+logger = logging.getLogger("regscale")
+
+
+class WizComplianceItem:
+    """Interface representing a compliance item from the main class."""
+
+    def __init__(self, compliance_item: Any) -> None:
+        """
+        Initialize wrapper for compliance item.
+
+        :param compliance_item: The compliance item object to wrap
+        """
+        self._item = compliance_item
+
+    @property
+    def resource_id(self) -> str:
+        return getattr(self._item, "resource_id", "")
+
+    @property
+    def control_id(self) -> str:
+        return getattr(self._item, "control_id", "")
+
+    @property
+    def is_fail(self) -> bool:
+        return getattr(self._item, "is_fail", False)
+
+    def get_all_control_ids(self) -> List[str]:
+        """
+        Get all control IDs this item maps to.
+
+        :return: List of control IDs this policy assessment affects
+        """
+        # This would be implemented by the main class
+        if hasattr(self._item, "_get_all_control_ids_for_compliance_item"):
+            return self._item._get_all_control_ids_for_compliance_item(self._item)
+        return [self.control_id] if self.control_id else []
+
+
+class FindingConsolidator:
+    """Consolidates multiple compliance items into control-centric findings."""
+
+    def __init__(self, integration_instance: Any) -> None:
+        """
+        Initialize the finding consolidator.
+
+        :param integration_instance: The Wiz integration instance
+        """
+        self.integration = integration_instance
+        self.asset_consolidator = AssetConsolidator()
+
+    def create_consolidated_findings(self, failed_compliance_items: List[Any]) -> Iterator[IntegrationFinding]:
+        """
+        Create consolidated findings grouped by control ID.
+
+        :param failed_compliance_items: List of failed compliance items
+        :yield: Consolidated findings
+        """
+        if not failed_compliance_items:
+            logger.debug("No failed compliance items to process")
+            return
+
+        logger.debug("Starting control-centric finding consolidation")
+
+        # Group compliance items by control ID
+        control_groups = self._group_by_control(failed_compliance_items)
+
+        if not control_groups:
+            logger.debug("No control groupings created")
+            return
+
+        logger.debug(f"Grouped into {len(control_groups)} controls with failing resources")
+
+        # Create consolidated findings for each control
+        findings_created = 0
+        for control_id, resources in control_groups.items():
+            finding = self._create_consolidated_finding_for_control(control_id, resources)
+            if finding:
+                findings_created += 1
+                yield finding
+
+        logger.debug(f"Generated {findings_created} consolidated findings")
+
+    def _group_by_control(self, compliance_items: List[Any]) -> Dict[str, Dict[str, Any]]:
+        """
+        Group compliance items by control ID.
+
+        :param compliance_items: List of compliance items to group
+        :return: Dict mapping control IDs to resource dictionaries
+        """
+        control_groups = defaultdict(dict)  # {control_id: {resource_id: compliance_item}}
+
+        for item in compliance_items:
+            wrapped_item = WizComplianceItem(item)
+            resource_id = wrapped_item.resource_id
+
+            if not resource_id:
+                continue
+
+            # Get all control IDs this item maps to
+            all_control_ids = wrapped_item.get_all_control_ids()
+            if not all_control_ids:
+                continue
+
+            # Add this resource to each control it fails
+            for control_id in all_control_ids:
+                normalized_control = control_id.upper()
+                resource_key = resource_id.lower()
+
+                # Use first occurrence for each resource-control pair
+                if resource_key not in control_groups[normalized_control]:
+                    control_groups[normalized_control][resource_key] = item
+
+        return dict(control_groups)
+
+    def _create_consolidated_finding_for_control(
+        self, control_id: str, resources: Dict[str, Any]
+    ) -> Optional[IntegrationFinding]:
+        """
+        Create a consolidated finding for a control with all affected resources.
+
+        :param control_id: Control identifier
+        :param resources: Dict of resource_id -> compliance_item
+        :return: Consolidated finding or None
+        """
+        logger.debug(f"Creating consolidated finding for control {control_id} with {len(resources)} resources")
+
+        # Filter to only resources that exist as assets in RegScale
+        asset_mappings = self._build_asset_mappings(list(resources.keys()))
+
+        if not asset_mappings:
+            logger.debug(f"No existing assets found for control {control_id}")
+            return None
+
+        logger.debug(f"Creating finding for control {control_id} with {len(asset_mappings)} existing assets")
+
+        # Use the first compliance item as the base for the finding
+        base_item = next(iter(resources.values()))
+
+        # Create the base finding
+        finding = self._create_base_finding(base_item, control_id)
+        if not finding:
+            return None
+
+        # Update with consolidated asset information
+        self._update_finding_with_assets(finding, asset_mappings)
+
+        return finding
+
+    def _build_asset_mappings(self, resource_ids: List[str]) -> Dict[str, Dict[str, str]]:
+        """
+        Build asset mappings for resources that exist in RegScale.
+
+        :param resource_ids: List of Wiz resource IDs to check
+        :return: Dict mapping resource_ids to asset info (name, wiz_id)
+        """
+        asset_mappings = {}
+
+        for resource_id in resource_ids:
+            if self.integration._asset_exists_in_regscale(resource_id):
+                asset = self.integration.get_asset_by_identifier(resource_id)
+                if asset and hasattr(asset, "name") and asset.name:
+                    asset_mappings[resource_id] = {"name": asset.name, "wiz_id": resource_id}
+                else:
+                    # Fallback to resource ID if asset name not found
+                    asset_mappings[resource_id] = {"name": resource_id, "wiz_id": resource_id}
+
+        return asset_mappings
+
+    def _create_base_finding(self, compliance_item: Any, control_id: str) -> Optional[IntegrationFinding]:
+        """Create a base finding from a compliance item for a specific control."""
+        try:
+            # Use the integration's existing method but for specific control
+            if hasattr(self.integration, "_create_finding_for_specific_control"):
+                return self.integration._create_finding_for_specific_control(compliance_item, control_id)
+            else:
+                # Fallback to generic method
+                return self.integration.create_finding_from_compliance_item(compliance_item)
+        except Exception as e:
+            logger.error(f"Error creating base finding for control {control_id}: {e}")
+            return None
+
+    def _update_finding_with_assets(
+        self, finding: IntegrationFinding, asset_mappings: Dict[str, Dict[str, str]]
+    ) -> None:
+        """
+        Update finding with consolidated asset information.
+
+        :param finding: Finding to update
+        :param asset_mappings: Asset mapping information
+        """
+        # Update asset identifier with all assets
+        consolidated_identifier = self.asset_consolidator.create_consolidated_asset_identifier(asset_mappings)
+        finding.asset_identifier = consolidated_identifier
+
+        # Update description for multiple assets
+        asset_names = [info["name"] for info in asset_mappings.values()]
+        self.asset_consolidator.update_finding_description_for_multiple_assets(finding, len(asset_names), asset_names)
+
+
+class FindingToIssueProcessor:
+    """Processes findings into RegScale issues."""
+
+    def __init__(self, integration_instance: Any) -> None:
+        """
+        Initialize the finding to issue processor.
+
+        :param integration_instance: The Wiz integration instance
+        """
+        self.integration = integration_instance
+
+    def process_findings_to_issues(self, findings: List[IntegrationFinding]) -> tuple[int, int]:
+        """
+        Process findings into issues and return counts.
+
+        :param findings: List of consolidated findings to process
+        :return: Tuple of (issues_created, issues_skipped)
+        """
+        issues_created = 0
+        issues_skipped = 0
+
+        for finding in findings:
+            try:
+                if self._process_single_finding(finding):
+                    issues_created += 1
+                else:
+                    issues_skipped += 1
+            except Exception as e:
+                logger.error(f"Error processing finding: {e}")
+                issues_skipped += 1
+
+        return issues_created, issues_skipped
+
+    def _process_single_finding(self, finding: IntegrationFinding) -> bool:
+        """
+        Process a single finding into an issue.
+
+        :param finding: Finding to process
+        :return: True if successful, False if skipped
+        """
+        # Verify assets exist
+        if not self._verify_assets_exist(finding):
+            logger.debug(f"Asset not found for finding {finding.external_id}")
+            return False
+
+        # Create or update the issue
+        try:
+            issue_title = self.integration.get_issue_title(finding)
+            issue = self.integration.create_or_update_issue_from_finding(title=issue_title, finding=finding)
+            return issue is not None
+        except Exception as e:
+            logger.error(f"Error creating/updating issue: {e}")
+            return False
+
+    def _verify_assets_exist(self, finding: IntegrationFinding) -> bool:
+        """
+        Verify that assets referenced in the finding exist.
+
+        :param finding: Finding with asset identifiers to verify
+        :return: True if all assets exist in RegScale, False otherwise
+        """
+        if not hasattr(finding, "asset_identifier") or not finding.asset_identifier:
+            return False
+
+        # For consolidated findings, asset_identifier may contain multiple assets
+        identifiers = finding.asset_identifier.split("\n")
+
+        for identifier in identifiers:
+            identifier = identifier.strip()
+            if not identifier:
+                continue
+
+            # Extract resource ID from format "Asset Name (resource-id)"
+            if "(" in identifier and identifier.endswith(")"):
+                resource_id = identifier.split("(")[-1].rstrip(")")
+            else:
+                resource_id = identifier
+
+            # Check if asset exists
+            if not self.integration._asset_exists_in_regscale(resource_id):
+                logger.debug(f"Asset {resource_id} does not exist in RegScale")
+                return False
+
+        return True