Support fetching GCS bucket encryption metadata.

7 files changed, 218 insertions, 7 deletions

diff --git a/boto/exception.py b/boto/exception.py
index 2f175979..6db5d48d 100644
--- a/boto/exception.py
+++ b/boto/exception.py
@@ -475,9 +475,12 @@ class InvalidCorsError(Exception):
         self.message = message
 
 
-class NoAuthHandlerFound(Exception):
-    """Is raised when no auth handlers were found ready to authenticate."""
-    pass
+class InvalidEncryptionConfigError(Exception):
+    """Exception raised when GCS encryption configuration XML is invalid."""
+
+    def __init__(self, message):
+        super(InvalidEncryptionConfigError, self).__init__(message)
+        self.message = message
 
 
 class InvalidLifecycleConfigError(Exception):
@@ -488,6 +491,11 @@ class InvalidLifecycleConfigError(Exception):
         self.message = message
 
 
+class NoAuthHandlerFound(Exception):
+    """Is raised when no auth handlers were found ready to authenticate."""
+    pass
+
+
 # Enum class for resumable upload failure disposition.
 class ResumableTransferDisposition(object):
     # START_OVER means an attempt to resume an existing transfer failed,
diff --git a/boto/gs/bucket.py b/boto/gs/bucket.py
index dcce1d08..8e56dedb 100644
--- a/boto/gs/bucket.py
+++ b/boto/gs/bucket.py
@@ -32,6 +32,7 @@ from boto.gs.acl import ACL, CannedACLStrings
 from boto.gs.acl import SupportedPermissions as GSPermissions
 from boto.gs.bucketlistresultset import VersionedBucketListResultSet
 from boto.gs.cors import Cors
+from boto.gs.encryptionconfig import EncryptionConfig
 from boto.gs.lifecycle import LifecycleConfig
 from boto.gs.key import Key as GSKey
 from boto.s3.acl import Policy
@@ -43,6 +44,7 @@ from boto.compat import six
 DEF_OBJ_ACL = 'defaultObjectAcl'
 STANDARD_ACL = 'acl'
 CORS_ARG = 'cors'
+ENCRYPTION_CONFIG_ARG = 'encryptionConfig'
 LIFECYCLE_ARG = 'lifecycle'
 STORAGE_CLASS_ARG='storageClass'
 ERROR_DETAILS_REGEX = re.compile(r'<Details>(?P<details>.*)</Details>')
@@ -51,12 +53,19 @@ class Bucket(S3Bucket):
     """Represents a Google Cloud Storage bucket."""
 
     BillingBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
-                   '<BillingConfiguration><RequesterPays>%s</RequesterPays>'
+                   '<BillingConfiguration>'
+                   '<RequesterPays>%s</RequesterPays>'
                    '</BillingConfiguration>')
+    EncryptionConfigBody = (
+        '<?xml version="1.0" encoding="UTF-8"?>\n'
+        '<EncryptionConfiguration>%s</EncryptionConfiguration>')
+    EncryptionConfigDefaultKeyNameFragment = (
+        '<DefaultKmsKeyName>%s</DefaultKmsKeyName>')
     StorageClassBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
                         '<StorageClass>%s</StorageClass>')
     VersioningBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
-                      '<VersioningConfiguration><Status>%s</Status>'
+                      '<VersioningConfiguration>'
+                      '<Status>%s</Status>'
                       '</VersioningConfiguration>')
     WebsiteBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
                    '<WebsiteConfiguration>%s%s</WebsiteConfiguration>')
@@ -1065,3 +1074,61 @@ class Bucket(S3Bucket):
         else:
             req_body = self.BillingBody % ('Disabled')
         self.set_subresource('billing', req_body, headers=headers)
