text: # snippet-comment:[These are tags for the AWS doc team's sample catalog. Do not remove.]
  # snippet-sourcedescription:[encrypt_decrypt_file.py demonstrates how to encrypt and decrypt a file using the AWS Key Management Service]
  # snippet-service:[kms]
  # snippet-keyword:[AWS Key Management Service (KMS)]
  # snippet-keyword:[Python]
  # snippet-sourcesyntax:[python]
  # snippet-sourcesyntax:[python]
  # snippet-sourcetype:[full-example]
  # snippet-sourcedate:[2019-03-12]
  # snippet-sourceauthor:[AWS]
  
  # Copyright 2010-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  #
  # This file is licensed under the Apache License, Version 2.0 (the "License").
  # You may not use this file except in compliance with the License. A copy of the
  # License is located at
  #
  # http://aws.amazon.com/apache2.0/
  #
  # This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS
  # OF ANY KIND, either express or implied. See the License for the specific
  # language governing permissions and limitations under the License.
  
  import base64
  import logging
  import boto3
  from botocore.exceptions import ClientError
  
  # To perform the optional file encryption/decryption operations, the Python
  # cryptography package must be installed.
  #       pip install cryptography
  from cryptography.fernet import Fernet
  
  
  # snippet-start:[kms.python.retrieve-cmk]
  def retrieve_cmk(desc):
      """Retrieve an existing KMS CMK based on its description
  
      :param desc: Description of CMK specified when the CMK was created
      :return Tuple(KeyId, KeyArn) where:
          KeyId: CMK ID
          KeyArn: Amazon Resource Name of CMK
      :return Tuple(None, None) if a CMK with the specified description was
      not found
      """
  
      # Retrieve a list of existing CMKs
      # If more than 100 keys exist, retrieve and process them in batches
      kms_client = boto3.client('kms')
      try:
          response = kms_client.list_keys()
      except ClientError as e:
          logging.error(e)
          return None, None
  
      done = False
      while not done:
          for cmk in response['Keys']:
              # Get info about the key, including its description
              try:
                  key_info = kms_client.describe_key(KeyId=cmk['KeyArn'])
              except ClientError as e:
                  logging.error(e)
                  return None, None
  
              # Is this the key we're looking for?
              if key_info['KeyMetadata']['Description'] == desc:
                  return cmk['KeyId'], cmk['KeyArn']
  
          # Are there more keys to retrieve?
          if not response['Truncated']:
              # No, the CMK was not found
              logging.debug('A CMK with the specified description was not found')
              done = True
          else:
              # Yes, retrieve another batch
              try:
                  response = kms_client.list_keys(Marker=response['NextMarker'])
              except ClientError as e:
                  logging.error(e)
                  return None, None
  
      # All existing CMKs were checked and the desired key was not found
      return None, None
  # snippet-end:[kms.python.retrieve-cmk]
  
  
  # snippet-start:[kms.python.create_cmk]
  def create_cmk(desc='Customer Master Key'):
      """Create a KMS Customer Master Key
  
      The created CMK is a Customer-managed key stored in AWS KMS.
  
      :param desc: key description
      :return Tuple(KeyId, KeyArn) where:
          KeyId: AWS globally-unique string ID
          KeyArn: Amazon Resource Name of the CMK
      :return Tuple(None, None) if error
      """
  
      # Create CMK
      kms_client = boto3.client('kms')
      try:
          response = kms_client.create_key(Description=desc)
      except ClientError as e:
          logging.error(e)
          return None, None
  
      # Return the key ID and ARN
      return response['KeyMetadata']['KeyId'], response['KeyMetadata']['Arn']
  # snippet-end:[kms.python.create_cmk]
  
  
  # snippet-start:[kms.python.create_data_key]
  def create_data_key(cmk_id, key_spec='AES_256'):
      """Generate a data key to use when encrypting and decrypting data
  
      :param cmk_id: KMS CMK ID or ARN under which to generate and encrypt the
      data key.
      :param key_spec: Length of the data encryption key. Supported values:
          'AES_128': Generate a 128-bit symmetric key
          'AES_256': Generate a 256-bit symmetric key
      :return Tuple(EncryptedDataKey, PlaintextDataKey) where:
          EncryptedDataKey: Encrypted CiphertextBlob data key as binary string
          PlaintextDataKey: Plaintext base64-encoded data key as binary string
      :return Tuple(None, None) if error
      """
  
      # Create data key
      kms_client = boto3.client('kms')
      try:
          response = kms_client.generate_data_key(KeyId=cmk_id, KeySpec=key_spec)
      except ClientError as e:
          logging.error(e)
          return None, None
  
      # Return the encrypted and plaintext data key
      return response['CiphertextBlob'], base64.b64encode(response['Plaintext'])
  # snippet-end:[kms.python.create_data_key]
  
  
  # snippet-start:[kms.python.decrypt_data_key]
  def decrypt_data_key(data_key_encrypted):
      """Decrypt an encrypted data key
  
      :param data_key_encrypted: Encrypted ciphertext data key.
      :return Plaintext base64-encoded binary data key as binary string
      :return None if error
      """
  
      # Decrypt the data key
      kms_client = boto3.client('kms')
      try:
          response = kms_client.decrypt(CiphertextBlob=data_key_encrypted)
      except ClientError as e:
          logging.error(e)
          return None
  
      # Return plaintext base64-encoded binary data key
      return base64.b64encode((response['Plaintext']))
  # snippet-end:[kms.python.decrypt_data_key]
  
  
  # Number of bytes used in the encrypted file to store the length of the
  # encrypted data key. Used by encrypt_file() and decrypt_file().
  NUM_BYTES_FOR_LEN = 4
  
  
  # snippet-start:[kms.python.encrypt_file]
  def encrypt_file(filename, cmk_id):
      """Encrypt a file using an AWS KMS CMK
  
      A data key is generated and associated with the CMK.
      The encrypted data key is saved with the encrypted file. This enables the
      file to be decrypted at any time in the future and by any program that
      has the credentials to decrypt the data key.
      The encrypted file is saved to <filename>.encrypted
      Limitation: The contents of filename must fit in memory.
  
      :param filename: File to encrypt
      :param cmk_id: AWS KMS CMK ID or ARN
      :return: True if file was encrypted. Otherwise, False.
      """
  
      # Read the entire file into memory
      try:
          with open(filename, 'rb') as file:
              file_contents = file.read()
      except IOError as e:
          logging.error(e)
          return False
  
      # Generate a data key associated with the CMK
      # The data key is used to encrypt the file. Each file can use its own
      # data key or data keys can be shared among files.
      # Specify either the CMK ID or ARN
      data_key_encrypted, data_key_plaintext = create_data_key(cmk_id)
      if data_key_encrypted is None:
          return False
      logging.info('Created new AWS KMS data key')
  
      # Encrypt the file
      f = Fernet(data_key_plaintext)
      file_contents_encrypted = f.encrypt(file_contents)
  
      # Write the encrypted data key and encrypted file contents together
      try:
          with open(filename + '.encrypted', 'wb') as file_encrypted:
              file_encrypted.write(len(data_key_encrypted).to_bytes(NUM_BYTES_FOR_LEN,
                                                                    byteorder='big'))
              file_encrypted.write(data_key_encrypted)
              file_encrypted.write(file_contents_encrypted)
      except IOError as e:
          logging.error(e)
          return False
  
      # For the highest security, the data_key_plaintext value should be wiped
      # from memory. Unfortunately, this is not possible in Python. However,
      # storing the value in a local variable makes it available for garbage
      # collection.
      return True
  # snippet-end:[kms.python.encrypt_file]
  
  
  # snippet-start:[kms.python.decrypt_file]
  def decrypt_file(filename):
      """Decrypt a file encrypted by encrypt_file()
  
      The encrypted file is read from <filename>.encrypted
      The decrypted file is written to <filename>.decrypted
  
      :param filename: File to decrypt
      :return: True if file was decrypted. Otherwise, False.
      """
  
      # Read the encrypted file into memory
      try:
          with open(filename + '.encrypted', 'rb') as file:
              file_contents = file.read()
      except IOError as e:
          logging.error(e)
          return False
  
      # The first NUM_BYTES_FOR_LEN bytes contain the integer length of the
      # encrypted data key.
      # Add NUM_BYTES_FOR_LEN to get index of end of encrypted data key/start
      # of encrypted data.
      data_key_encrypted_len = int.from_bytes(file_contents[:NUM_BYTES_FOR_LEN],
                                              byteorder='big') \
                               + NUM_BYTES_FOR_LEN
      data_key_encrypted = file_contents[NUM_BYTES_FOR_LEN:data_key_encrypted_len]
  
      # Decrypt the data key before using it
      data_key_plaintext = decrypt_data_key(data_key_encrypted)
      if data_key_plaintext is None:
          logging.error("Cannot decrypt the data key")
          return False
  
      # Decrypt the rest of the file
      f = Fernet(data_key_plaintext)
      file_contents_decrypted = f.decrypt(file_contents[data_key_encrypted_len:])
  
      # Write the decrypted file contents
      try:
          with open(filename + '.decrypted', 'wb') as file_decrypted:
              file_decrypted.write(file_contents_decrypted)
      except IOError as e:
          logging.error(e)
          return False
  
      # The same security issue described at the end of encrypt_file() exists
      # here, too, i.e., the wish to wipe the data_key_plaintext value from
      # memory.
      return True
  # snippet-end:[kms.python.decrypt_file]
  
  
  def main():
      """Exercise AWS KMS operations retrieve_cmk(), create_cmk(),
      create_data_key(), and decrypt_data_key().
  
      Also, use the various KMS keys to encrypt and decrypt a file.
      """
  
      # Specify the description for the CMK. If an existing CMK with this
      # description is found, it is used for subsequent operations.
      # Otherwise, a new CMK is created.
      cmk_description = 'My sample CMK'
  
      # Specify a filename to encrypt/decrypt. To skip these operations,
      # specify an empty string.
      file_to_encrypt = 'sample.png'
  
      # Set up logging
      logging.basicConfig(level=logging.DEBUG,
                          format='%(levelname)s: %(asctime)s: %(message)s')
  
      # Does the desired CMK already exist?
      cmk_id, cmk_arn = retrieve_cmk(cmk_description)
      if cmk_id is None:
          # No, create it
          cmk_id, cmk_arn = create_cmk(cmk_description)
          if cmk_id is None:
              exit(1)
          logging.info('Created new AWS KMS CMK')
      else:
          logging.info('Retrieved existing AWS KMS CMK')
  
      # Use the key to encrypt and decrypt a file
      if file_to_encrypt:
          # Encrypted file contents are written to <file_to_encrypt>.encrypted
          # An encrypted data key is also written to the output file.
          # The encrypted file can be decrypted at any time and by any program
          # that has the credentials to decrypt the data key.
          if encrypt_file(file_to_encrypt, cmk_arn):
              logging.info(f'{file_to_encrypt} encrypted to '
                           f'{file_to_encrypt}.encrypted')
  
              # Decrypt the file
              if decrypt_file(file_to_encrypt):
                  # Decrypted file contents are written to <file_to_encrypt>.decrypted
                  logging.info(f'{file_to_encrypt}.encrypted decrypted to '
                               f'{file_to_encrypt}.decrypted')
  
  
  if __name__ == '__main__':
      main()
  