rsa_encryption.py

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


def save_file(path, bytes):
    answ = "y"
    if os.path.exists(path):
        answ = input("File %s exists, replace? (y/n) " % path)
    if "y" in answ.lower():
        with open(path, "wb") as f:
            f.write(bytes)


def read_file(path):
    with open(path, "rb") as f:
        return f.read()


def generate_data(msg, folder="data"):
    folder = "%s/" % folder
    if not os.path.exists(folder):
        os.makedirs(folder)

    key_pair = RSA.generate(4096)

    pub_key = key_pair.publickey()
    pub_key_pem = pub_key.exportKey()
    save_file(folder + "pub.key", pub_key_pem)

    key_pair_pem = key_pair.exportKey()
    save_file(folder + "priv.key", key_pair_pem)

    encryptor = PKCS1_OAEP.new(pub_key)
    encrypted = encryptor.encrypt(bytes(msg, "utf-8"))
    save_file(folder + "code", encrypted)

    print("Encoded data and keys were saved on: " + folder)


def read_data(folder="data"):
    priv_key_loaded = open("%s/priv.key" % folder).read()
    key_pair = RSA.import_key(priv_key_loaded)
    encrypted = read_file("%s/%s" % (folder, "code"))

    decryptor = PKCS1_OAEP.new(key_pair)
    decrypted = decryptor.decrypt(encrypted)
    return decrypted.decode("utf-8")


if __name__ == "__main__":
    generate_data("A very secret message")
    print(read_data())