From 2adfb593d5ad19320affe480455568c1410b9d93 Mon Sep 17 00:00:00 2001
From: cojenco <59401799+cojenco@users.noreply.github.com>
Date: Tue, 20 Apr 2021 11:33:53 -0700
Subject: [PATCH] feat: add getters and setters for encryption_key and
 kms_key_name (#409)

* fix: add getters and setters for attributes encryption_key and kms_key_name

* fix docstring formatting

* revise docstring

Co-authored-by: Tres Seaver <tseaver@palladion.com>
---
 google/cloud/storage/blob.py | 48 ++++++++++++++++++++++++++++++------
 tests/unit/test_blob.py      | 36 +++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 7 deletions(-)

diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
index 04212ce24..66cc1d153 100644
--- a/google/cloud/storage/blob.py
+++ b/google/cloud/storage/blob.py
@@ -254,6 +254,31 @@ def chunk_size(self, value):
             )
         self._chunk_size = value
 
+    @property
+    def encryption_key(self):
+        """Retrieve the customer-supplied encryption key for the object.
+
+        :rtype: bytes or ``NoneType``
+        :returns:
+            The encryption key or ``None`` if no customer-supplied encryption key was used,
+            or the blob's resource has not been loaded from the server.
+        """
+        return self._encryption_key
+
+    @encryption_key.setter
+    def encryption_key(self, value):
+        """Set the blob's encryption key.
+
+        See https://cloud.google.com/storage/docs/encryption#customer-supplied
+
+        To perform a key rotation for an encrypted blob, use :meth:`rewrite`.
+        See https://cloud.google.com/storage/docs/encryption/using-customer-supplied-keys?hl=ca#rotating
+
+        :type value: bytes
+        :param value: 32 byte encryption key for customer-supplied encryption.
+        """
+        self._encryption_key = value
+
     @staticmethod
     def path_helper(bucket_path, blob_name):
         """Relative URL path for a blob.
@@ -347,25 +372,25 @@ def public_url(self):
     def from_string(cls, uri, client=None):
         """Get a constructor for blob object by URI.
 
-         :type uri: str
-         :param uri: The blob uri pass to get blob object.
+        :type uri: str
+        :param uri: The blob uri pass to get blob object.
 
         :type client: :class:`~google.cloud.storage.client.Client`
         :param client:
             (Optional) The client to use.  If not passed, falls back to the
             ``client`` stored on the blob's bucket.
 
-         :rtype: :class:`google.cloud.storage.blob.Blob`
-         :returns: The blob object created.
+        :rtype: :class:`google.cloud.storage.blob.Blob`
+        :returns: The blob object created.
 
-         Example:
-            Get a constructor for blob object by URI..
+        Example:
+            Get a constructor for blob object by URI.
 
             >>> from google.cloud import storage
             >>> from google.cloud.storage.blob import Blob
             >>> client = storage.Client()
             >>> blob = Blob.from_string("gs://bucket/object")
-         """
+        """
         from google.cloud.storage.bucket import Bucket
 
         scheme, netloc, path, query, frag = urlsplit(uri)
@@ -3839,6 +3864,15 @@ def kms_key_name(self):
         """
         return self._properties.get("kmsKeyName")
 
+    @kms_key_name.setter
+    def kms_key_name(self, value):
+        """Set KMS encryption key for object.
+
+        :type value: str or ``NoneType``
+        :param value: new KMS key name (None to clear any existing key).
+        """
+        self._patch_property("kmsKeyName", value)
+
     storage_class = _scalar_property("storageClass")
     """Retrieve the storage class for the object.
 
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
index e8573ce21..50732a7f0 100644
--- a/tests/unit/test_blob.py
+++ b/tests/unit/test_blob.py
@@ -262,6 +262,42 @@ def test_acl_property(self):
         self.assertIsInstance(acl, ObjectACL)
         self.assertIs(acl, blob._acl)
 
+    def test_encryption_key_getter(self):
+        BLOB_NAME = "blob-name"
+        BUCKET = object()
+        blob = self._make_one(BLOB_NAME, bucket=BUCKET)
+        self.assertIsNone(blob.encryption_key)
+        VALUE = object()
+        blob._encryption_key = VALUE
+        self.assertIs(blob.encryption_key, VALUE)
+
+    def test_encryption_key_setter(self):
+        BLOB_NAME = "blob-name"
+        BUCKET = object()
+        blob = self._make_one(BLOB_NAME, bucket=BUCKET)
+        self.assertIsNone(blob._encryption_key)
+        key = b"12345678901234567890123456789012"
+        blob.encryption_key = key
+        self.assertEqual(blob._encryption_key, key)
+
+    def test_kms_key_name_getter(self):
+        BLOB_NAME = "blob-name"
+        BUCKET = object()
+        blob = self._make_one(BLOB_NAME, bucket=BUCKET)
+        self.assertIsNone(blob.kms_key_name)
+        VALUE = object()
+        blob._patch_property("kmsKeyName", VALUE)
+        self.assertIs(blob.kms_key_name, VALUE)
+
+    def test_kms_key_name_setter(self):
+        BLOB_NAME = "blob-name"
+        BUCKET = object()
+        blob = self._make_one(BLOB_NAME, bucket=BUCKET)
+        self.assertIsNone(blob._properties.get("kmsKeyName"))
+        kms_key_name = "cryptoKeys/test-key"
+        blob.kms_key_name = kms_key_name
+        self.assertEqual(blob._properties.get("kmsKeyName"), kms_key_name)
+
     def test_path_bad_bucket(self):
         fake_bucket = object()
         name = u"blob-name"