fix lint and address comments

Author: vijaykanthm

securitycenter/snippets_management_api/noxfile_config.py

@@ -23,6 +23,9 @@
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
     "ignored_versions": ["2.7", "3.7", "3.9", "3.10", "3.11"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,    
     # An envvar key for determining the project id to use. Change it
     # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
     # build specific Cloud project. You can also use your own string

securitycenter/snippets_management_api/security_health_analytics_custom_module_test.py

@@ -14,20 +14,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-import uuid
+
 import backoff
-from google.api_core.exceptions import InternalServerError, ServiceUnavailable, NotFound
+
+from google.api_core.exceptions import InternalServerError, NotFound, ServiceUnavailable 
 from google.cloud import securitycentermanagement_v1
-import security_health_analytics_custom_modules
-import pytest
-import time
+
 import random
+import time
+
+import pytest
+
+import security_health_analytics_custom_modules
 
-#Replace these variables before running the sample.
+# Replace these variables before running the sample.
+# GCLOUD_ORGANIZATION: The organization ID.
 ORGANIZATION_ID = os.environ["GCLOUD_ORGANIZATION"]
 LOCATION = "global"
 PREFIX = "python_sample_sha_custom_module"  # Prefix used for identifying test modules
 
+
 @pytest.fixture(scope="session", autouse=True)
 def setup_environment():
     """Fixture to ensure a clean environment by removing test modules before running tests."""
@@ -37,6 +43,7 @@ def setup_environment():
     print(f"Cleaning up existing custom modules for organization: {ORGANIZATION_ID}")
     cleanup_existing_custom_modules(ORGANIZATION_ID)
 
+
 def cleanup_existing_custom_modules(org_id: str):
     """
     Deletes all custom modules matching a specific naming pattern.
@@ -59,6 +66,7 @@ def cleanup_existing_custom_modules(org_id: str):
     except NotFound:
         print(f"Custom Module not found for deletion: {module.name}")
 
+
 def add_custom_module(org_id: str):
 
     parent = f"organizations/{org_id}/locations/global"
@@ -99,17 +107,17 @@ def add_custom_module(org_id: str):
     }
 
     request = securitycentermanagement_v1.CreateSecurityHealthAnalyticsCustomModuleRequest(
-        parent= parent,
-        security_health_analytics_custom_module= custom_module,
+        parent=parent,
+        security_health_analytics_custom_module=custom_module,
     )
-
     response = client.create_security_health_analytics_custom_module(request=request)
     print(f"Created Security Health Analytics Custom Module: {response.name}")
     module_name = response.name
     module_id = module_name.split("/")[-1]
 
     return module_name, module_id
 
+
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
@@ -124,6 +132,7 @@ def test_create_security_health_analytics_custom_module():
     assert response.display_name.startswith(PREFIX)
     assert response.enablement_state == securitycentermanagement_v1.SecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
 
+
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
@@ -139,9 +148,9 @@ def test_get_security_health_analytics_custom_module():
     # Verify that the custom module was created
     assert response.display_name.startswith(PREFIX)
     assert response.enablement_state == securitycentermanagement_v1.SecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
-    
     print(f"Retrieved Custom Module: {response.name}")
 
+
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
@@ -160,6 +169,7 @@ def test_delete_security_health_analytics_custom_module():
 
     print(f"Custom module was deleted successfully: {module_id}")
 
+
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
@@ -184,6 +194,7 @@ def test_list_security_health_analytics_custom_module():
         == securitycentermanagement_v1.SecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
     )
 
+
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )

securitycenter/snippets_management_api/security_health_analytics_custom_modules.py

@@ -14,15 +14,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict
-import time
 import random
+import time
+from typing import Dict
+
 from google.cloud import securitycentermanagement_v1
+from google.api_core.exceptions import NotFound
+
 
 # [START securitycenter_create_security_health_analytics_custom_module]
 def create_security_health_analytics_custom_module(parent: str) -> Dict:
     """
     Creates a Security Health Analytics custom module.
+
+    This custom module evaluates Cloud KMS CryptoKeys to ensure their rotation period exceeds 30 days (2592000 seconds),
+    as per security best practices. A shorter rotation period helps reduce the risk of exposure in the event of a compromise.
+
     Args:
         parent: Use any one of the following options:
                 - organizations/{organization_id}/locations/{location_id}
@@ -42,24 +49,33 @@ def create_security_health_analytics_custom_module(parent: str) -> Dict:
         "display_name": display_name,
         "enablement_state": "ENABLED",
         "custom_config": {
-            "description": "Sample custom module for testing purpose. Please do not delete.",
+            "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": "GCE Instance High Severity",
-                "description": "Custom module to detect high severity issues on GCE instances.",
+                "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": "Ensure proper security configurations on GCE instances.",
+            "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 name of the instance",
+                            "description": "The resource name of the CryptoKey being evaluated.",
                             "expression": "resource.name",
                             "location": "global",
-                            "title": "Instance Name",
+                            "title": "CryptoKey Resource Name",
                         },
                     }
                 ]
