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

regscale/integrations/commercial/wizv2/utils/__init__.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Utils module for Wiz integration - re-exports from main.py for clean imports."""
+
+# Import all util functions from the main utils module
+from regscale.integrations.commercial.wizv2.utils.main import (
+    create_asset_type,
+    get_notes_from_wiz_props,
+    handle_management_type,
+    map_category,
+    is_report_expired,
+    convert_first_seen_to_days,
+    fetch_report_by_id,
+    download_file,
+    fetch_sbom_report,
+    compliance_job_progress,
+    get_report_url_and_status,
+    get_or_create_report_id,
+    create_compliance_report,
+    download_report,
+    rerun_expired_report,
+)
+
+# Import constants needed by tests and other modules
+from regscale.integrations.commercial.wizv2.core.constants import (
+    CHECK_INTERVAL_FOR_DOWNLOAD_REPORT,
+    MAX_RETRIES,
+)
+
+__all__ = [
+    "create_asset_type",
+    "get_notes_from_wiz_props",
+    "handle_management_type",
+    "map_category",
+    "is_report_expired",
+    "convert_first_seen_to_days",
+    "fetch_report_by_id",
+    "download_file",
+    "fetch_sbom_report",
+    "compliance_job_progress",
+    "get_report_url_and_status",
+    "get_or_create_report_id",
+    "create_compliance_report",
+    "download_report",
+    "rerun_expired_report",
+    "CHECK_INTERVAL_FOR_DOWNLOAD_REPORT",
+    "MAX_RETRIES",
+]

regscale/integrations/commercial/wizv2/{utils.py → utils/main.py}

@@ -28,7 +28,7 @@ from regscale.core.app.utils.app_utils import (
 )
 from regscale.core.utils.date import datetime_obj
 from regscale.integrations.commercial.cpe import extract_product_name_and_version
-from regscale.integrations.commercial.wizv2.constants import (
+from regscale.integrations.commercial.wizv2.core.constants import (
     BEARER,
     CHECK_INTERVAL_FOR_DOWNLOAD_REPORT,
     CONTENT_TYPE,
@@ -42,7 +42,7 @@ from regscale.integrations.commercial.wizv2.constants import (
 )
 from regscale.models.integration_models.wizv2 import ComplianceReport, ComplianceCheckStatus
 from regscale.integrations.commercial.wizv2.variables import WizVariables
-from regscale.integrations.commercial.wizv2.wiz_auth import wiz_authenticate
+from regscale.integrations.commercial.wizv2.core.auth import wiz_authenticate
 from regscale.models import (
     File,
     Sbom,
@@ -167,34 +167,57 @@ def map_category(node: dict[str, Any]) -> regscale_models.AssetCategory:
     :return: RegScale AssetCategory
     :rtype: regscale_models.AssetCategory
     """
-
     # First check if there is a CPE which can tell us the category directly.
+    if category := _get_category_from_cpe(node):
+        return category
+
+    # Then try mapping by the configured Wiz hardware asset and technology deployment model types.
+    asset_type = node.get("type", "")
+    if category := _get_category_from_hardware_types(node, asset_type):
+        return category
+
+    # Finally try matching the asset type directly by name.
+    if category := _get_category_from_asset_type(asset_type, node):
+        return category
+
+    # If all else fails, default to software.
+    return regscale_models.AssetCategory.Software
+
+
+def _get_category_from_cpe(node: dict[str, Any]) -> Optional[regscale_models.AssetCategory]:
+    """Get asset category from CPE information."""
     cpe = node.get("graphEntity", {}).get("properties", {}).get("cpe", "")
     cpe_part = extract_product_name_and_version(cpe).get("part", "")
     if cpe_part and cpe_part.lower() in CPE_PART_TO_CATEGORY_MAPPING:
         return CPE_PART_TO_CATEGORY_MAPPING[cpe_part]
+    return None
 
-    # Then try mapping by the configured Wiz hardware asset and technology deployment model types.
-    asset_type = node.get("type", "")
-    if WizVariables.useWizHardwareAssetTypes:
-        if asset_type in WizVariables.wizHardwareAssetTypes:
-            return regscale_models.AssetCategory.Hardware
-        if (graph_entity := node.get("graphEntity", {})) and (techs := graph_entity.get("technologies", [])):
-            for tech in techs:
-                # We double check just in case we get an explicit None for the technologies.
-                if tech and tech.get("deploymentModel", None) in WizVariables.wizHardwareAssetTypes:
-                    return regscale_models.AssetCategory.Hardware
-        else:
-            logger.debug("No graphEntity set for node %r, default to Software.", node)
 
