crypt_file.py

import base64
from Crypto.Cipher import AES     # pip install pycrypto or pycryptodome
from Crypto.Hash import SHA256
from Crypto import Random


__all__ = ['encrypt', 'decrypt']


def encrypt(key, text, **kwargs):
    """Encrypt the given text so no one can read it.

    Args:
        key (str/bytes): Key to encrypt the text with.
        text (str/bytes): Text to encrypt.

    Returns:
        encrypted (str/bytes): Encrypted text.
    """
    is_bytes = True
    if isinstance(text, str):
        is_bytes = False
        text = text.encode('utf-8')
    if isinstance(key, str):
        key = key.encode('utf-8')
    key = SHA256.new(key).digest()  # use SHA-256 over our key to get a proper-sized AES key
    IV = Random.new().read(AES.block_size)  # generate IV
    encryptor = AES.new(key, AES.MODE_CBC, IV)
    padding = AES.block_size - len(text) % AES.block_size  # calculate needed padding
    text += bytes([padding]) * padding  # Python 2.x: source += chr(padding) * padding
    data = IV + encryptor.encrypt(text)  # store the IV at the beginning and encrypt
    data = base64.b64encode(data)
    if is_bytes:
        return data
    return data.decode("latin-1")


def decrypt(key, encrypted, **kwargs):
    """Decrypt the given text to make it human readable.

    Args:
        key (str/bytes): Key to encrypt the text with.
        encrypted (str/bytes): Encrypted text.

    Returns:
        text (str/bytes): Human readable text.
    """
    if isinstance(encrypted, str):
        is_bytes = False
        encrypted = base64.b64decode(encrypted.encode("latin-1"))
    else:
        is_bytes = True
        encrypted = base64.b64decode(encrypted)
    if isinstance(key, str):
        key = key.encode('utf-8')
    key = SHA256.new(key).digest()  # use SHA-256 over our key to get a proper-sized AES key
    IV = encrypted[:AES.block_size]  # extract the IV from the beginning
    decryptor = AES.new(key, AES.MODE_CBC, IV)
    data = decryptor.decrypt(encrypted[AES.block_size:])  # decrypt
    padding = data[-1]  # pick the padding value from the end; Python 2.x: ord(data[-1])
    if data[-padding:] != bytes([padding]) * padding:  # Python 2.x: chr(padding) * padding
        raise ValueError("Invalid padding...")
    data = data[:-padding]  # remove the padding

    if is_bytes:
        return data
    return data.decode('utf-8')


if __name__ == '__main__':
    import os
    import traceback
    import glob
    import argparse

    P = argparse.ArgumentParser(description='Encrypt/Decrypt a file or folder.')
    P.add_argument('path', type=str, help='File or folder to encrypt/decrypt.')
    P.add_argument('key', type=str, help='Key to use for the encryption')
    P.add_argument('--encrypt', '-e', action='store_const', const=True,
                   help='If given encrypt the given file or folder path else decrypt. By default decrypt is used.')
    P.add_argument('--decrypt', '-d', action='store_const', const=True, default=None,
                   help='If given decrypt the given file or folder path. By default decrypt is used.')

    ARGS = P.parse_args()
    PATH = ARGS.path
    KEY = ARGS.key
    ENCRYPT = ARGS.encrypt or ARGS.decrypt is False

    if os.path.isdir(PATH):
        PATH = glob.iglob(os.path.join(PATH, '*'))
    else:
        PATH = [PATH]

    CRYPT_IDENT = b'Crypto-'
    for FILE in PATH:
        with open(FILE, 'rb') as F:
            DATA = F.read()

        # Run Encryption
        try:
            if ENCRYPT:
                assert not DATA.startswith(CRYPT_IDENT), 'The file "{}" was already encrypted!'.format(FILE)
                DATA = CRYPT_IDENT + encrypt(KEY, DATA)
            else:
                assert DATA.startswith(CRYPT_IDENT), 'The file "{}" was not encrypted!'.format(FILE)
                DATA = DATA[len(CRYPT_IDENT):]
                DATA = decrypt(KEY, DATA)
        except AssertionError:
            traceback.print_exc()

        with open(FILE, 'wb') as F:
            F.write(DATA)