Storage Encryption Key Rotation Sample using Veneer + Tests (GoogleCloudPlatform#672)

ryanmats · dpebot · commit 9867964b7d25 · 2016-11-22T12:12:06.000-08:00
diff --git a/storage/cloud-client/encryption.py b/storage/cloud-client/encryption.py
@@ -97,8 +97,27 @@ def rotate_encryption_key(bucket_name, blob_name, base64_encryption_key,
                           base64_new_encryption_key):
     """Performs a key rotation by re-writing an encrypted blob with a new
     encryption key."""
-    raise NotImplementedError(
-        'This is currently not available using the Cloud Client Library.')
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+    current_encryption_key = base64.b64decode(base64_encryption_key)
+    new_encryption_key = base64.b64decode(base64_new_encryption_key)
+
+    # Both source_blob and destination_blob refer to the same storage object,
+    # but destination_blob has the new encryption key.
+    source_blob = Blob(
+        blob_name, bucket, encryption_key=current_encryption_key)
+    destination_blob = Blob(
+        blob_name, bucket, encryption_key=new_encryption_key)
+
+    token = None
+
+    while True:
+        token, bytes_rewritten, total_bytes = destination_blob.rewrite(
+            source_blob, token=token)
+        if token is None:
+            break
+
+    print('Key rotation complete for Blob {}'.format(blob_name))
 
 
 if __name__ == '__main__':
diff --git a/storage/cloud-client/encryption_test.py b/storage/cloud-client/encryption_test.py
@@ -16,13 +16,18 @@
 import tempfile
 
 from google.cloud import storage
+from google.cloud.storage import Blob
 import pytest
 
 import encryption
 
+
 TEST_ENCRYPTION_KEY = 'brtJUWneL92g5q0N2gyDSnlPSYAiIVZ/cWgjyZNeMy0='
 TEST_ENCRYPTION_KEY_DECODED = base64.b64decode(TEST_ENCRYPTION_KEY)
 
+TEST_ENCRYPTION_KEY_2 = 'o4OD7SWCaPjfeEGhAY+YCgMdY9UW+OJ8mvfWD9lNtO4='
+TEST_ENCRYPTION_KEY_2_DECODED = base64.b64decode(TEST_ENCRYPTION_KEY_2)
+
 
 def test_generate_encryption_key(capsys):
     encryption.generate_encryption_key()
@@ -47,9 +52,9 @@ def test_upload_encrypted_blob(cloud_config):
 def test_blob(cloud_config):
     """Provides a pre-existing blob in the test bucket."""
     bucket = storage.Client().bucket(cloud_config.storage_bucket)
-    blob = bucket.blob('encryption_test_sigil')
+    blob = Blob('encryption_test_sigil',
+                bucket, encryption_key=TEST_ENCRYPTION_KEY_DECODED)
     content = 'Hello, is it me you\'re looking for?'
-    blob.encryption_key = TEST_ENCRYPTION_KEY_DECODED
     blob.upload_from_string(content)
     return blob.name, content
 
@@ -65,3 +70,22 @@ def test_download_blob(test_blob, cloud_config):
 
         downloaded_content = dest_file.read().decode('utf-8')
         assert downloaded_content == test_blob_content
+
+
+def test_rotate_encryption_key(test_blob, cloud_config):
+    test_blob_name, test_blob_content = test_blob
+    encryption.rotate_encryption_key(
+        cloud_config.storage_bucket,
+        test_blob_name,
+        TEST_ENCRYPTION_KEY,
+        TEST_ENCRYPTION_KEY_2)
+
+    with tempfile.NamedTemporaryFile() as dest_file:
+        encryption.download_encrypted_blob(
+            cloud_config.storage_bucket,
+            test_blob_name,
+            dest_file.name,
+            TEST_ENCRYPTION_KEY_2)
+
+        downloaded_content = dest_file.read().decode('utf-8')
+        assert downloaded_content == test_blob_content