-    # Finally try matching the asset type directly by name.
+def _get_category_from_hardware_types(node: dict[str, Any], asset_type: str) -> Optional[regscale_models.AssetCategory]:
+    """Get asset category from configured hardware types."""
+    if not WizVariables.useWizHardwareAssetTypes:
+        return None
+
+    if asset_type in WizVariables.wizHardwareAssetTypes:
+        return regscale_models.AssetCategory.Hardware
+
+    if (graph_entity := node.get("graphEntity", {})) and (techs := graph_entity.get("technologies", [])):
+        for tech in techs:
+            if tech and tech.get("deploymentModel", None) in WizVariables.wizHardwareAssetTypes:
+                return regscale_models.AssetCategory.Hardware
+    else:
+        logger.debug("No graphEntity set for node %r, default to Software.", node)
+
+    return None
+
+
+def _get_category_from_asset_type(asset_type: str, node: dict[str, Any]) -> Optional[regscale_models.AssetCategory]:
+    """Get asset category from asset type name."""
     if hasattr(regscale_models.AssetCategory, asset_type):
         if asset_category := getattr(regscale_models.AssetCategory, asset_type):
             return asset_category
         logger.debug("Unknown AssetType %r for node %r. Defaulting to Software.", asset_type, node)
-
-    # If all else fails, default to software.
-    return regscale_models.AssetCategory.Software
+    return None
 
 
 def convert_first_seen_to_days(first_seen: str) -> int:
@@ -753,27 +776,41 @@ def get_report_url_and_status(report_id: str) -> str:
             raise requests.RequestException("Failed to download report")
 
         response_json = response.json()
-        if errors := response_json.get("errors"):
-            message = errors[0]["message"]
-            if RATE_LIMIT_MSG in message:
-                rate = errors[0]["extensions"]["retryAfter"]
-                logger.warning("Sleeping %i seconds due to rate limit", rate)
-                time.sleep(rate)
-                continue
-
-            logger.error(errors)
-        else:
-            status = response_json.get("data", {}).get("report", {}).get("lastRun", {}).get("status")
-            if status == "COMPLETED":
-                return response_json["data"]["report"]["lastRun"]["url"]
-            elif status == "EXPIRED":
-                logger.warning("Report %s is expired, rerunning report...", report_id)
-                rerun_expired_report({"reportId": report_id})
-                return get_report_url_and_status(report_id)
+        if url := _handle_report_response(response_json, report_id):
+            return url
 
     raise requests.RequestException("Download failed, exceeding the maximum number of retries")
 
 
