Upgrade aws-encryption-sdk to work with newer cryptography

ApplauseOSS · May 18, 2022 · f421878 · f421878
1 parent 1379a4d
commit f421878
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 12 deletions.
diff --git a/kmsencryption/lib.py b/kmsencryption/lib.py
@@ -10,13 +10,15 @@
 
 
 def get_key_provider(cmk_arn, profile):
+    cls = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider
     if cmk_arn:
         kms_kwargs = dict(key_ids=[cmk_arn])
     else:
         kms_kwargs = dict()
+        cls = aws_encryption_sdk.DiscoveryAwsKmsMasterKeyProvider
     if profile is not None:
         kms_kwargs['botocore_session'] = botocore.session.Session(profile=profile)
-    return aws_encryption_sdk.KMSMasterKeyProvider(**kms_kwargs)
+    return cls(**kms_kwargs)
 
 
 def decrypt_value(data, prefix, key_provider):
@@ -26,14 +28,20 @@ def decrypt_value(data, prefix, key_provider):
         data = data[len(prefix):]
 
     raw_data = base64.b64decode(data)
-    decrypted_plaintext, decryptor_header = aws_encryption_sdk.decrypt(
+    client = aws_encryption_sdk.EncryptionSDKClient(
+        commitment_policy=aws_encryption_sdk.CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+    )
+    decrypted_plaintext, decryptor_header = client.decrypt(
         source=raw_data,
         key_provider=key_provider)
     return decrypted_plaintext
 
 
 def encrypt_value(data, prefix, key_provider):
-    ciphertext, encryptor_header = aws_encryption_sdk.encrypt(
+    client = aws_encryption_sdk.EncryptionSDKClient(
+        commitment_policy=aws_encryption_sdk.CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+    )
+    ciphertext, encryptor_header = client.encrypt(
         source=data,
         key_provider=key_provider)
     return prefix + base64.b64encode(ciphertext).decode('utf-8')

diff --git a/setup.py b/setup.py
@@ -12,7 +12,7 @@ def finalize_options(self):
 
 setup(
     name='kms-encryption-toolbox',
-    version='0.2.2',
+    version='0.2.3',
     url='https://github.com/ApplauseOSS/kms-encryption-toolbox',
     license='MIT',
     description='Encryption toolbox to be used with the Amazon Key Management Service for securing your deployment secrets. It encapsulates the aws-encryption-sdk package to expose cmdline actions.',
@@ -22,7 +22,7 @@ def finalize_options(self):
     packages=['kmsencryption'],
     install_requires=[
         'cffi>=1.10.0',
-        'aws-encryption-sdk>=1.2.0,<2.0',
+        'aws-encryption-sdk>=3,<4',
         'click>=6.6',
         'cryptography>=1.8.1,!=3.4',
         'future>=0.16.0'

diff --git a/test.py b/test.py
@@ -1,16 +1,20 @@
 import aws_encryption_sdk
 
-kms_key_provider = aws_encryption_sdk.KMSMasterKeyProvider(key_ids=[
+kms_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(key_ids=[
    'arn:aws:kms:us-east-1:873559269338:key/1e1a6a81-93e0-4b9a-954b-cc09802bf3ce'
 ])
 my_plaintext = 'This is some super secret data!  Yup, sure is!'
 
-my_ciphertext, encryptor_header = aws_encryption_sdk.encrypt(
+client = aws_encryption_sdk.EncryptionSDKClient(
+    commitment_policy=aws_encryption_sdk.CommitmentPolicy.FORBID_ENCRYPT_ALLOW_DECRYPT,
+)
+
+my_ciphertext, encryptor_header = client.encrypt(
     source=my_plaintext,
     key_provider=kms_key_provider
 )
 
-decrypted_plaintext, decryptor_header = aws_encryption_sdk.decrypt(
+decrypted_plaintext, decryptor_header = client.decrypt(
     source=my_ciphertext,
     key_provider=kms_key_provider
 )

diff --git a/tox.ini b/tox.ini
@@ -1,6 +1,6 @@
 [tox]
 downloadcache = {toxworkdir}/cache/
-envlist = py27, py3
+envlist = py3
 
 [testenv]
 usedevelop = False
@@ -13,9 +13,6 @@ commands =
     py.cleanup -p -q
     pytest kmsencryption --flake8
 
-[testenv:py27]
-basepython = python2.7
-
 [testenv:py36]
 basepython = python3.6