Add customer-supplied encryption to storage.

gcloud/storage/blob.py

@@ -14,7 +14,9 @@
 
 """Create / interact with Google Cloud Storage blobs."""
 
+import base64
 import copy
+import hashlib
 from io import BytesIO
 import json
 import mimetypes
@@ -109,6 +111,22 @@ def path_helper(bucket_path, blob_name):
         """
         return bucket_path + '/o/' + quote(blob_name, safe='')
 
+    @staticmethod
+    def _get_customer_encryption_headers(key):
+        """Builds customer encyrption key headers
+
+        :type key: str
+        :param key: 32 byte key to build request key and hash.
+        """
+        headers = {}
+        key_hash = base64.encodestring(hashlib.sha256(key).digest())
+        encoded_key = base64.encodestring(key)
+        headers['X-Goog-Encryption-Algorithm'] = 'AES256'
+        headers['X-Goog-Encryption-Key'] = encoded_key.rstrip()
+        headers['X-Goog-Encryption-Key-Sha256'] = key_hash.rstrip()
+
+        return headers
+
     @property
     def acl(self):
         """Create our ACL on demand."""
@@ -276,17 +294,34 @@ def delete(self, client=None):
         """
         return self.bucket.delete_blob(self.name, client=client)
 
-    def download_to_file(self, file_obj, client=None):
+    def download_to_file(self, file_obj, key=None, client=None):
         """Download the contents of this blob into a file-like object.
 
         .. note::
 
            If the server-set property, :attr:`media_link`, is not yet
            initialized, makes an additional API request to load it.
 
+         Downloading a file that has been `customer-supplied
+         <https://cloud.google.com/storage/docs/encryption#customer-supplied>`_
+         encryption::
+
+            >>> from gcloud import storage
+            >>> from gcloud.storage import Blob
+
+            >>> sc = storage.Client(project='my-project')
+            >>> bucket = sc.get_bucket('my-bucket')
+            >>> key = 'aa426195405adee2c8081bb9e7e74b19'
+            >>> blob = Blob('secure-data', bucket)
+            >>> with open('/tmp/my-secure-file', 'w') as file_obj:
+            >>>     blob.download_to_file(file_obj, key=key)
+
         :type file_obj: file
         :param file_obj: A file handle to which to write the blob's data.
 
+        :type key: str
+        :param key: Optional 32 byte key for customer-supplied encryption.
+
         :type client: :class:`gcloud.storage.client.Client` or ``NoneType``
         :param client: Optional. The client to use.  If not passed, falls back
                        to the ``client`` stored on the blob's bucket.
@@ -305,7 +340,11 @@ def download_to_file(self, file_obj, client=None):
         if self.chunk_size is not None:
             download.chunksize = self.chunk_size
 
-        request = Request(download_url, 'GET')
+        headers = {}
+        if key:
+            headers.update(self._get_customer_encryption_headers(key))
+
+        request = Request(download_url, 'GET', headers)
 
         # Use the private ``_connection`` rather than the public
         # ``.connection``, since the public connection may be a batch. A
@@ -315,27 +354,46 @@ def download_to_file(self, file_obj, client=None):
         # it has all three (http, API_BASE_URL and build_api_url).
         download.initialize_download(request, client._connection.http)
 
-    def download_to_filename(self, filename, client=None):
+    def download_to_filename(self, filename, key=None, client=None):
         """Download the contents of this blob into a named file.
 
         :type filename: string
         :param filename: A filename to be passed to ``open``.
 
+        :type key: str
+        :param key: Optional 32 byte key for customer-supplied encryption.
+
         :type client: :class:`gcloud.storage.client.Client` or ``NoneType``
         :param client: Optional. The client to use.  If not passed, falls back
                        to the ``client`` stored on the blob's bucket.
 
         :raises: :class:`gcloud.exceptions.NotFound`
         """
         with open(filename, 'wb') as file_obj:
-            self.download_to_file(file_obj, client=client)
+            self.download_to_file(file_obj, key=key, client=client)
 
         mtime = time.mktime(self.updated.timetuple())
         os.utime(file_obj.name, (mtime, mtime))
 
-    def download_as_string(self, client=None):
+    def download_as_string(self, key=None, client=None):
         """Download the contents of this blob as a string.
 
+         Downloading a blob that has been `customer-supplied
+         <https://cloud.google.com/storage/docs/encryption#customer-supplied>`_
+         encryption::
+
+            >>> from gcloud import storage
+            >>> from gcloud.storage import Blob
+
+            >>> sc = storage.Client(project='my-project')
+            >>> bucket = sc.get_bucket('my-bucket')
+            >>> key = 'aa426195405adee2c8081bb9e7e74b19'
+            >>> blob = Blob('secure-data', bucket)
+            >>> data = blob.download_as_string(file_obj, key=key)
+
+        :type key: str
+        :param key: Optional 32 byte key for customer-supplied encryption.
+
         :type client: :class:`gcloud.storage.client.Client` or ``NoneType``
         :param client: Optional. The client to use.  If not passed, falls back
                        to the ``client`` stored on the blob's bucket.
