RSACryptor/rsacryptor.py

# /usr/bin/env python3
import os
import datetime
import time
import sys
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP, AES
import os
import getpass



class Cryptor:
    def __init__(self):
        self.path = f"/home/{getpass.getuser()}/FileCryptor/"
        if os.path.exists(self.path):
            pass
        else:
            os.system("mkdir "+self.path)
    def encrypt(self, dataFile, publicKeyFile):
        '''
        use EAX mode to allow detection of unauthorized modifications
        '''
        # read data from file
        with open(dataFile, 'rb') as f:
            data = f.read()

        # convert data to bytes
        data = bytes(data)
        # read public key from file
        with open(publicKeyFile, 'rb') as f:
            publicKey = f.read()

        # create public key object
        key = RSA.import_key(publicKey)
        sessionKey = os.urandom(16)

        # encrypt the session key with the public key
        cipher = PKCS1_OAEP.new(key)
        encryptedSessionKey = cipher.encrypt(sessionKey)

        # encrypt the data with the session key
        cipher = AES.new(sessionKey, AES.MODE_EAX)
        ciphertext, tag = cipher.encrypt_and_digest(data)
        []

        # save the encrypted data to file
        [fileName, fileExtension] = dataFile.split('.')
        encryptedFile = fileName + '.' + fileExtension
        with open(encryptedFile, 'wb+') as f:
            [f.write(x) for x in (encryptedSessionKey, cipher.nonce, tag, ciphertext)]
        print('Encrypted file saved to ' + encryptedFile)


    def decrypt(self, dataFile, privateKeyFile):
        '''
        use EAX mode to allow detection of unauthorized modifications
        '''

        # read private key from file
        with open(privateKeyFile, 'rb') as f:
            privateKey = f.read()
            # create private key object
            key = RSA.import_key(privateKey)

        # read data from file
        with open(dataFile, 'rb') as f:
            # convert data to bytes

            # read the session key
            encryptedSessionKey, nonce, tag, ciphertext = [f.read(x) for x in (key.size_in_bytes(), 16, 16, -1)]

        # decrypt the session key
        cipher = PKCS1_OAEP.new(key)
        sessionKey = cipher.decrypt(encryptedSessionKey)

        # decrypt the data with the session key
        cipher = AES.new(sessionKey, AES.MODE_EAX, nonce)
        data = cipher.decrypt_and_verify(ciphertext, tag)

        # save the decrypted data to file
        [fileName, fileExtension] = dataFile.split('.')
        decryptedFile = fileName + '.' + fileExtension
        with open(decryptedFile, 'wb') as f:
            f.write(data)

        print('Decrypted file saved to ' + decryptedFile)

    def genKey(self,nameFile):
        key = RSA.generate(2048)
        privateKey = key.export_key()
        publicKey = key.publickey().export_key()
        with open(self.path+"private_" + nameFile + ".pem", 'wb+') as f:
            f.write(privateKey)

        # save public key to file
        with open(self.path+"public_" + nameFile + ".pem", 'wb+') as f:
            f.write(publicKey)
        print("Keys saved in "+self.path)
        return [self.path+"private_" + nameFile + ".pem",self.path+"public_" + nameFile + ".pem"]




CRYPTOR = Cryptor()

def clear_screen():
    os.system('clear')

def commands():
    print("1:Encrypt file")
    print("2:Decrypt file")
    print("3:Genereate standalone pair keys (pub/private)")
    print("4:About")
    print("0:Quit")

def logo():     
    print("╭━━━┳━━━┳━━━┳━━━┳╮╱╱╭┳━━━┳━━━━┳━━━┳━━━╮")
    print("┃╭━╮┃╭━━┫╭━╮┃╭━╮┃╰╮╭╯┃╭━╮┃╭╮╭╮┃╭━╮┃╭━╮┃")
    print("┃┃╱┃┃╰━━┫┃╱╰┫╰━╯┣╮╰╯╭┫╰━╯┣╯┃┃╰┫┃╱┃┃╰━╯┃")
    print("┃╰━╯┃╭━━┫┃╱╭┫╭╮╭╯╰╮╭╯┃╭━━╯╱┃┃╱┃┃╱┃┃╭╮╭╯")
    print("┃╭━╮┃╰━━┫╰━╯┃┃┃╰╮╱┃┃╱┃┃╱╱╱╱┃┃╱┃╰━╯┃┃┃╰╮")
    print("╰╯╱╰┻━━━┻━━━┻╯╰━╯╱╰╯╱╰╯╱╱╱╱╰╯╱╰━━━┻╯╰━╯")

def encrypt():
    clear_screen()
    file = input("Enter full path to file:")
    print("Select:\n 1:Use exist public key \n 2:Generate new keys")
    q = int(input("Select:"))
    if q==1:
        clear_screen()
        path_to_key = input("Enter full path to public key:")
        clear_screen()
        CRYPTOR.encrypt(file,path_to_key)
        time.sleep(5)
    else:
        keys = CRYPTOR.genKey(str(datetime.date.today()))
        clear_screen()
        CRYPTOR.encrypt(file,keys[1])
        time.sleep(5)

def decrypt():
    clear_screen()
    file = input("Enter full path to file:")
    key = input("Enter full path to private key:")
    CRYPTOR.decrypt(file,key)
    time.sleep(5)

def genkey():
    clear_screen()
    CRYPTOR.genKey(str(datetime.date.today()))
    print("Keys generated in your home folder.")
    time.sleep(5)

def about():
    clear_screen()
    logo()
    print("This program created for encryption/decryption files with AES.\nIMPORTANT: if you will try encrypt photos (.png and etc.) you can see preview of this file.")
    print("Created by Ebobalik. URL:https://notabug.org/Ebobalik")
    input()

def main_loop():
    clear_screen()
    while True:
        logo()
        commands()
        query = int(input("Select:"))
        match query:
            case 0:
                break
            case 1:
                encrypt()
            case 2:
                decrypt()
            case 3:
                genkey()
            case 4:
                about()
        clear_screen()
    
if __name__=="__main__":
    params = sys.argv
    params.pop(0)
    if len(params) == 1:
        if params[0]=="-G":
            main_loop() 
        if params[0]=="--help" or params[0]=="-H":
            logo()
            print("Keys for usage: \n-E: encrypt file. \n-D: decrypt file. \n-H or --help: show this message. \n-k: path to private or public keys. For encryption file you should use public key. For decryption - private key.")
    elif len(params)>1:
        if params[0]=="-D":
            if len(params)==4:
                CRYPTOR.decrypt(params[1],params[3])
            else:
                print("Incorrect command. Run with key --help or -H")
        if params[0]=="-E":
            if len(params)==4:
                CRYPTOR.decrypt(params[1],params[3])
            elif len(params)==2:
                keys = CRYPTOR.genKey(str(datetime.date.today()))
                CRYPTOR.encrypt(params[1],keys[1])
            else:
                print("Incorrect command. Run with key --help or -H")
        if params[0]=="-GSK":
            CRYPTOR.genKey("StandaloneKeys")