From 02ac5ca111828f71928f6fc7dd7de56c23e8e5db Mon Sep 17 00:00:00 2001
From: Macwan Nevil <macnev2013@gmail.com>
Date: Mon, 18 Jan 2021 23:47:13 +0530
Subject: [PATCH] added sse support for s3 (#3592)

* added sse support for s3

* lint fixed for py37
---
 moto/s3/models.py        | 29 ++++++++++++++++++++++++++++-
 moto/s3/responses.py     | 18 +++++++++++++++++-
 tests/test_s3/test_s3.py | 26 ++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/moto/s3/models.py b/moto/s3/models.py
index 17282739a..d86821c4f 100644
--- a/moto/s3/models.py
+++ b/moto/s3/models.py
@@ -99,6 +99,9 @@ class FakeKey(BaseModel):
         max_buffer_size=DEFAULT_KEY_BUFFER_SIZE,
         multipart=None,
         bucket_name=None,
+        encryption=None,
+        kms_key_id=None,
+        bucket_key_enabled=None,
     ):
         self.name = name
         self.last_modified = datetime.datetime.utcnow()
@@ -118,6 +121,10 @@ class FakeKey(BaseModel):
         self.value = value
         self.lock = threading.Lock()
 
+        self.encryption = encryption
+        self.kms_key_id = kms_key_id
+        self.bucket_key_enabled = bucket_key_enabled
+
     @property
     def version_id(self):
         return self._version_id
@@ -229,6 +236,14 @@ class FakeKey(BaseModel):
             "last-modified": self.last_modified_RFC1123,
             "content-length": str(self.size),
         }
+        if self.encryption is not None:
+            res["x-amz-server-side-encryption"] = self.encryption
+            if self.encryption == "aws:kms" and self.kms_key_id is not None:
+                res["x-amz-server-side-encryption-aws-kms-key-id"] = self.kms_key_id
+        if self.bucket_key_enabled is not None:
+            res[
+                "x-amz-server-side-encryption-bucket-key-enabled"
+            ] = self.bucket_key_enabled
         if self._storage_class != "STANDARD":
             res["x-amz-storage-class"] = self._storage_class
         if self._expiry is not None:
@@ -1404,7 +1419,16 @@ class S3Backend(BaseBackend):
         return self.account_public_access_block
 
     def set_object(
-        self, bucket_name, key_name, value, storage=None, etag=None, multipart=None
+        self,
+        bucket_name,
+        key_name,
+        value,
+        storage=None,
+        etag=None,
+        multipart=None,
+        encryption=None,
+        kms_key_id=None,
+        bucket_key_enabled=None,
     ):
         key_name = clean_key_name(key_name)
         if storage is not None and storage not in STORAGE_CLASS:
@@ -1420,6 +1444,9 @@ class S3Backend(BaseBackend):
             is_versioned=bucket.is_versioned,
             version_id=str(uuid.uuid4()) if bucket.is_versioned else None,
             multipart=multipart,
+            encryption=encryption,
+            kms_key_id=kms_key_id,
+            bucket_key_enabled=bucket_key_enabled,
         )
 
         keys = [
diff --git a/moto/s3/responses.py b/moto/s3/responses.py
index 9aa576e6a..d34efc690 100644
--- a/moto/s3/responses.py
+++ b/moto/s3/responses.py
@@ -1261,6 +1261,16 @@ class ResponseObject(_TemplateEnvironmentMixin, ActionAuthenticatorMixin):
             return 200, response_headers, response
 
         storage_class = request.headers.get("x-amz-storage-class", "STANDARD")
+        encryption = request.headers.get("x-amz-server-side-encryption", None)
+        kms_key_id = request.headers.get(
+            "x-amz-server-side-encryption-aws-kms-key-id", None
+        )
+        bucket_key_enabled = request.headers.get(
+            "x-amz-server-side-encryption-bucket-key-enabled", None
+        )
+        if bucket_key_enabled is not None:
+            bucket_key_enabled = str(bucket_key_enabled).lower()
+
         acl = self._acl_from_headers(request.headers)
         if acl is None:
             acl = self.backend.get_bucket(bucket_name).acl
@@ -1343,7 +1353,13 @@ class ResponseObject(_TemplateEnvironmentMixin, ActionAuthenticatorMixin):
         else:
             # Initial data
             new_key = self.backend.set_object(
-                bucket_name, key_name, body, storage=storage_class
+                bucket_name,
+                key_name,
+                body,
+                storage=storage_class,
+                encryption=encryption,
+                kms_key_id=kms_key_id,
+                bucket_key_enabled=bucket_key_enabled,
             )
             request.streaming = True
             metadata = metadata_from_headers(request.headers)
diff --git a/tests/test_s3/test_s3.py b/tests/test_s3/test_s3.py
index fa64c88e4..56fe1698b 100644
--- a/tests/test_s3/test_s3.py
+++ b/tests/test_s3/test_s3.py
@@ -4973,3 +4973,29 @@ def test_request_partial_content_without_specifying_range_should_return_full_obj
     file = s3.Object(bucket, object_key)
     response = file.get(Range="")
     response["ContentLength"].should.equal(30)
+
+
+@mock_s3
+def test_object_headers():
+    bucket = "my-bucket"
+    s3 = boto3.client("s3")
+    s3.create_bucket(Bucket=bucket)
+
+    res = s3.put_object(
+        Bucket=bucket,
+        Body=b"test",
+        Key="file.txt",
+        ServerSideEncryption="aws:kms",
+        SSEKMSKeyId="test",
+        BucketKeyEnabled=True,
+    )
+    res.should.have.key("ETag")
+    res.should.have.key("ServerSideEncryption")
+    res.should.have.key("SSEKMSKeyId")
+    res.should.have.key("BucketKeyEnabled")
+
+    res = s3.get_object(Bucket=bucket, Key="file.txt")
+    res.should.have.key("ETag")
+    res.should.have.key("ServerSideEncryption")
+    res.should.have.key("SSEKMSKeyId")
+    res.should.have.key("BucketKeyEnabled")