+
+    def get_encryption_config(self, headers=None):
+        """Returns a bucket's EncryptionConfig.
+
+        :param dict headers: Additional headers to send with the request.
+        :rtype: :class:`~.encryption_config.EncryptionConfig`
+        """
+        response = self.connection.make_request(
+            'GET', self.name, query_args=ENCRYPTION_CONFIG_ARG, headers=headers)
+        body = response.read()
+        if response.status == 200:
+            # Success - parse XML and return EncryptionConfig object.
+            encryption_config = EncryptionConfig()
+            h = handler.XmlHandler(encryption_config, self)
+            xml.sax.parseString(body, h)
+            return encryption_config
+        else:
+            raise self.connection.provider.storage_response_error(
+                response.status, response.reason, body)
+
+    def _construct_encryption_config_xml(self, default_kms_key_name=None):
+        """Creates an XML document for setting a bucket's EncryptionConfig.
+
+        This method is internal as it's only here for testing purposes. As
+        managing Cloud KMS resources for testing is complex, we settle for
+        testing that we're creating correctly-formed XML for setting a bucket's
+        encryption configuration.
+
+        :param str default_kms_key_name: A string containing a fully-qualified
+            Cloud KMS key name.
+        :rtype: str
+        """
+        if default_kms_key_name:
+            default_kms_key_name_frag = (
+                self.EncryptionConfigDefaultKeyNameFragment %
+                default_kms_key_name)
+        else:
+            default_kms_key_name_frag = ''
+
+        return self.EncryptionConfigBody % default_kms_key_name_frag
+
+
+    def set_encryption_config(self, default_kms_key_name=None, headers=None):
+        """Sets a bucket's EncryptionConfig XML document.
+
+        :param str default_kms_key_name: A string containing a fully-qualified
+            Cloud KMS key name.
+        :param dict headers: Additional headers to send with the request.
+        """
+        body = self._construct_encryption_config_xml(
+            default_kms_key_name=default_kms_key_name)
+        response = self.connection.make_request(
+            'PUT', get_utf8_value(self.name), data=get_utf8_value(body),
+            query_args=ENCRYPTION_CONFIG_ARG, headers=headers)
+        body = response.read()
+        if response.status != 200:
+            raise self.connection.provider.storage_response_error(
+                response.status, response.reason, body)
diff --git a/boto/gs/encryptionconfig.py b/boto/gs/encryptionconfig.py
new file mode 100755
index 00000000..b6dd18b4
--- /dev/null
+++ b/boto/gs/encryptionconfig.py
@@ -0,0 +1,76 @@
+# Copyright 2018 Google Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish, dis-
+# tribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the fol-
+# lowing conditions:
+#
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
+# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+
+import types
+from boto.gs.user import User
+from boto.exception import InvalidEncryptionConfigError
+from xml.sax import handler
+
+# Relevant tags for the EncryptionConfiguration XML document.
+DEFAULT_KMS_KEY_NAME = 'DefaultKmsKeyName'
+ENCRYPTION_CONFIG = 'EncryptionConfiguration'
+
+class EncryptionConfig(handler.ContentHandler):
+    """Encapsulates the EncryptionConfiguration XML document"""
+    def __init__(self):
+        # Valid items in an EncryptionConfiguration XML node.
+        self.default_kms_key_name = None
+
+        self.parse_level = 0
+
+    def validateParseLevel(self, tag, level):
+        """Verify parse level for a given tag."""
+        if self.parse_level != level:
+            raise InvalidEncryptionConfigError(
+                'Invalid tag %s at parse level %d: ' % (tag, self.parse_level))
+
+    def startElement(self, name, attrs, connection):
+        """SAX XML logic for parsing new element found."""
+        if name == ENCRYPTION_CONFIG:
+            self.validateParseLevel(name, 0)
+            self.parse_level += 1;
+        elif name == DEFAULT_KMS_KEY_NAME:
+            self.validateParseLevel(name, 1)
+            self.parse_level += 1;
+        else:
+            raise InvalidEncryptionConfigError('Unsupported tag ' + name)
+
+    def endElement(self, name, value, connection):
+        """SAX XML logic for parsing new element found."""
+        if name == ENCRYPTION_CONFIG:
+            self.validateParseLevel(name, 1)
+            self.parse_level -= 1;
+        elif name == DEFAULT_KMS_KEY_NAME:
+            self.validateParseLevel(name, 2)
+            self.parse_level -= 1;
+            self.default_kms_key_name = value.strip()
+        else:
+            raise InvalidEncryptionConfigError('Unsupported end tag ' + name)
+
+    def to_xml(self):
+        """Convert EncryptionConfig object into XML string representation."""
+        s = ['<%s>' % ENCRYPTION_CONFIG]
+        if self.default_kms_key_name:
+            s.append('<%s>%s</%s>' % (DEFAULT_KMS_KEY_NAME,
+                                      self.default_kms_key_name,
+                                      DEFAULT_KMS_KEY_NAME))
+        s.append('</%s>' % ENCRYPTION_CONFIG)
+        return ''.join(s)
diff --git a/boto/storage_uri.py b/boto/storage_uri.py
index feecc0f0..5202b9ff 100755
--- a/boto/storage_uri.py
+++ b/boto/storage_uri.py
@@ -821,6 +821,25 @@ class BucketStorageUri(StorageUri):
         bucket = self.get_bucket(validate, headers)
         bucket.configure_billing(requester_pays=requester_pays, headers=headers)
 
+    def get_encryption_config(self, validate=False, headers=None):
+        """Returns a GCS bucket's encryption configuration."""
+        self._check_bucket_uri('get_encryption_config')
+        # EncryptionConfiguration is defined as a bucket param for GCS, but not
+        # for S3.
+        if self.scheme != 'gs':
+            raise ValueError('get_encryption_config() not supported for %s '
+                             'URIs.' % self.scheme)
+        bucket = self.get_bucket(validate, headers)
+        return bucket.get_encryption_config(headers=headers)
+
+    def set_encryption_config(self, default_kms_key_name=None, validate=False,
+                              headers=None):
+        """Sets a GCS bucket's encryption configuration."""
+        self._check_bucket_uri('set_encryption_config')
+        bucket = self.get_bucket(validate, headers)
+        bucket.set_encryption_config(default_kms_key_name=default_kms_key_name,
+                                     headers=headers)
+
     def exists(self, headers=None):
         """Returns True if the object exists or False if it doesn't"""
         if not self.object_name:
