Python - AES 256 crypt/decrypt compatible with 'openssl enc' format

aes256-compat-openssl-enc-dec.py

#!python
"""OpenSSL enc compatible script by Python.
"""
""" Based on:
    https://gist.github.com/chrono-meter/d122cbefc6f6248a0af554995f072460
    Added: base64
    Need: python3-crypto
    
    # Encrypt
    python3 ./aes256-compat-openssl-enc-dec.py -in <filein> -e -salt -aes-256-cbc

    # Decrypt
    python3 ./aes256-compat-openssl-enc-dec.py -in <filein> -d -salt -aes-256-cbc
"""
import sys
import os
import io
import hashlib
import binascii
import argparse
import logging
import getpass
import base64
from Crypto.Cipher import (AES, ARC2, ARC4, Blowfish, CAST, DES, DES3)
from Crypto.Cipher.blockalgo import (MODE_ECB, MODE_CBC, MODE_CFB, MODE_OFB, MODE_CTR, MODE_OPENPGP)

OPENSSL_ENC_MAGIC = b'Salted__'
PKCS5_SALT_LEN = 8

def EVP_BytesToKey(key_length: int, iv_length: int, md, salt: bytes, data: bytes, count: int=1) -> (bytes, bytes):
    """
    Usage:
        key, iv = EVP_BytesToKey(
            32,  # 256 bits
            Crypto.Cipher.AES.block_size,
            hashlib.sha256,
            salt,
            password.encode('utf-8'),
            )

    See:
        https://github.com/openssl/openssl/blob/6f0ac0e2f27d9240516edb9a23b7863e7ad02898/crypto/evp/evp_key.c#L74
    """
    assert data
    assert salt == b'' or len(salt) == PKCS5_SALT_LEN

    md_buf = b''
    key = b''
    iv = b''
    # key_length = 32  # type.key_size
    # iv_length = type.block_size

    addmd = 0

    while key_length > len(key) or iv_length > len(iv):
        c = md()
        if addmd:
            c.update(md_buf)
        addmd += 1
        c.update(data)
        c.update(salt)
        md_buf = c.digest()
        for i in range(1, count):
            md_buf = md(md_buf)

        md_buf2 = md_buf

        if key_length > len(key):
            key, md_buf2 = key + md_buf2[:key_length - len(key)], md_buf2[key_length - len(key):]

        if iv_length > len(iv):
            iv = iv + md_buf2[:iv_length - len(iv)]

    return key, iv


def openssl_enc_encrypt(input: io.BytesIO, output: io.BytesIO, passphrase: bytes, cipher_factory, key_length, block_size, md=hashlib.sha256, salt: bytes=None):
    if salt is None:
        # TODO: Python 3.6's secrets module
        salt = os.urandom(PKCS5_SALT_LEN)

    key, iv = EVP_BytesToKey(key_length, block_size, md, salt, passphrase)

    cipher = cipher_factory(key, iv)

    encrypted = b""
    encrypted += OPENSSL_ENC_MAGIC
    encrypted += salt

    padding = b''
    while 1:
        chunk = input.read(cipher.block_size)

        # PKCS#7 padding method
        if len(chunk) < cipher.block_size:
            remaining = cipher.block_size - len(chunk)
            padding = bytes([remaining] * remaining)

        encrypted += cipher.encrypt(chunk + padding)

        if padding:
            break

    output.write(base64.b64encode(encrypted))


def openssl_enc_decrypt(input: io.BytesIO, output: io.BytesIO, passphrase: bytes, cipher_factory, key_length, block_size, md=hashlib.sha256):
    assert input.read(len(OPENSSL_ENC_MAGIC)) == OPENSSL_ENC_MAGIC, 'bad magic number'
    salt = input.read(PKCS5_SALT_LEN)

    key, iv = EVP_BytesToKey(key_length, block_size, md, salt, passphrase)

    cipher = cipher_factory(key, iv)

    prefetch = chunk = None
    while 1:
        chunk = prefetch
        prefetch = input.read(cipher.block_size)

        if chunk:
            chunk = cipher.decrypt(chunk)
            if not prefetch:
                chunk = chunk[:-chunk[-1]]
            output.write(chunk)

        if not prefetch:
            break