@@ -68,15 +84,16 @@ def create_security_health_analytics_custom_module(parent: str) -> Dict:
     }
 
     request = securitycentermanagement_v1.CreateSecurityHealthAnalyticsCustomModuleRequest(
-        parent= parent,
-        security_health_analytics_custom_module= custom_module,
+        parent=parent,
+        security_health_analytics_custom_module=custom_module,
     )
 
     response = client.create_security_health_analytics_custom_module(request=request)
     print(f"Created Security Health Analytics Custom Module: {response.name}")
     return response
 # [END securitycenter_create_security_health_analytics_custom_module]
 
+
 # [START securitycenter_get_security_health_analytics_custom_module]
 def get_security_health_analytics_custom_module(parent: str, module_id: str):
     """
@@ -95,7 +112,7 @@ def get_security_health_analytics_custom_module(parent: str, module_id: str):
 
     try:
         request = securitycentermanagement_v1.GetSecurityHealthAnalyticsCustomModuleRequest(
-            name= f"{parent}/securityHealthAnalyticsCustomModules/{module_id}",
+            name=f"{parent}/securityHealthAnalyticsCustomModules/{module_id}",
         )
 
         response = client.get_security_health_analytics_custom_module(request=request)
@@ -107,7 +124,6 @@ def get_security_health_analytics_custom_module(parent: str, module_id: str):
 # [END securitycenter_get_security_health_analytics_custom_module]
 
 
-
 # [START securitycenter_list_security_health_analytics_custom_module]
 def list_security_health_analytics_custom_module(parent: str):
     """
@@ -122,22 +138,20 @@ def list_security_health_analytics_custom_module(parent: str):
     Raises:
         NotFound: If the specified custom module does not exist.
     """
-    from google.api_core.exceptions import NotFound
 
     client = securitycentermanagement_v1.SecurityCenterManagementClient()
 
     try:
         request = securitycentermanagement_v1.ListSecurityHealthAnalyticsCustomModulesRequest(
-            parent= parent,
+            parent=parent,
         )
 
         response = client.list_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)
-        
+            custom_modules.append(custom_module)        
         return custom_modules
     except NotFound as e:
         print(f"Parent resource not found: {parent}")
@@ -147,6 +161,7 @@ def list_security_health_analytics_custom_module(parent: str):
         raise e
 # [END securitycenter_list_security_health_analytics_custom_module]
 
+
 # [START securitycenter_delete_security_health_analytics_custom_module]
 def delete_security_health_analytics_custom_module(parent: str, module_id: str):
     """
@@ -165,7 +180,7 @@ def delete_security_health_analytics_custom_module(parent: str, module_id: str):
 
     try:
         request = securitycentermanagement_v1.DeleteSecurityHealthAnalyticsCustomModuleRequest(
-            name= f"{parent}/securityHealthAnalyticsCustomModules/{module_id}",
+            name=f"{parent}/securityHealthAnalyticsCustomModules/{module_id}",
         )
 
         client.delete_security_health_analytics_custom_module(request=request)
@@ -175,6 +190,7 @@ def delete_security_health_analytics_custom_module(parent: str, module_id: str):
         raise e
 # [END securitycenter_delete_security_health_analytics_custom_module]
 
+
 # [START securitycenter_update_security_health_analytics_custom_module]
 def update_security_health_analytics_custom_module(parent: str, module_id: str):
     """
@@ -190,7 +206,6 @@ def update_security_health_analytics_custom_module(parent: str, module_id: str):
         NotFound: If the specified custom module does not exist.
     """
     from google.protobuf.field_mask_pb2 import FieldMask
-    from google.api_core.exceptions import NotFound, InternalServerError, ServiceUnavailable
 
     client = securitycentermanagement_v1.SecurityCenterManagementClient()
     try: