feat: add samples for bucket encryption enforcement config

Co-authored-by: nidhiii-27 <224584462+nidhiii-27@users.noreply.github.com>

Author: google-labs-jules[bot]

samples/snippets/encryption_test.py

@@ -27,6 +27,10 @@
 import storage_object_csek_to_cmek
 import storage_rotate_encryption_key
 import storage_upload_encrypted_file
+import storage_get_bucket_encryption_enforcement_config
+import storage_set_bucket_encryption_enforcement_config
+import storage_update_encryption_enforcement_config
+import storage_remove_all_bucket_encryption_enforcement_config
 
 BUCKET = os.environ["CLOUD_STORAGE_BUCKET"]
 KMS_KEY = os.environ["MAIN_CLOUD_KMS_KEY"]
@@ -126,3 +130,50 @@ def test_object_csek_to_cmek(test_blob):
     )
 
     assert cmek_blob.download_as_bytes(), test_blob_content
+
+def test_bucket_encryption_enforcement_config(capsys):
+    bucket_name = f"test_encryption_enforcement_{uuid.uuid4().hex}"
+
+    try:
+        # Create
+        storage_set_bucket_encryption_enforcement_config.set_bucket_encryption_enforcement_config(bucket_name)
+        out, _ = capsys.readouterr()
+        assert f"Created bucket {bucket_name} with Encryption Enforcement Config." in out
+
+        # Get
+        storage_get_bucket_encryption_enforcement_config.get_bucket_encryption_enforcement_config(bucket_name)
+        out, _ = capsys.readouterr()
+        assert f"Encryption Enforcement Config for bucket {bucket_name}:" in out
+        assert "Customer-managed encryption enforcement config restriction mode: NOT_RESTRICTED" in out
+        assert "Customer-supplied encryption enforcement config restriction mode: FULLY_RESTRICTED" in out
+        assert "Google-managed encryption enforcement config restriction mode: FULLY_RESTRICTED" in out
+
+        # Update
+        storage_update_encryption_enforcement_config.update_encryption_enforcement_config(bucket_name)
+        out, _ = capsys.readouterr()
+        assert f"Encryption enforcement policy updated for bucket {bucket_name}." in out
+
+        # Get after update
+        storage_get_bucket_encryption_enforcement_config.get_bucket_encryption_enforcement_config(bucket_name)
+        out, _ = capsys.readouterr()
+        assert "Customer-managed encryption enforcement config restriction mode: NOT_RESTRICTED" in out
+        assert "Customer-supplied encryption enforcement config restriction mode: None" in out
+        assert "Google-managed encryption enforcement config restriction mode: FULLY_RESTRICTED" in out
+
+        # Remove
+        storage_remove_all_bucket_encryption_enforcement_config.remove_all_bucket_encryption_enforcement_config(bucket_name)
+        out, _ = capsys.readouterr()
+        assert f"Removed Encryption Enforcement Config from bucket {bucket_name}." in out
+
+        # Get after remove
+        storage_get_bucket_encryption_enforcement_config.get_bucket_encryption_enforcement_config(bucket_name)
+        out, _ = capsys.readouterr()
+        assert "Customer-managed encryption enforcement config restriction mode: None" in out
+        assert "Customer-supplied encryption enforcement config restriction mode: None" in out
+        assert "Google-managed encryption enforcement config restriction mode: None" in out
+
+    finally:
+        try:
+            storage.Client().get_bucket(bucket_name).delete(force=True)
+        except Exception:
+            pass

samples/snippets/storage_get_bucket_encryption_enforcement_config.py

