v0.0.13 - adding support for referencing the library in Python code

Author: Tomasz Rogozik

kmsencryption/__init__.py

@@ -0,0 +1,6 @@
+from kmsencryption.lib import encrypt
+from kmsencryption.lib import decrypt
+from kmsencryption.lib import decrypt_json
+from kmsencryption.lib import encrypt_json
+
+__all__ = ['encrypt', 'decrypt', 'decrypt_json', 'encrypt_json']

kmsencryption/__main__.py

@@ -1,59 +1,23 @@
-import aws_encryption_sdk
-import base64
-import botocore.session
+
 import click
-import json
-import os
 import sys
 
+from kmsencryption import lib
+
 
 @click.group(context_settings={"help_option_names": ['-h', '--help']})
 def main():
     pass
 
 
-def get_key_provider(cmk_arn, profile):
-    if cmk_arn:
-        kms_kwargs = dict(key_ids=[cmk_arn])
-    else:
-        kms_kwargs = dict()
-    if profile is not None:
-        kms_kwargs['botocore_session'] = botocore.session.Session(profile=profile)
-    return aws_encryption_sdk.KMSMasterKeyProvider(**kms_kwargs)
-
-
-def decrypt_value(data, prefix, key_provider):
-    if data.startswith(prefix):
-        data = data[len(prefix):]
-
-    raw_data = base64.b64decode(data)
-    decrypted_plaintext, decryptor_header = aws_encryption_sdk.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(
-        source=data,
-        key_provider=key_provider)
-    return prefix + base64.b64encode(ciphertext).decode('utf-8')
-
-
 @main.command(help='Encrypts data with a new data key and returns a base64-encoded result.')
 @click.option('--cmk-arn', 'cmk_arn', prompt=True, help='ARN of an existing Customer Master Key in KMS')
 @click.option('--data', 'data', envvar='DATA', help='Data to be encrypted. Use to pass it as a named argument.')
 @click.option('--env', 'env', help='Name of an environment variable that contains data to be encrypted.')
 @click.option('--profile', 'profile', default=None, help='Name of an AWS CLI profile to be used when contacting AWS.')
 @click.option('--prefix', 'prefix', default='', help='An output prefix to be added to the generated result.')
 def encrypt(cmk_arn, data, env, profile, prefix):
-    kms_key_provider = get_key_provider(cmk_arn, profile)
-    if env is not None:
-        data = os.getenv(env, data)
-    if not data:
-        raise ValueError('No data provided via --data or in a variable name passed with --env')
-
-    click.echo(encrypt_value(data, prefix, kms_key_provider))
+    click.echo(lib.encrypt(cmk_arn, data, env, profile, prefix))
 
 
 @main.command(help='Decrypts a base64-encoded data.')
@@ -63,13 +27,7 @@ def encrypt(cmk_arn, data, env, profile, prefix):
 @click.option('--prefix', 'prefix', default='',
               help='An input prefix to be trimmed from the beginning before a value is decrypted.')
 def decrypt(data, env, profile, prefix):
-    kms_key_provider = get_key_provider(None, profile)
-    if env is not None:
-        data = os.getenv(env, data)
-    if not data:
-        raise ValueError('No data provided via --data or in a variable name passed with --env')
-
-    click.echo(decrypt_value(data, prefix, kms_key_provider))
+    click.echo(lib.decrypt(data, env, profile, prefix))
 
 
 @main.command('decrypt-json',
@@ -80,12 +38,7 @@ def decrypt(data, env, profile, prefix):
 @click.option('--prefix', 'prefix', default='',
               help='An input prefix to be trimmed from the beginning before a value is decrypted.')
 def decrypt_json(input, profile, prefix):
-    kms_key_provider = get_key_provider(None, profile)
-    input_map = json.load(input)
-    output = {}
-    for name, value in input_map.iteritems():
-        output[name] = decrypt_value(value, prefix, kms_key_provider) if value.startswith(prefix) else value
-    click.echo(json.dumps(output))
+    click.echo(lib.decrypt_json(input, profile, prefix))
 
 
 @main.command('encrypt-json',
@@ -97,12 +50,7 @@ def decrypt_json(input, profile, prefix):
 @click.option('--prefix', 'prefix', default='',
               help='An output prefix to be added to the beginning of an encrypted value.')
 def encrypt_json(input, cmk_arn, profile, prefix):
-    kms_key_provider = get_key_provider(cmk_arn, profile)
-    input_map = json.load(input)
-    output = {}
-    for name, value in input_map.iteritems():
-        output[name] = encrypt_value(value, prefix, kms_key_provider)
-    click.echo(json.dumps(output))
+    click.echo(lib.encrypt_json(input, cmk_arn, profile, prefix))
 
 
 if __name__ == '__main__':

kmsencryption/lib.py

@@ -0,0 +1,70 @@
+import aws_encryption_sdk
+import base64
+import botocore.session
+import json
+import os
+
+def get_key_provider(cmk_arn, profile):
+    if cmk_arn:
+        kms_kwargs = dict(key_ids=[cmk_arn])
+    else:
+        kms_kwargs = dict()
+    if profile is not None:
+        kms_kwargs['botocore_session'] = botocore.session.Session(profile=profile)
+    return aws_encryption_sdk.KMSMasterKeyProvider(**kms_kwargs)
+
+
+def decrypt_value(data, prefix, key_provider):
+    if data.startswith(prefix):
+        data = data[len(prefix):]
+
+    raw_data = base64.b64decode(data)
+    decrypted_plaintext, decryptor_header = aws_encryption_sdk.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(
+        source=data,
+        key_provider=key_provider)
+    return prefix + base64.b64encode(ciphertext).decode('utf-8')
+
+
+def encrypt(cmk_arn, data, env, profile, prefix):
+    kms_key_provider = get_key_provider(cmk_arn, profile)
+    if env is not None:
+        data = os.getenv(env, data)
+    if not data:
+        raise ValueError('No data provided via --data or in a variable name passed with --env')
+
+    return encrypt_value(data, prefix, kms_key_provider)
+
+
+def decrypt(data, env, profile, prefix):
+    kms_key_provider = get_key_provider(None, profile)
+    if env is not None:
+        data = os.getenv(env, data)
+    if not data:
+        raise ValueError('No data provided via --data or in a variable name passed with --env')
+
+    return decrypt_value(data, prefix, kms_key_provider)
+
+
+def decrypt_json(input, profile, prefix):
+    kms_key_provider = get_key_provider(None, profile)
+    input_map = json.load(input)
+    output = {}
+    for name, value in input_map.iteritems():
+        output[name] = decrypt_value(value, prefix, kms_key_provider) if value.startswith(prefix) else value
+    return json.dumps(output)
+
+
+def encrypt_json(input, cmk_arn, profile, prefix):
+    kms_key_provider = get_key_provider(cmk_arn, profile)
+    input_map = json.load(input)
+    output = {}
+    for name, value in input_map.iteritems():
+        output[name] = encrypt_value(value, prefix, kms_key_provider)
+    return json.dumps(output)

setup.py

@@ -3,14 +3,14 @@
 
 setup(
     name='kms-encryption-toolbox',
-    version='0.0.12',
+    version='0.0.13',
     url='https://github.com/ApplauseOSS/kms-encryption-toolbox',
     license='Applause',
     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.',
     author='Applause',
     author_email='[email protected]',
     zip_safe=False,
-    packages=find_packages(),
+    packages=['kmsencryption'],
     install_requires=[
         'cffi>=1.10.0',
         'aws-encryption-sdk>=1.2.0',