Python functions to generate asymmetric keys and encrypt/decrypt data

encrypt_decrypt.py

# pip install pycrypto

import base64

from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA

PUBLIC_KEY = '/home/public_key.pem'
PRIVATE_KEY = '/home/private_key.pem'


def encrypt(message: str, key_file: str = None) -> bytes:
    if not key_file:
        k_file = PUBLIC_KEY
    else:
        k_file = key_file
    with open(k_file, 'rb') as fp:
        key = RSA.importKey(fp.read())
    cyp = PKCS1_OAEP.new(key)
    encrypted = cyp.encrypt(message.encode())
    encoded = base64.b64encode(encrypted)
    return encoded


def decrypt(encrypted: bytes, key_file: str = None) -> str:
    if not key_file:
        k_file = PRIVATE_KEY
    else:
        k_file = key_file
    with open(k_file, 'rb') as fp:
        key = RSA.importKey(fp.read())
    decyp = PKCS1_OAEP.new(key)
    decoded = base64.b64decode(encrypted)
    decrypted = decyp.decrypt(decoded)
    return decrypted.decode()


def generate_keys(n_bits: int = 2056, filepath: str = None) -> tuple:
    # n_bits must be multiple of 4
    private = RSA.generate(n_bits)
    public = private.publickey()
    if filepath is not None:
        # save keys in files
        with open(filepath + '_public.pem', 'w') as fp:
            fp.write(public.exportKey().decode())
        with open(filepath + '_private.pem', 'w') as fp:
            fp.write(private.exportKey().decode())
    return private, public