googleapis · tseaver · Apr 27, 2018 · Apr 27, 2018 · Apr 27, 2018
diff --git a/storage/google/cloud/storage/bucket.py b/storage/google/cloud/storage/bucket.py
@@ -171,7 +171,8 @@ def user_project(self):
         """
         return self._user_project
 
-    def blob(self, blob_name, chunk_size=None, encryption_key=None):
+    def blob(self, blob_name, chunk_size=None,
+             encryption_key=None, kms_key_name=None):
         """Factory constructor for blob object.
 
         .. note::
@@ -190,11 +191,15 @@ def blob(self, blob_name, chunk_size=None, encryption_key=None):
         :param encryption_key:
             Optional 32 byte encryption key for customer-supplied encryption.
 
+        :type kms_key_name: str
+        :param kms_key_name:
+            Optional resource name of KMS key used to encrypt blob's content.
+
         :rtype: :class:`google.cloud.storage.blob.Blob`
         :returns: The blob object created.
         """
         return Blob(name=blob_name, bucket=self, chunk_size=chunk_size,
-                    encryption_key=encryption_key)
+                    encryption_key=encryption_key, kms_key_name=kms_key_name)
 
     def notification(self, topic_name,
                      topic_project=None,

diff --git a/storage/tests/system.py b/storage/tests/system.py
@@ -990,3 +990,131 @@ def test_access_to_public_bucket(self):
         blob, = bucket.list_blobs(max_results=1)
         with tempfile.TemporaryFile() as stream:
             blob.download_to_file(stream)
+
+
+class TestKMSIntegration(TestStorageFiles):
+
+    FILENAMES = (
+        'file01.txt',
+    )
+
+    KEYRING_NAME = 'gcs-test'
+    KEY_NAME = 'gcs-test'
+    ALT_KEY_NAME = 'gcs-test-alternate'
+
+    def _kms_key_name(self, key_name=None):
+        if key_name is None:
+            key_name = self.KEY_NAME
+
+        return (
+            "projects/{}/"
+            "locations/{}/"
+            "keyRings/{}/"
+            "cryptoKeys/{}"
+        ).format(
+            Config.CLIENT.project,
+            'global',  # will be 'self.bucket.location' after launch
+            self.KEYRING_NAME,
+            key_name,
+        )
+
+    def test_blob_w_explicit_kms_key_name(self):
+        BLOB_NAME = 'explicit-kms-key-name'
+        file_data = self.FILES['simple']
+        kms_key_name = self._kms_key_name()
+        blob = self.bucket.blob(BLOB_NAME, kms_key_name=kms_key_name)
+        blob.upload_from_filename(file_data['path'])
+        self.case_blobs_to_delete.append(blob)
+        with open(file_data['path'], 'rb') as _file_data:
+            self.assertEqual(blob.download_as_string(), _file_data.read())
+        # We don't know the current version of the key.
+        self.assertTrue(blob.kms_key_name.startswith(kms_key_name))
+
+    def test_bucket_w_default_kms_key_name(self):
+        BLOB_NAME = 'default-kms-key-name'
+        OVERRIDE_BLOB_NAME = 'override-default-kms-key-name'
+        ALT_BLOB_NAME = 'alt-default-kms-key-name'
+        CLEARTEXT_BLOB_NAME = 'cleartext'
+
+        file_data = self.FILES['simple']
+
+        with open(file_data['path'], 'rb') as _file_data:
+            contents = _file_data.read()
+
+        kms_key_name = self._kms_key_name()
+        self.bucket.default_kms_key_name = kms_key_name
+        self.bucket.patch()
+        self.assertEqual(self.bucket.default_kms_key_name, kms_key_name)
+
+        defaulted_blob = self.bucket.blob(BLOB_NAME)
+        defaulted_blob.upload_from_filename(file_data['path'])
+        self.case_blobs_to_delete.append(defaulted_blob)
+
+        self.assertEqual(defaulted_blob.download_as_string(), contents)
+        # We don't know the current version of the key.
+        self.assertTrue(defaulted_blob.kms_key_name.startswith(kms_key_name))
+
+        alt_kms_key_name = self._kms_key_name(self.ALT_KEY_NAME)
+
+        override_blob = self.bucket.blob(
+            OVERRIDE_BLOB_NAME, kms_key_name=alt_kms_key_name)
+        override_blob.upload_from_filename(file_data['path'])
+        self.case_blobs_to_delete.append(override_blob)
+
+        self.assertEqual(override_blob.download_as_string(), contents)
+        # We don't know the current version of the key.
+        self.assertTrue(
+            override_blob.kms_key_name.startswith(alt_kms_key_name))
+
+        self.bucket.default_kms_key_name = alt_kms_key_name
+        self.bucket.patch()
+
+        alt_blob = self.bucket.blob(ALT_BLOB_NAME)
+        alt_blob.upload_from_filename(file_data['path'])
+        self.case_blobs_to_delete.append(alt_blob)
+
+        self.assertEqual(alt_blob.download_as_string(), contents)
+        # We don't know the current version of the key.
+        self.assertTrue(alt_blob.kms_key_name.startswith(alt_kms_key_name))
+
+        self.bucket.default_kms_key_name = None
+        self.bucket.patch()
+
+        cleartext_blob = self.bucket.blob(CLEARTEXT_BLOB_NAME)
+        cleartext_blob.upload_from_filename(file_data['path'])
+        self.case_blobs_to_delete.append(cleartext_blob)
+
+        self.assertEqual(cleartext_blob.download_as_string(), contents)
+        self.assertIsNone(cleartext_blob.kms_key_name)
+
+    def test_rewrite_rotate_csek_to_cmek(self):
+        BLOB_NAME = 'rotating-keys'
+        file_data = self.FILES['simple']
+
+        SOURCE_KEY = os.urandom(32)
+        source = self.bucket.blob(BLOB_NAME, encryption_key=SOURCE_KEY)
+        source.upload_from_filename(file_data['path'])
+        self.case_blobs_to_delete.append(source)
+        source_data = source.download_as_string()
+
+        kms_key_name = self._kms_key_name()
+
+        # We can't verify it, but ideally we would check that the following
+        # URL was resolvable with our credentals
+        # KEY_URL = 'https://cloudkms.googleapis.com/v1/{}'.format(
+        #     kms_key_name)
+
+        dest = self.bucket.blob(BLOB_NAME, kms_key_name=kms_key_name)
+        token, rewritten, total = dest.rewrite(source)
+
+        while token is not None:
+            token, rewritten, total = dest.rewrite(source, token=token)
+
+        # Not adding 'dest' to 'self.case_blobs_to_delete':  it is the
+        # same object as 'source'.
+
+        self.assertIsNone(token)
+        self.assertEqual(rewritten, len(source_data))
+        self.assertEqual(total, len(source_data))
+
+        self.assertEqual(dest.download_as_string(), source_data)
diff --git a/storage/tests/unit/test_bucket.py b/storage/tests/unit/test_bucket.py
@@ -77,7 +77,25 @@ def test_ctor_w_user_project(self):
         self.assertFalse(bucket._default_object_acl.loaded)
         self.assertIs(bucket._default_object_acl.bucket, bucket)
 
-    def test_blob(self):
+    def test_blob_wo_keys(self):
+        from google.cloud.storage.blob import Blob
+
+        BUCKET_NAME = 'BUCKET_NAME'
+        BLOB_NAME = 'BLOB_NAME'
+        CHUNK_SIZE = 1024 * 1024
+
+        bucket = self._make_one(name=BUCKET_NAME)
+        blob = bucket.blob(
+            BLOB_NAME, chunk_size=CHUNK_SIZE)
+        self.assertIsInstance(blob, Blob)
+        self.assertIs(blob.bucket, bucket)
+        self.assertIs(blob.client, bucket.client)
+        self.assertEqual(blob.name, BLOB_NAME)
+        self.assertEqual(blob.chunk_size, CHUNK_SIZE)
+        self.assertIsNone(blob._encryption_key)
+        self.assertIsNone(blob.kms_key_name)
+
+    def test_blob_w_encryption_key(self):
         from google.cloud.storage.blob import Blob
 
         BUCKET_NAME = 'BUCKET_NAME'
@@ -94,6 +112,31 @@ def test_blob(self):
         self.assertEqual(blob.name, BLOB_NAME)
         self.assertEqual(blob.chunk_size, CHUNK_SIZE)
         self.assertEqual(blob._encryption_key, KEY)
+        self.assertIsNone(blob.kms_key_name)
+
+    def test_blob_w_kms_key_name(self):
+        from google.cloud.storage.blob import Blob
+
+        BUCKET_NAME = 'BUCKET_NAME'
+        BLOB_NAME = 'BLOB_NAME'
+        CHUNK_SIZE = 1024 * 1024
+        KMS_RESOURCE = (
+            "projects/test-project-123/"
+            "locations/global/"
+            "keyRings/test-ring/"
+            "cryptoKeys/test-key/"
+        )
+
+        bucket = self._make_one(name=BUCKET_NAME)
+        blob = bucket.blob(
+            BLOB_NAME, chunk_size=CHUNK_SIZE, kms_key_name=KMS_RESOURCE)
+        self.assertIsInstance(blob, Blob)
+        self.assertIs(blob.bucket, bucket)
+        self.assertIs(blob.client, bucket.client)
+        self.assertEqual(blob.name, BLOB_NAME)
+        self.assertEqual(blob.chunk_size, CHUNK_SIZE)
+        self.assertIsNone(blob._encryption_key)
+        self.assertEqual(blob.kms_key_name, KMS_RESOURCE)
 
     def test_notification_defaults(self):
         from google.cloud.storage.notification import BucketNotification