+def _handle_report_response(response_json: dict, report_id: str) -> Optional[str]:
+    """Handle report response and return URL if ready."""
+    if errors := response_json.get("errors"):
+        if _handle_rate_limit_error(errors):
+            return None
+        logger.error(errors)
+        return None
+
+    status = response_json.get("data", {}).get("report", {}).get("lastRun", {}).get("status")
+    if status == "COMPLETED":
+        return response_json["data"]["report"]["lastRun"]["url"]
+    if status == "EXPIRED":
+        logger.warning("Report %s is expired, rerunning report...", report_id)
+        rerun_expired_report({"reportId": report_id})
+        return get_report_url_and_status(report_id)
+    return None
+
+
+def _handle_rate_limit_error(errors: list) -> bool:
+    """Handle rate limit error and return True if rate limited."""
+    message = errors[0]["message"]
+    if RATE_LIMIT_MSG in message:
+        rate = errors[0]["extensions"]["retryAfter"]
+        logger.warning("Sleeping %i seconds due to rate limit", rate)
+        time.sleep(rate)
+        return True
+    return False
+
+
 def download_report(variables: Dict) -> requests.Response:
     """
     Return a download URL for a provided Wiz report id
@@ -1264,35 +1301,53 @@ def report_result_to_implementation_status(result: str) -> str:
     compliance_settings = get_wiz_compliance_settings()
 
     if compliance_settings:
-        try:
-            # Get implementation status labels from compliance settings
-            status_labels = compliance_settings.get_field_labels("implementationStatus")
-
-            # Map compliance check result to implementation status
-            result_lower = result.lower()
-            for label in status_labels:
-                label_lower = label.lower()
-                if result_lower == ComplianceCheckStatus.PASS.value.lower():
-                    if label_lower in ["implemented", "complete", "compliant"]:
-                        return label
-                elif result_lower == ComplianceCheckStatus.FAIL.value.lower():
-                    if label_lower in ["inremediation", "in remediation", "remediation", "failed", "non-compliant"]:
-                        return label
-                else:  # Not implemented or other status
-                    if label_lower in ["notimplemented", "not implemented", "pending", "planned"]:
-                        return label
-
-            logger.debug(f"No matching compliance setting found for result: {result}")
-        except Exception as e:
-            logger.debug(f"Error using compliance settings for implementation status mapping: {e}")
+        if status := _try_get_status_from_settings(compliance_settings, result):
+            return status
 
     # Fallback to default mapping
+    return _get_default_status_mapping(result)
+
+
+def _try_get_status_from_settings(compliance_settings, result: str) -> Optional[str]:
+    """Try to get status from compliance settings."""
+    try:
+        status_labels = compliance_settings.get_field_labels("implementationStatus")
+        result_lower = result.lower()
+
+        for label in status_labels:
+            if status := _match_label_to_result(label, result_lower):
+                return status
+
+        logger.debug(f"No matching compliance setting found for result: {result}")
+    except Exception as e:
+        logger.debug(f"Error using compliance settings for implementation status mapping: {e}")
+    return None
+
+
+def _match_label_to_result(label: str, result_lower: str) -> Optional[str]:
+    """Match a label to a result status."""
+    label_lower = label.lower()
+
+    if result_lower == ComplianceCheckStatus.PASS.value.lower():
+        if label_lower in ["implemented", "complete", "compliant"]:
+            return label
+    elif result_lower == ComplianceCheckStatus.FAIL.value.lower():
+        if label_lower in ["inremediation", "in remediation", "remediation", "failed", "non-compliant"]:
+            return label
+    else:  # Not implemented or other status
+        if label_lower in ["notimplemented", "not implemented", "pending", "planned"]:
+            return label
+
+    return None
+
+
+def _get_default_status_mapping(result: str) -> str:
+    """Get default status mapping for result."""
     if result == ComplianceCheckStatus.PASS.value:
         return ControlImplementationStatus.Implemented.value
-    elif result == ComplianceCheckStatus.FAIL.value:
+    if result == ComplianceCheckStatus.FAIL.value:
         return ControlImplementationStatus.InRemediation.value
-    else:
-        return ControlImplementationStatus.NotImplemented.value
+    return ControlImplementationStatus.NotImplemented.value
 
 
 def create_vulnerabilities_from_wiz_findings(
@@ -1385,7 +1440,7 @@ def create_single_vulnerability_from_wiz_data(
     """
     # Import here to avoid circular imports
     from regscale.integrations.commercial.wizv2.scanner import WizVulnerabilityIntegration
-    from regscale.integrations.commercial.wizv2.constants import WizVulnerabilityType
+    from regscale.integrations.commercial.wizv2.core.constants import WizVulnerabilityType
 
     try:
         # Create integration instance

regscale/integrations/commercial/wizv2/variables.py

@@ -3,7 +3,93 @@
 """Wiz Variables"""
 
 from regscale.core.app.utils.variables import RsVariableType, RsVariablesMeta
