addressed PR comments

GoogleCloudPlatform · Sep 27, 2018 · e85865b · e85865b
1 parent 09927c4
commit e85865b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/kms/api-client/asymmetric.py b/kms/api-client/asymmetric.py
@@ -45,10 +45,10 @@ def getAsymmetricPublicKey(client, key_path):
 # [START kms_decrypt_rsa]
 def decryptRSA(ciphertext, client, key_path):
     """
-    Decrypt the given ciphertext using an 'RSA_DECRYPT_OAEP_2048_SHA256' private
-    key stored on Cloud KMS
+    Decrypt the input ciphertext (bytes) using an
+    'RSA_DECRYPT_OAEP_2048_SHA256' private key stored on Cloud KMS
     """
-    request_body = {'ciphertext': ciphertext.decode()}
+    request_body = {'ciphertext': base64.b64encode(ciphertext).decode()}
     request = client.projects() \
                     .locations() \
                     .keyRings() \
@@ -65,15 +65,14 @@ def decryptRSA(ciphertext, client, key_path):
 # [START kms_encrypt_rsa]
 def encryptRSA(plaintext, client, key_path):
     """
-    Encrypt data locally using an 'RSA_DECRYPT_OAEP_2048_SHA256' public
-    key retrieved from Cloud KMS
+    Encrypt the input plaintext (bytes) locally using an
+    'RSA_DECRYPT_OAEP_2048_SHA256' public key retrieved from Cloud KMS
     """
     public_key = getAsymmetricPublicKey(client, key_path)
     pad = padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None)
-    ciphertext = public_key.encrypt(plaintext, pad)
-    return base64.b64encode(ciphertext)
+    return public_key.encrypt(plaintext, pad)
 # [END kms_encrypt_rsa]
 
 

diff --git a/kms/api-client/asymmetric_test.py b/kms/api-client/asymmetric_test.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
 from os import environ
 from time import sleep
 
@@ -102,11 +103,10 @@ def test_rsa_encrypt_decrypt(self):
         ciphertext_bytes = sample.encryptRSA(self.message_bytes,
                                              self.client,
                                              self.rsaDecrypt)
-        ciphertext = ciphertext_bytes.decode('utf-8')
+        ciphertext = base64.b64encode(ciphertext_bytes).decode()
         # ciphertext should be 344 characters with base64 and RSA 2048
         assert len(ciphertext) == 344, \
             'ciphertext should be 344 chars; got {}'.format(len(ciphertext))
-        assert ciphertext[-2:] == '==', 'cipher text should end with =='
         plaintext_bytes = sample.decryptRSA(ciphertext_bytes,
                                             self.client,
                                             self.rsaDecrypt)
@@ -127,7 +127,7 @@ def test_rsa_sign_verify(self):
                                             self.client,
                                             self.rsaSign)
         assert success is True, 'RSA verification failed'
-        changed_bytes = (self.message+".").encode('utf-8')
+        changed_bytes = self.message_bytes + b'.'
         success = sample.verifySignatureRSA(sig,
                                             changed_bytes,
                                             self.client,
@@ -145,7 +145,7 @@ def test_ec_sign_verify(self):
                                            self.client,
                                            self.ecSign)
         assert success is True, 'EC verification failed'
-        changed_bytes = (self.message+".").encode('utf-8')
+        changed_bytes = self.message_bytes + b'.'
         success = sample.verifySignatureEC(sig,
                                            changed_bytes,
                                            self.client,