From 53fbd7dca02f294afb0ffaeac92e793eb8021a95 Mon Sep 17 00:00:00 2001
From: Jack Danger <github@jackcanty.com>
Date: Thu, 9 Feb 2017 19:36:24 -0800
Subject: [PATCH] KMS encryption under Python 3 (#826)

This upgrades the KMS encrypt and decrypt endpoints to work
under both Python 2 and 3
---
 moto/kms/responses.py      |  8 ++++++--
 tests/test_kms/test_kms.py | 32 ++++++++++++++------------------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/moto/kms/responses.py b/moto/kms/responses.py
index fb5c8590..bc928f6f 100644
--- a/moto/kms/responses.py
+++ b/moto/kms/responses.py
@@ -3,6 +3,7 @@ from __future__ import unicode_literals
 import base64
 import json
 import re
+import six
 
 from boto.exception import JSONResponseError
 from boto.kms.exceptions import AlreadyExistsException, NotFoundException
@@ -220,11 +221,14 @@ class KmsResponse(BaseResponse):
         decode it in decrypt().
         """
         value = self.parameters.get("Plaintext")
-        return json.dumps({"CiphertextBlob": base64.b64encode(value).encode("utf-8")})
+        if isinstance(value, six.text_type):
+            value = value.encode('utf-8')
+        return json.dumps({"CiphertextBlob": base64.b64encode(value).decode("utf-8")})
 
     def decrypt(self):
         value = self.parameters.get("CiphertextBlob")
-        return json.dumps({"Plaintext": base64.b64decode(value).encode("utf-8")})
+        print("value 3", value)
+        return json.dumps({"Plaintext": base64.b64decode(value).decode("utf-8")})
 
 
 def _assert_valid_key_id(key_id):
diff --git a/tests/test_kms/test_kms.py b/tests/test_kms/test_kms.py
index 9ec4ffce..04e6fbb4 100644
--- a/tests/test_kms/test_kms.py
+++ b/tests/test_kms/test_kms.py
@@ -1,6 +1,5 @@
 from __future__ import unicode_literals
 import re
-import six
 
 import boto.kms
 from boto.exception import JSONResponseError
@@ -137,25 +136,22 @@ def test_disable_key_rotation():
     conn.get_key_rotation_status(key_id)['KeyRotationEnabled'].should.equal(False)
 
 
-# Scoping encryption/decryption to only Python 2 because our test suite
-# hardcodes a dependency on boto version 2.36.0 which is not compatible with
-# Python 3 (2.40+, however, passes these tests).
-if six.PY2:
-    @mock_kms
-    def test_encrypt():
-        """
-        Using base64 encoding to merely test that the endpoint was called
-        """
-        conn = boto.kms.connect_to_region("us-west-2")
-        response = conn.encrypt('key_id', 'encryptme'.encode('utf-8'))
-        response['CiphertextBlob'].should.equal('ZW5jcnlwdG1l')
+@mock_kms
+def test_encrypt():
+    """
+    test_encrypt
+    Using base64 encoding to merely test that the endpoint was called
+    """
+    conn = boto.kms.connect_to_region("us-west-2")
+    response = conn.encrypt('key_id', 'encryptme'.encode('utf-8'))
+    response['CiphertextBlob'].should.equal(b'ZW5jcnlwdG1l')
 
 
-    @mock_kms
-    def test_decrypt():
-        conn = boto.kms.connect_to_region('us-west-2')
-        response = conn.decrypt('ZW5jcnlwdG1l'.encode('utf-8'))
-        response['Plaintext'].should.equal('encryptme')
+@mock_kms
+def test_decrypt():
+    conn = boto.kms.connect_to_region('us-west-2')
+    response = conn.decrypt('ZW5jcnlwdG1l'.encode('utf-8'))
+    response['Plaintext'].should.equal(b'encryptme')
 
 
 @mock_kms