meduza.py - semisymmetric file encryption with pycryptodome

#!/usr/bin/python

RSA_KEY_SIZE = 2048


import sys
from getpass import getpass

from Crypto.Hash import SHA256
from Crypto.Cipher import ChaCha20, PKCS1_OAEP
from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes

def make_prng(seed):
    key = SHA256.SHA256Hash(seed).digest()
    nonce = SHA256.SHA256Hash(key).digest()[:12]
    cc = ChaCha20.new(key=key, nonce=nonce)
    def prng(n):
        return cc.encrypt(b'0'*n)
    return prng

def emit_pub_key(seed, size=2048):
    prng = make_prng(seed)
    key = RSA.generate(size, randfunc=prng)
    return key.public_key().export_key('PEM')

def encrypt(file, pub_key):
    pub_key = RSA.import_key(pub_key)
    rsa = PKCS1_OAEP.new(pub_key)

    sess_key = get_random_bytes(32)
    cc = ChaCha20.new(key=sess_key)
    sess_key += cc.nonce

    enc_key = rsa.encrypt(sess_key)
    sys.stdout.buffer.write(len(enc_key).to_bytes(4, 'little'))
    sys.stdout.buffer.write(enc_key)

    f = open(file, 'rb')
    while True:
        b = f.read(1024**2)
        if not b: break
        e = cc.encrypt(b)
        sys.stdout.buffer.write(e)

def decrypt(file, seed, size=2048):
    f = open(file, 'rb')
    key_size = int.from_bytes(f.read(4), 'little')
    sess_key = f.read(key_size)

    prng = make_prng(seed)
    rsa_key = RSA.generate(size, randfunc=prng)
    rsa = PKCS1_OAEP.new(rsa_key)
    sess_key = rsa.decrypt(sess_key)

    sess_key, nonce = sess_key[:32], sess_key[32:]
    cc = ChaCha20.new(key=sess_key, nonce=nonce)

    while True:
        b = f.read(1024**2)
        if not b: break
        d = cc.decrypt(b)
        sys.stdout.buffer.write(d)

def get_seed(prompt):
    pswd = getpass(prompt, stream=sys.stderr)
    seed = pswd.encode()
    return seed



if len(sys.argv) < 2:
    print(f'''Usage:
    {sys.argv[0]} <file> <public_key.pem>   - to encrypt a file;
    {sys.argv[0]} <file>                    - to decrypt a file;
    {sys.argv[0]} -k                        - to generate a public key.

    Everything goes to stdout.

    Examples:
        {sys.argv[0]} -k > key.pem
        {sys.argv[0]} secret.txt key.pem > secret.txt.enc

        {sys.argv[0]} secret.txt.enc > secret_dec.txt
    ''')
    exit(0)


if sys.argv[1] == '-k':
    seed = get_seed('Enter new passphrase: ')
    print(emit_pub_key(seed, RSA_KEY_SIZE).decode())
    exit(0)

if len(sys.argv) == 2:
    seed = get_seed('Enter passphrase to decrypt: ')
    decrypt(sys.argv[1], seed, RSA_KEY_SIZE)
    exit(0)

if len(sys.argv) == 3:
    pub_key = open(sys.argv[2], 'rb').read()
    encrypt(sys.argv[1], pub_key)
    exit(0)


print('Wrong argument count.\nRun program without args for syntax help.', file=sys.stderr)
exit(1)