if __name__ == '__main__':
    # https://github.com/openssl/openssl/blob/master/apps/enc.c
    # https://wiki.openssl.org/index.php/Command_Line_Utilities#Basic_file
    parser = argparse.ArgumentParser()

    parser.add_argument('-in', dest='input', nargs='?', type=argparse.FileType('rb'), default=sys.stdin.buffer, help='Input file')
    parser.add_argument('-out', dest='output', nargs='?', type=argparse.FileType('wb'), default=sys.stdout.buffer, help='Output file')
    # {"pass", OPT_PASS, 's', "Passphrase source"},

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-e', dest='encrypt', action='store_true', default=True, help='Encrypt')
    group.add_argument('-d', dest='encrypt', action='store_false', help='Decrypt')

    # group = parser.add_mutually_exclusive_group()
    # group.add_argument('-p', dest='print', action='store_true', help='Print the iv/key')
    # group.add_argument('-P', dest='print_exit', action='store_true', help='Print the iv/key and exit')

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-v', dest='loglevel', action='count', help='Verbose output')
    group.add_argument('-debug', dest='loglevel', action='store_const', const=4, help='Print debug info')

    # {"nopad", OPT_NOPAD, '-', "Disable standard block padding"},

    # parser.add_argument('-a', '-base64', dest='base64', action='store_true', default=True, help='Base64 encode/decode, depending on encryption flag')
    # {"A", OPT_UPPER_A, '-',
     # "Used with -[base64|a] to specify base64 buffer as a single line"},

    # {"bufsize", OPT_BUFSIZE, 's', "Buffer size"},

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-k', dest='passphrase', help='Passphrase')
    group.add_argument('-K', dest='passphrase', type=binascii.a2b_hex, help='Raw key, in hex')

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-salt', dest='dummy', action='store_true', help='Use salt in the KDF (default)')
    group.add_argument('-nosalt', dest='salt', action='store_const', const=b'', help='Do not use salt in the KDF')
    parser.add_argument('-S', dest='salt', type=binascii.a2b_hex, help='Salt, in hex')

    # parser.add_argument('-iv', dest='iv', type=binascii.a2b_hex, help='IV in hex')

    parser.add_argument('-md', dest='md', type=lambda v: functools.partial(hashlib.new, v), default=hashlib.sha256, help='Use specified digest to create a key from the passphrase')

    group = parser.add_mutually_exclusive_group(required=True)

    for key_size in AES.key_size:
        prefix = '-aes-%d' % (key_size * 8, )
        group.add_argument(prefix + '-cbc', dest='cipher', action='store_const', const=(lambda key, iv: AES.new(key, MODE_CBC, iv), key_size, AES.block_size))
        # -aes-128-ccm
        group.add_argument(prefix + '-cfb', dest='cipher', action='store_const', const=(lambda key, iv: AES.new(key, MODE_CFB, iv), key_size, AES.block_size))
        # -aes-128-cfb1
        # -aes-128-cfb8
        group.add_argument(prefix + '-ctr', dest='cipher', action='store_const', const=(lambda key, iv: AES.new(key, MODE_CTR, iv), key_size, AES.block_size))
        group.add_argument(prefix + '-ecb', dest='cipher', action='store_const', const=(lambda key, iv: AES.new(key, MODE_ECB, iv), key_size, AES.block_size))
        # -aes-128-gcm
        # -aes-128-ocb
        group.add_argument(prefix + '-ofb', dest='cipher', action='store_const', const=(lambda key, iv: AES.new(key, MODE_OFB, iv), key_size, AES.block_size))
        # -aes-128-xts
        group.add_argument('-aes%d' % (key_size * 8, ), dest='cipher', action='store_const', const=(lambda key, iv: AES.new(key, MODE_CBC, iv), key_size, AES.block_size))
        # -aes128-wrap

    # -bf
    # -bf-cbc
    # -bf-cfb
    # -bf-ecb
    # -bf-ofb
    # -blowfish

    # -camellia-128-cbc
    # -camellia-128-cfb
    # -camellia-128-cfb1
    # -camellia-128-cfb8
    # -camellia-128-ctr
    # -camellia-128-ecb
    # -camellia-128-ofb
    # -camellia-192-cbc
    # -camellia-192-cfb
    # -camellia-192-cfb1
    # -camellia-192-cfb8
    # -camellia-192-ctr
    # -camellia-192-ecb
    # -camellia-192-ofb
    # -camellia-256-cbc
    # -camellia-256-cfb
    # -camellia-256-cfb1
    # -camellia-256-cfb8
    # -camellia-256-ctr
    # -camellia-256-ecb
    # -camellia-256-ofb
    # -camellia128
    # -camellia192
    # -camellia256

    # -cast
    # -cast-cbc
    # -cast5-cbc
    # -cast5-cfb
    # -cast5-ecb
    # -cast5-ofb

    # -chacha20
    # -chacha20-poly1305

    # -des
    # -des-cbc
    # -des-cfb
    # -des-cfb1
    # -des-cfb8
    # -des-ecb
    # -des-ede
    # -des-ede-cbc
    # -des-ede-cfb
    # -des-ede-ecb
    # -des-ede-ofb
    # -des-ede3
    # -des-ede3-cbc
    # -des-ede3-cfb
    # -des-ede3-cfb1
    # -des-ede3-cfb8
    # -des-ede3-ecb
    # -des-ede3-ofb
    # -des-ofb

    # -des3
    # -des3-wrap
    # -desx
    # -desx-cbc

    # -id-aes128-CCM
    # -id-aes128-GCM
    # -id-aes128-wrap
    # -id-aes128-wrap-pad
    # -id-aes192-CCM
    # -id-aes192-GCM
    # -id-aes192-wrap
    # -id-aes192-wrap-pad
    # -id-aes256-CCM
    # -id-aes256-GCM
    # -id-aes256-wrap
    # -id-aes256-wrap-pad
    # -id-smime-alg-CMS3DESwrap

    # -idea
    # -idea-cbc
    # -idea-cfb
    # -idea-ecb
    # -idea-ofb

    # -rc2
    # -rc2-128
    # -rc2-40
    # -rc2-40-cbc
    # -rc2-64
    # -rc2-64-cbc
    # -rc2-cbc
    # -rc2-cfb
    # -rc2-ecb
    # -rc2-ofb

    # -rc4
    # -rc4-40
    # -rc4-hmac-md5

    # -seed
    # -seed-cbc
    # -seed-cfb
    # -seed-ecb
    # -seed-ofb

    # {"none", OPT_NONE, '-', "Don't encrypt"},
    # {"z", OPT_Z, '-', "Use zlib as the 'encryption'"},

    arg = parser.parse_args()
    arg.loglevel = max([logging.CRITICAL, logging.ERROR, logging.WARNING, logging.INFO, logging.DEBUG, logging.NOTSET][arg.loglevel:] + [logging.NOTSET])

    if arg.passphrase is None:
        arg.passphrase = getpass.getpass('password: ')
    if isinstance(arg.passphrase, str):
        arg.passphrase = arg.passphrase.encode('utf-8')

    if arg.encrypt:
        openssl_enc_encrypt(arg.input, arg.output, arg.passphrase, arg.cipher[0], arg.cipher[1], arg.cipher[2], arg.md, arg.salt)
    else:
        openssl_enc_decrypt(arg.input, arg.output, arg.passphrase, arg.cipher[0], arg.cipher[1], arg.cipher[2], arg.md)