diff --git a/boto/utils.py b/boto/utils.py
index f8801817..22f97110 100644
--- a/boto/utils.py
+++ b/boto/utils.py
@@ -93,7 +93,10 @@ qsa_of_interest = ['acl', 'cors', 'defaultObjectAcl', 'location', 'logging',
                    # billing is a QSA for buckets in Google Cloud Storage.
                    'billing',
                    # userProject is a QSA for requests in Google Cloud Storage.
-                   'userProject']
+                   'userProject',
+                   # encryptionConfig is a QSA for requests in Google Cloud
+                   # Storage.
+                   'encryptionConfig']
 
 
 _first_cap_regex = re.compile('(.)([A-Z][a-z]+)')
diff --git a/tests/integration/gs/test_basic.py b/tests/integration/gs/test_basic.py
index f26f87a2..eb2f2707 100644
--- a/tests/integration/gs/test_basic.py
+++ b/tests/integration/gs/test_basic.py
@@ -51,6 +51,12 @@ CORS_DOC = ('<CorsConfig><Cors><Origins><Origin>origin1.example.com'
             '<ResponseHeader>bar</ResponseHeader></ResponseHeaders>'
             '</Cors></CorsConfig>')
 
+ENCRYPTION_CONFIG_WITH_KEY = (
+    '<?xml version="1.0" encoding="UTF-8"?>\n'
+    '<EncryptionConfiguration>'
+    '<DefaultKmsKeyName>%s</DefaultKmsKeyName>'
+    '</EncryptionConfiguration>')
+
 LIFECYCLE_EMPTY = ('<?xml version="1.0" encoding="UTF-8"?>'
                    '<LifecycleConfiguration></LifecycleConfiguration>')
 LIFECYCLE_DOC = ('<?xml version="1.0" encoding="UTF-8"?>'
@@ -491,3 +497,34 @@ class GSBasicTest(GSTestCase):
         uri.configure_billing(requester_pays=False)
         billing = uri.get_billing_config()
         self.assertEqual(billing, BILLING_DISABLED)
+
+    def test_encryption_config_bucket(self):
+        """Test setting and getting of EncryptionConfig on gs Bucket objects."""
+        # Create a new bucket.
+        bucket = self._MakeBucket()
+        bucket_name = bucket.name
+        # Get EncryptionConfig and make sure it's empty.
+        encryption_config = bucket.get_encryption_config()
+        self.assertIsNone(encryption_config.default_kms_key_name)
+        # Testing set functionality would require having an existing Cloud KMS
+        # key. Since we can't hardcode a key name or dynamically create one, we
+        # only test here that we're creating the correct XML document to send to
+        # GCS.
+        xmldoc = bucket._construct_encryption_config_xml(
+            default_kms_key_name='dummykey')
+        self.assertEqual(xmldoc, ENCRYPTION_CONFIG_WITH_KEY % 'dummykey')
+        # Test that setting an empty encryption config works.
+        bucket.set_encryption_config()
+
+    def test_encryption_config_storage_uri(self):
+        """Test setting and getting of EncryptionConfig with storage_uri."""
+        # Create a new bucket.
+        bucket = self._MakeBucket()
+        bucket_name = bucket.name
+        uri = storage_uri('gs://' + bucket_name)
+        # Get EncryptionConfig and make sure it's empty.
+        encryption_config = uri.get_encryption_config()
+        self.assertIsNone(encryption_config.default_kms_key_name)
+
+        # Test that setting an empty encryption config works.
+        uri.set_encryption_config()
diff --git a/tests/integration/s3/test_key.py b/tests/integration/s3/test_key.py
index 471857a7..9fb0db94 100644
--- a/tests/integration/s3/test_key.py
+++ b/tests/integration/s3/test_key.py
@@ -423,7 +423,8 @@ class S3KeyTest(unittest.TestCase):
             check.cache_control,
             ('public,%20max-age=500', 'public, max-age=500')
         )
-        self.assertEqual(remote_metadata['cache-control'], 'public,%20max-age=500')
+        self.assertIn(remote_metadata['cache-control'],
+                      ('public,%20max-age=500', 'public, max-age=500'))
         self.assertEqual(check.get_metadata('test-plus'), 'A plus (+)')
         self.assertEqual(check.content_disposition, 'filename=Sch%C3%B6ne%20Zeit.txt')
         self.assertEqual(remote_metadata['content-disposition'], 'filename=Sch%C3%B6ne%20Zeit.txt')