-from regscale.integrations.commercial.wizv2.constants import RECOMMENDED_WIZ_INVENTORY_TYPES, DEFAULT_WIZ_HARDWARE_TYPES
+
+# Define constants locally to avoid circular import with core.constants
+_RECOMMENDED_WIZ_INVENTORY_TYPES = [
+    # Compute resources
+    "CONTAINER",
+    "CONTAINER_GROUP",
+    "CONTAINER_IMAGE",
+    "POD",
+    "SERVERLESS",
+    "SERVERLESS_PACKAGE",
+    "VIRTUAL_DESKTOP",
+    "VIRTUAL_MACHINE",
+    "VIRTUAL_MACHINE_IMAGE",
+    # Network and exposure
+    "API_GATEWAY",
+    "CDN",
+    "CERTIFICATE",
+    "DNS_RECORD",
+    "ENDPOINT",
+    "FIREWALL",
+    "GATEWAY",
+    "LOAD_BALANCER",
+    "MANAGED_CERTIFICATE",
+    "NETWORK_ADDRESS",
+    "NETWORK_INTERFACE",
+    "PRIVATE_ENDPOINT",
+    "PRIVATE_LINK",
+    "PROXY",
+    "WEB_SERVICE",
+    # Storage and data
+    "BACKUP_SERVICE",
+    "BUCKET",
+    "DATABASE",
+    "DATA_WORKLOAD",
+    "DB_SERVER",
+    "FILE_SYSTEM_SERVICE",
+    "SECRET",
+    "SECRET_CONTAINER",
+    "STORAGE_ACCOUNT",
+    "VOLUME",
+    # Identity and access management
+    "ACCESS_ROLE",
+    "AUTHENTICATION_CONFIGURATION",
+    "IAM_BINDING",
+    "RAW_ACCESS_POLICY",
+    "SERVICE_ACCOUNT",
+    # Development and CI/CD
+    "APPLICATION",
+    "CICD_SERVICE",
+    "CONFIG_MAP",
+    "CONTAINER_REGISTRY",
+    "CONTAINER_SERVICE",
+    # Kubernetes resources
+    "CONTROLLER_REVISION",
+    "KUBERNETES_CLUSTER",
+    "KUBERNETES_INGRESS",
+    "KUBERNETES_NODE",
+    "KUBERNETES_SERVICE",
+    "NAMESPACE",
+    # Infrastructure and management
+    "CLOUD_LOG_CONFIGURATION",
+    "CLOUD_ORGANIZATION",
+    "DOMAIN",
+    "EMAIL_SERVICE",
+    "ENCRYPTION_KEY",
+    "MANAGEMENT_SERVICE",
+    "MESSAGING_SERVICE",
+    "REGISTERED_DOMAIN",
+    "RESOURCE_GROUP",
+    "SERVICE_CONFIGURATION",
+    "SUBNET",
+    "SUBSCRIPTION",
+    "VIRTUAL_NETWORK",
+]
+
+_DEFAULT_WIZ_HARDWARE_TYPES = [
+    # CloudResource types
+    "VIRTUAL_MACHINE",
+    "VIRTUAL_MACHINE_IMAGE",
+    "CONTAINER",
+    "CONTAINER_IMAGE",
+    "DB_SERVER",
+    # technology deploymentModels
+    "SERVER_APPLICATION",
+    "CLIENT_APPLICATION",
+    "VIRTUAL_APPLIANCE",
+]
 
 
 class WizVariables(metaclass=RsVariablesMeta):
@@ -23,7 +109,7 @@ class WizVariables(metaclass=RsVariablesMeta):
     wizInventoryFilterBy: RsVariableType(
         str,
         '{"projectId": ["xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"], "type": ["API_GATEWAY"]}',
-        default="""{"type": ["%s"] }""" % '","'.join(RECOMMENDED_WIZ_INVENTORY_TYPES),  # type: ignore
+        default="""{"type": ["%s"] }""" % '","'.join(_RECOMMENDED_WIZ_INVENTORY_TYPES),  # type: ignore
     )  # type: ignore
     wizAccessToken: RsVariableType(str, "", sensitive=True, required=False)  # type: ignore
     wizClientId: RsVariableType(str, "", sensitive=True)  # type: ignore
@@ -34,7 +120,7 @@ class WizVariables(metaclass=RsVariablesMeta):
         list,
         '["CONTAINER", "CONTAINER_IMAGE", "VIRTUAL_MACHINE", "VIRTUAL_MACHINE_IMAGE", "DB_SERVER", '
         '"CLIENT_APPLICATION", "SERVER_APPLICATION", "VIRTUAL_APPLIANCE"]',
-        default=DEFAULT_WIZ_HARDWARE_TYPES,
+        default=_DEFAULT_WIZ_HARDWARE_TYPES,
         required=False,
     )  # type: ignore
     wizReportAge: RsVariableType(int, "14", default=14, required=False)  # type: ignore