filecrypt.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import secrets
import binascii

# Generate keys using:
# $ openssl genrsa -out privatekey.pem 4096
# $ openssl rsa -in privatekey.pem -pubout -out publickey.pem


def main_encrypt(pubkey_path: str, filename: str) -> int:
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization, hashes
    from cryptography.hazmat.primitives.asymmetric import padding
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

    backend = default_backend()

    # Check destination
    dst = filename + ".enc"
    if os.path.exists(dst):
        print("{!r} already exists".format(dst), file=sys.stderr)
        return 1

    # Select preferred hash algorithm
    PaddingHash = hashes.SHA256
    # OpenSSL compatibility: PaddingHash = hashes.SHA1

    # Read the public key
    with open(pubkey_path, "rb") as fp:
        public_key = serialization.load_pem_public_key(
            fp.read(),
            backend=backend,
        )

    # Generate key and initialization vector (CTR nonce) for symmetric encryption
    aes_key = secrets.token_bytes(32)
    aes_iv = secrets.token_bytes(16)

    keys_info = b"%s\n%s\n" % (
        binascii.hexlify(aes_key),  # secret
        binascii.hexlify(aes_iv),  # actually not a secret, but why not
    )

    # And store it using asymmetric encryption
    keys_info_encrypted = public_key.encrypt(
        keys_info,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=PaddingHash()),
            algorithm=PaddingHash(),
            label=None,
        )
    )

    # Prepare AES-256-CTR encryption
    cipher = Cipher(algorithms.AES(aes_key), modes.CTR(nonce=aes_iv), backend=backend)
    encryptor = cipher.encryptor()

    with open(filename, "rb") as rfp:
        with open(dst, "wb") as wfp:
            # Write format identifier
            wfp.write("encrypted:RSA-{},OAEP-MGF1-{}:AES-256-CTR\n".format(
                public_key.key_size,
                PaddingHash.name.upper(),
            ).encode("ascii"))
            # write keys
            wfp.write(keys_info_encrypted)

            while True:
                # encrypt!
                chunk = rfp.read(65536)
                if not chunk:
                    break

                enc_chunk = encryptor.update(chunk)
                wfp.write(enc_chunk)

            enc_fin = encryptor.finalize()
            wfp.write(enc_fin)

    return 0


def main_decrypt(privkey_path: str, filename: str) -> int:
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization, hashes
    from cryptography.hazmat.primitives.asymmetric import padding
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

    backend = default_backend()

    # Check destination
    if filename.lower().endswith(".enc"):
        dst = filename[:-4]
    else:
        dst = filename + ".dec"
    if os.path.exists(dst):
        print("{!r} already exists".format(dst), file=sys.stderr)
        return 1

    # Read the private key
    with open(privkey_path, "rb") as fp:
        private_key = serialization.load_pem_private_key(
            fp.read(),
            password=None,  # TODO: implement?
            backend=backend,
        )

    with open(filename, "rb") as rfp:
        # Detect format
        header = rfp.readline().decode("ascii")
        if not header.startswith("encrypted:"):
            print("Unknown file format", file=sys.stderr)
            return 1
        _, asym_info, sym_info = header.rstrip().split(":")

        # Check asymmetric params support
        asym_algo, asym_padding = asym_info.split(",")
        if asym_algo != "RSA-{}".format(private_key.key_size):
            print("Cannot decode {}".format(asym_algo), file=sys.stderr)
            return 1
        if asym_padding == "OAEP-MGF1-SHA1":
            PaddingHash = hashes.SHA1
        elif asym_padding == "OAEP-MGF1-SHA256":
            PaddingHash = hashes.SHA256
        else:
            print("Unknown padding {}".format(asym_padding))
            return 1

        # Check symmetric params support
        if sym_info != "AES-256-CTR":
            print("Unknown symmetric algorithm {}".format(sym_info))
            return 1

        # Everything is okay, read encrypted key and iv
        keys_info_encrypted = rfp.read(private_key.key_size // 8)
        keys_info = private_key.decrypt(
            keys_info_encrypted,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=PaddingHash()),
                algorithm=PaddingHash(),
                label=None,
            )
        )

        aes_key, aes_iv = [binascii.unhexlify(x) for x in keys_info.strip().split(b"\n")]
        if len(aes_key) != 32:
            print("Unsupported AES key size {}".format(len(aes_key * 8)))
            return 1

        # Prepare AES-256-CTR decryption
        cipher = Cipher(algorithms.AES(aes_key), modes.CTR(nonce=aes_iv), backend=backend)
        decryptor = cipher.decryptor()

        with open(dst, "wb") as wfp:
            while True:
                # decrypt!
                enc_chunk = rfp.read(65536)
                if not enc_chunk:
                    break
                chunk = decryptor.update(enc_chunk)
                wfp.write(chunk)

            dec_fin = decryptor.finalize()
            wfp.write(dec_fin)

    return 0


def print_help() -> int:
    print((
        "Usage:\n"
        "{0} encrypt publickey.pem filename\n"
        "{0} decrypt privatekey.pem filename.enc\n"
    ).format(sys.argv[0]), file=sys.stderr)
    return 1


def main() -> int:
    if len(sys.argv) != 4:
        return print_help()

    cmd, key_path, filename = sys.argv[1:]

    if cmd in ("encrypt", "e"):
        return main_encrypt(key_path, filename)
    if cmd in ("decrypt", "d"):
        return main_decrypt(key_path, filename)
    return print_help()


if __name__ == "__main__":
    sys.exit(main())