@@ -345,7 +403,7 @@ def download_as_string(self, client=None):
         :raises: :class:`gcloud.exceptions.NotFound`
         """
         string_buffer = BytesIO()
-        self.download_to_file(string_buffer, client=client)
+        self.download_to_file(string_buffer, key=key, client=client)
         return string_buffer.getvalue()
 
     @staticmethod
@@ -358,7 +416,8 @@ def _check_response_error(request, http_response):
             raise make_exception(faux_response, http_response.content,
                                  error_info=request.url)
 
-    def upload_from_file(self, file_obj, rewind=False, size=None,
+    # pylint: disable=too-many-arguments,too-many-locals
+    def upload_from_file(self, file_obj, rewind=False, size=None, key=None,
                          content_type=None, num_retries=6, client=None):
         """Upload the contents of this blob from a file-like object.
 
@@ -391,6 +450,9 @@ def upload_from_file(self, file_obj, rewind=False, size=None,
                      :func:`os.fstat`. (If the file handle is not from the
                      filesystem this won't be possible.)
 
+        :type key: str
+        :param key: Optional 32 byte key for customer-supplied encryption.
+
         :type content_type: string or ``NoneType``
         :param content_type: Optional type of content being uploaded.
 
@@ -434,6 +496,9 @@ def upload_from_file(self, file_obj, rewind=False, size=None,
             'User-Agent': connection.USER_AGENT,
         }
 
+        if key:
+            headers.update(self._get_customer_encryption_headers(key))
+
         upload = Upload(file_obj, content_type, total_bytes,
                         auto_transfer=False)
 
@@ -474,7 +539,7 @@ def upload_from_file(self, file_obj, rewind=False, size=None,
             response_content = response_content.decode('utf-8')
         self._set_properties(json.loads(response_content))
 
-    def upload_from_filename(self, filename, content_type=None,
+    def upload_from_filename(self, filename, content_type=None, key=None,
                              client=None):
         """Upload this blob's contents from the content of a named file.
 
@@ -500,6 +565,9 @@ def upload_from_filename(self, filename, content_type=None,
         :type content_type: string or ``NoneType``
         :param content_type: Optional type of content being uploaded.
 
+        :type key: str
+        :param key: Optional 32 byte key for customer-supplied encryption.
+
         :type client: :class:`gcloud.storage.client.Client` or ``NoneType``
         :param client: Optional. The client to use.  If not passed, falls back
                        to the ``client`` stored on the blob's bucket.
@@ -509,10 +577,10 @@ def upload_from_filename(self, filename, content_type=None,
             content_type, _ = mimetypes.guess_type(filename)
 
         with open(filename, 'rb') as file_obj:
-            self.upload_from_file(file_obj, content_type=content_type,
+            self.upload_from_file(file_obj, content_type=content_type, key=key,
                                   client=client)
 
-    def upload_from_string(self, data, content_type='text/plain',
+    def upload_from_string(self, data, content_type='text/plain', key=None,
                            client=None):
         """Upload contents of this blob from the provided string.
 
@@ -527,6 +595,19 @@ def upload_from_string(self, data, content_type='text/plain',
            `lifecycle <https://cloud.google.com/storage/docs/lifecycle>`_
            API documents for details.
 
+        Uploading a string that with `customer-supplied
+        <https://cloud.google.com/storage/docs/encryption#customer-supplied>`_
+        encryption::
+
+            >>> from gcloud import storage
+            >>> from gcloud.storage import Blob
+
+            >>> sc = storage.Client(project='my-project')
+            >>> bucket = sc.get_bucket('my-bucket')
+            >>> key = 'aa426195405adee2c8081bb9e7e74b19'
+            >>> blob = Blob('secure-data', bucket)
+            >>> b.upload_from_string('my secure string', key=key)
+
         :type data: bytes or text
         :param data: The data to store in this blob.  If the value is
                      text, it will be encoded as UTF-8.
@@ -535,6 +616,9 @@ def upload_from_string(self, data, content_type='text/plain',
         :param content_type: Optional type of content being uploaded. Defaults
                              to ``'text/plain'``.
 
+        :type key: str
+        :param key: Optional 32 byte key for customer-supplied encryption.
+
         :type client: :class:`gcloud.storage.client.Client` or ``NoneType``
         :param client: Optional. The client to use.  If not passed, falls back
                        to the ``client`` stored on the blob's bucket.
@@ -545,7 +629,7 @@ def upload_from_string(self, data, content_type='text/plain',
         string_buffer.write(data)
         self.upload_from_file(file_obj=string_buffer, rewind=True,
                               size=len(data), content_type=content_type,
-                              client=client)
+                              key=key, client=client)
 
     def make_public(self, client=None):
         """Make this blob public giving all users read access.