Allow new S3 keys to be encrypted

* Pulls encrypt_uploads bool from settings
CenterForOpenScience · Aug 3, 2015 · 1e3f29f · 1e3f29f
1 parent 1c6b112
commit 1e3f29f
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 6 deletions.
diff --git a/tests/providers/s3/test_provider.py b/tests/providers/s3/test_provider.py
@@ -35,7 +35,10 @@ def credentials():
 
 @pytest.fixture
 def settings():
-    return {'bucket': 'that kerning'}
+    return {
+        'bucket': 'that kerning',
+        'encrypt_uploads': False
+    }
 
 
 @pytest.fixture
@@ -181,7 +184,8 @@ def file_metadata():
         'Content-Length': 9001,
         'Last-Modified': 'SomeTime',
         'Content-Type': 'binary/octet-stream',
-        'ETag': '"fba9dede5f27731c9771645a39863328"'
+        'ETag': '"fba9dede5f27731c9771645a39863328"',
+        'X-AMZ-SERVER-SIDE-ENCRYPTION': 'AES256'
     }
 
 
@@ -344,6 +348,34 @@ def test_upload_update(self, provider, file_content, file_stream, file_metadata)
         assert aiohttpretty.has_call(method='PUT', uri=url)
         assert aiohttpretty.has_call(method='HEAD', uri=metadata_url)
 
+    @async
+    @pytest.mark.aiohttpretty
+    def test_upload_encrypted(self, provider, file_content, file_stream, file_metadata):
+
+        # Set trigger for encrypt_key=True in s3.provider.upload
+        provider.encrypt_uploads = True
+        path = WaterButlerPath('/foobah')
+        content_md5 = hashlib.md5(file_content).hexdigest()
+        url = provider.bucket.new_key(path.path).generate_url(100, 'PUT', encrypt_key=True)
+        metadata_url = provider.bucket.new_key(path.path).generate_url(100, 'HEAD')
+        aiohttpretty.register_uri(
+            'HEAD',
+            metadata_url,
+            responses=[
+                {'status': 404},
+                {'headers': file_metadata},
+            ],
+        )
+        aiohttpretty.register_uri('PUT', url, status=200, headers={'ETag': '"{}"'.format(content_md5)})
+
+        metadata, created = yield from provider.upload(file_stream, path)
+
+        assert metadata.kind == 'file'
+        assert metadata.extra['encryption'] == 'AES256'
+        assert created
+        assert aiohttpretty.has_call(method='PUT', uri=url)
+        assert aiohttpretty.has_call(method='HEAD', uri=metadata_url)
+
     @async
     @pytest.mark.aiohttpretty
     def test_delete(self, provider):

diff --git a/waterbutler/providers/s3/metadata.py b/waterbutler/providers/s3/metadata.py
@@ -45,7 +45,8 @@ def etag(self):
     @property
     def extra(self):
         return {
-            'md5': self.raw['ETAG'].replace('"', '')
+            'md5': self.raw['ETAG'].replace('"', ''),
+            'encryption': self.raw.get('X-AMZ-SERVER-SIDE-ENCRYPTION', '')
         }
 
 

diff --git a/waterbutler/providers/s3/provider.py b/waterbutler/providers/s3/provider.py
@@ -51,6 +51,7 @@ def __init__(self, auth, credentials, settings):
         self.connection = S3Connection(credentials['access_key'],
                 credentials['secret_key'], calling_format=calling_format)
         self.bucket = self.connection.get_bucket(settings['bucket'], validate=False)
+        self.encrypt_uploads = self.settings.get('encrypt_uploads', False)
 
     @asyncio.coroutine
     def validate_path(self, path, **kwargs):
@@ -140,18 +141,20 @@ def upload(self, stream, path, conflict='replace', **kwargs):
         :rtype: dict, bool
         """
         path, exists = yield from self.handle_name_conflict(path, conflict=conflict)
-
         stream.add_writer('md5', streams.HashStreamWriter(hashlib.md5))
 
         resp = yield from self.make_request(
             'PUT',
-            self.bucket.new_key(path.path).generate_url(settings.TEMP_URL_SECS, 'PUT'),
+            self.bucket.new_key(path.path).generate_url(
+                settings.TEMP_URL_SECS,
+                'PUT',
+                encrypt_key=self.encrypt_uploads
+            ),
             data=stream,
             headers={'Content-Length': str(stream.size)},
             expects=(200, 201, ),
             throws=exceptions.UploadError,
         )
-
         # md5 is returned as ETag header as long as server side encryption is not used.
         # TODO: nice assertion error goes here
         assert resp.headers['ETag'].replace('"', '') == stream.writers['md5'].hexdigest