@@ -0,0 +1,40 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import storage
+
+# [START storage_get_bucket_encryption_enforcement_config]
+def get_bucket_encryption_enforcement_config(bucket_name):
+    """Gets the bucket encryption enforcement configuration."""
+    # The ID of your GCS bucket
+    # bucket_name = "your-unique-bucket-name"
+
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+
+    print(f"Encryption Enforcement Config for bucket {bucket.name}:")
+
+    cmek_config = bucket.customer_managed_encryption_enforcement_config
+    csek_config = bucket.customer_supplied_encryption_enforcement_config
+    gmek_config = bucket.google_managed_encryption_enforcement_config
+
+    print(f"Customer-managed encryption enforcement config restriction mode: {cmek_config.restriction_mode if cmek_config else None}")
+    print(f"Customer-supplied encryption enforcement config restriction mode: {csek_config.restriction_mode if csek_config else None}")
+    print(f"Google-managed encryption enforcement config restriction mode: {gmek_config.restriction_mode if gmek_config else None}")
+
+
+# [END storage_get_bucket_encryption_enforcement_config]
+
+if __name__ == "__main__":
+    get_bucket_encryption_enforcement_config(bucket_name="your-unique-bucket-name")

samples/snippets/storage_remove_all_bucket_encryption_enforcement_config.py

@@ -0,0 +1,36 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import storage
+
+# [START storage_remove_all_bucket_encryption_enforcement_config]
+def remove_all_bucket_encryption_enforcement_config(bucket_name):
+    """Removes all bucket encryption enforcement configuration."""
+    # The ID of your GCS bucket
+    # bucket_name = "your-unique-bucket-name"
+
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+
+    bucket.customer_managed_encryption_enforcement_config = None
+    bucket.customer_supplied_encryption_enforcement_config = None
+    bucket.google_managed_encryption_enforcement_config = None
+    bucket.patch()
+
+    print(f"Removed Encryption Enforcement Config from bucket {bucket.name}.")
+
+# [END storage_remove_all_bucket_encryption_enforcement_config]
+
+if __name__ == "__main__":
+    remove_all_bucket_encryption_enforcement_config(bucket_name="your-unique-bucket-name")

samples/snippets/storage_set_bucket_encryption_enforcement_config.py

@@ -0,0 +1,40 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import storage
+
+# [START storage_set_bucket_encryption_enforcement_config]
+def set_bucket_encryption_enforcement_config(bucket_name):
+    """Creates a bucket with encryption enforcement configuration."""
+    # The ID of your GCS bucket
+    # bucket_name = "your-unique-bucket-name"
+
+    storage_client = storage.Client()
+    bucket = storage_client.bucket(bucket_name)
+
+    # Restriction mode can be "FULLY_RESTRICTED" or "NOT_RESTRICTED"
+    from google.cloud.storage.bucket import EncryptionEnforcementConfig
+
+    bucket.customer_managed_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="NOT_RESTRICTED")
+    bucket.customer_supplied_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="FULLY_RESTRICTED")
+    bucket.google_managed_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="FULLY_RESTRICTED")
+
+    bucket.create()
+
+    print(f"Created bucket {bucket.name} with Encryption Enforcement Config.")
+
+# [END storage_set_bucket_encryption_enforcement_config]
+
+if __name__ == "__main__":
+    set_bucket_encryption_enforcement_config(bucket_name="your-unique-bucket-name")

samples/snippets/storage_update_encryption_enforcement_config.py

@@ -0,0 +1,42 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import storage
+
+# [START storage_update_encryption_enforcement_config]
+def update_encryption_enforcement_config(bucket_name):
+    """Updates the encryption enforcement policy for a bucket."""
+    # The ID of your GCS bucket
+    # bucket_name = "your-unique-bucket-name"
+
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+
+    # 1. Update a specific type (e.g., change GMEK to FULLY_RESTRICTED)
+    from google.cloud.storage.bucket import EncryptionEnforcementConfig
+    bucket.google_managed_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="FULLY_RESTRICTED")
+    bucket.customer_managed_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="NOT_RESTRICTED")
+
+    # 2. Remove a specific type (e.g., remove CSEK enforcement)
+    bucket.customer_supplied_encryption_enforcement_config = None
+
+    bucket.patch()
+
+    print(f"Encryption enforcement policy updated for bucket {bucket.name}.")
+    print("GMEK is now fully restricted, CMEK is now not restricted, and CSEK enforcement has been removed.")
+
+# [END storage_update_encryption_enforcement_config]
+
+if __name__ == "__main__":
+    update_encryption_enforcement_config(bucket_name="your-unique-bucket-name")