feat(securitycenter): Add Resource SCC Org Mgt API SHA Cust Modules (GetEff, ListEff, ListDesc, Simulate)

vijaykanthm · vijaykanthm · commit 5b8d9a68b3ef · 2024-12-30T03:49:34.000Z
diff --git a/securitycenter/snippets_management_api/security_health_analytics_custom_modules.py b/securitycenter/snippets_management_api/security_health_analytics_custom_modules.py
@@ -14,14 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import random
-import time
-from typing import Dict
-
-# [START securitycenter_get_effective_security_health_analytics_custom_module]
 from google.api_core.exceptions import NotFound
 from google.cloud import securitycentermanagement_v1
 
+# [START securitycenter_get_effective_security_health_analytics_custom_module]
+
+
 def get_effective_security_health_analytics_custom_module(parent: str, module_id: str):
     """
     Retrieves a Security Health Analytics custom module.
@@ -49,3 +47,154 @@ def get_effective_security_health_analytics_custom_module(parent: str, module_id
         print(f"Custom Module not found: {response.name}")
         raise e
 # [END securitycenter_get_effective_security_health_analytics_custom_module]
+
+# [START securitycenter_list_descendant_security_health_analytics_custom_module]
+
+
+def list_descendant_security_health_analytics_custom_module(parent: str):
+    """
+    Retrieves list of all resident Security Health Analytics custom modules and all of its descendants.
+    Args:
+        parent: Use any one of the following options:
+                - organizations/{organization_id}/locations/{location_id}
+                - folders/{folder_id}/locations/{location_id}
+                - projects/{project_id}/locations/{location_id}
+    Returns:
+        List of retrieved all resident Security Health Analytics custom modules and all of its descendants.
+    Raises:
+        NotFound: If the specified custom module does not exist.
+    """
+
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+
+    try:
+        request = securitycentermanagement_v1.ListDescendantSecurityHealthAnalyticsCustomModulesRequest(
+            parent=parent,
+        )
+
+        response = client.list_descendant_security_health_analytics_custom_modules(request=request)
+
+        custom_modules = []
+        for custom_module in response:
+            print(f"Custom Module: {custom_module.name}")
+            custom_modules.append(custom_module)
+        return custom_modules
+    except NotFound as e:
+        print(f"Parent resource not found: {parent}")
+        raise e
+    except Exception as e:
+        print(f"An error occurred while listing custom modules: {e}")
+        raise e
+# [END securitycenter_list_descendant_security_health_analytics_custom_module]
+
+# [START securitycenter_list_effective_security_health_analytics_custom_module]
+
+
+def list_effective_security_health_analytics_custom_module(parent: str):
+    """
+    Retrieves list of all Security Health Analytics custom modules.
+    This includes resident modules defined at the scope of the parent,
+    and inherited modules, inherited from ancestor organizations, folders, and projects (no descendants).
+
+    Args:
+        parent: Use any one of the following options:
+                - organizations/{organization_id}/locations/{location_id}
+                - folders/{folder_id}/locations/{location_id}
+                - projects/{project_id}/locations/{location_id}
+    Returns:
+        List of retrieved all Security Health Analytics custom modules.
+    Raises:
+        NotFound: If the specified custom module does not exist.
+    """
+
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+
+    try:
+        request = securitycentermanagement_v1.ListEffectiveSecurityHealthAnalyticsCustomModulesRequest(
+            parent=parent,
+        )
+
+        response = client.list_effective_security_health_analytics_custom_modules(request=request)
+
+        custom_modules = []
+        for custom_module in response:
+            print(f"Custom Module: {custom_module.name}")
+            custom_modules.append(custom_module)
+        return custom_modules
+    except NotFound as e:
+        print(f"Parent resource not found: {parent}")
+        raise e
+    except Exception as e:
+        print(f"An error occurred while listing custom modules: {e}")
+        raise e
+# [END securitycenter_list_effective_security_health_analytics_custom_module]
+
+# [START securitycenter_simulate_security_health_analytics_custom_module]
+
+
+def simulate_security_health_analytics_custom_module(parent: str):
+    """
+    Simulates the result of using a SecurityHealthAnalyticsCustomModule to check a resource.
+
+    Args:
+        parent: Use any one of the following options:
+                - organizations/{organization_id}/locations/{location_id}
+                - folders/{folder_id}/locations/{location_id}
+                - projects/{project_id}/locations/{location_id}
+    Returns:
+        Simulated Security Health Analytics custom module.
+    """
+
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+
+    # Define the custom config configuration
+    custom_config = {
+        "description": (
+            "Sample custom module for testing purposes. This custom module evaluates "
+            "Cloud KMS CryptoKeys to ensure their rotation period exceeds 30 days (2592000 seconds)."
+        ),
+        "predicate": {
+            "expression": "has(resource.rotationPeriod) && (resource.rotationPeriod > duration('2592000s'))",
+            "title": "Cloud KMS CryptoKey Rotation Period",
+            "description": (
+                "Evaluates whether the rotation period of a Cloud KMS CryptoKey exceeds 30 days. "
+                "A longer rotation period might increase the risk of exposure."
+            ),
+        },
+        "recommendation": (
+            "Review and adjust the rotation period for Cloud KMS CryptoKeys to align with your security policies. "
+            "Consider setting a shorter rotation period if possible."
+        ),
+        "resource_selector": {"resource_types": ["cloudkms.googleapis.com/CryptoKey"]},
+        "severity": "CRITICAL",
+        "custom_output": {
+            "properties": [
+                {
+                    "name": "example_property",
+                    "value_expression": {
+                        "description": "The resource name of the CryptoKey being evaluated.",
+                        "expression": "resource.name",
+                        "location": "global",
+                        "title": "CryptoKey Resource Name",
+                    },
+                }
+            ]
+        },
+    }
+
+    # Initialize request argument(s)
+    resource = securitycentermanagement_v1.types.SimulateSecurityHealthAnalyticsCustomModuleRequest.SimulatedResource()
+    resource.resource_type = "cloudkms.googleapis.com/CryptoKey"  # Replace with the correct resource type
+
+    request = securitycentermanagement_v1.SimulateSecurityHealthAnalyticsCustomModuleRequest(
+        parent=parent,
+        custom_config=custom_config,
+        resource=resource,
+    )
+
+    response = client.simulate_security_health_analytics_custom_module(request=request)
+
+    print(f"Simulated Security Health Analytics Custom Module: {response}")
+    return response
+
+# [END securitycenter_simulate_security_health_analytics_custom_module]
diff --git a/securitycenter/snippets_management_api/security_health_analytics_custom_modules_test.py b/securitycenter/snippets_management_api/security_health_analytics_custom_modules_test.py
@@ -121,6 +121,7 @@ def add_custom_module(org_id: str):
     module_id = module_name.split("/")[-1]
     return module_name, module_id
 
+
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
@@ -137,3 +138,69 @@ def test_get_effective_security_health_analytics_custom_module():
     assert response.display_name.startswith(PREFIX)
     assert response.enablement_state == securitycentermanagement_v1.EffectiveSecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
     print(f"Retrieved Custom Module: {response.name}")
+
+
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
+def test_list_descendant_security_health_analytics_custom_module():
+
+    module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
+    # Retrieve the list descendant custom modules
+    custom_modules = security_health_analytics_custom_modules.list_descendant_security_health_analytics_custom_module(parent)
+
+    assert custom_modules is not None, "Failed to retrieve the custom modules."
+    assert len(custom_modules) > 0, "No custom modules were retrieved."
+
+    # Verify the created module is in the list
+    created_module = next(
+        (module for module in custom_modules if module.name == module_name), None
+    )
+    assert created_module is not None, "Created custom module not found in the list."
+    assert created_module.display_name.startswith(PREFIX)
+    assert (
+        created_module.enablement_state
+        == securitycentermanagement_v1.SecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
+    )
+
+
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
+def test_list_effective_security_health_analytics_custom_module():
+
+    module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
+    # Retrieve the list of custom modules
+    custom_modules = security_health_analytics_custom_modules.list_effective_security_health_analytics_custom_module(parent)
+
+    assert custom_modules is not None, "Failed to retrieve the custom modules."
+    assert len(custom_modules) > 0, "No custom modules were retrieved."
+
+    # Verify the created module is in the list
+    created_module = next(
+        (module for module in custom_modules if (module.name.split("/")[-1]) == module_id), None
+    )
+    assert created_module is not None, "Created custom module not found in the list."
+    assert created_module.display_name.startswith(PREFIX)
+    assert (
+        created_module.enablement_state
+        == securitycentermanagement_v1.EffectiveSecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
+    )
+
+
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
+def test_simulate_security_health_analytics_custom_module():
+
+    module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
+
+    simulated_custom_module = security_health_analytics_custom_modules.simulate_security_health_analytics_custom_module(parent)
+
+    assert simulated_custom_module is not None, "Failed to retrieve the simulated custom module."
+    assert simulated_custom_module.result.no_violation is not None, (
+        f"Expected no_violation to be present, got {simulated_custom_module.result}."
+    )
