import time
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Random import get_random_bytes
import os

# RSA key generation
def generate_keys():
    key = RSA.generate(2048)
    private_key = key.export_key()
    public_key = key.publickey().export_key()
    return private_key, public_key

# Function for encrypting a file
def encrypt_file(input_file, output_file, public_key):
    try:
        # Loading the public key
        rsa_public_key = RSA.import_key(public_key)
        cipher = PKCS1_OAEP.new(rsa_public_key)
        
        with open(input_file, 'rb') as f:
            file_data = f.read()

        # Encrypting the data
        encrypted_data = cipher.encrypt(file_data)

        with open(output_file, 'wb') as f:
            f.write(encrypted_data)

        print(f'Encrypted file successfully created: {output_file}')

    except FileNotFoundError:
        print(f'Error: File {input_file} was not found.')
    except Exception as e:
        print(f'An error occurred during encryption: {str(e)}')

# Function for decrypting a file
def decrypt_file(input_file, output_file, private_key):
    try:
        # Loading the private key
        rsa_private_key = RSA.import_key(private_key)
        cipher = PKCS1_OAEP.new(rsa_private_key)

        with open(input_file, 'rb') as f:
            encrypted_data = f.read()

        # Decrypting the data
        decrypted_data = cipher.decrypt(encrypted_data)

        with open(output_file, 'wb') as f:
            f.write(decrypted_data)

        print(f'Decrypted file successfully created: {output_file}')

    except FileNotFoundError:
        print(f'Error: File {input_file} was not found.')
    except Exception as e:
        print(f'An error occurred during decryption: {str(e)}')

# Function to measure encryption and decryption times
def measure_encryption_decryption_time(input_file):
    private_key, public_key = generate_keys()

    start_time = time.time()
    encrypt_file(input_file, 'encrypted_file.bin', public_key)
    encryption_time = time.time() - start_time

    start_time = time.time()
    decrypt_file('encrypted_file.bin', 'decrypted_file.txt', private_key)
    decryption_time = time.time() - start_time

    return encryption_time, decryption_time

# Main
if __name__ == "__main__":

    input_file = input('Enter the name of the file you want to encrypt: ')

    if not os.path.isfile(input_file):
        print(f'The file {input_file} does not exist. Please ensure the name and path are correct.')
    else:
        encryption_time, decryption_time = measure_encryption_decryption_time(input_file)

        print(f'Encryption time: {encryption_time:.6f} seconds')
        print(f'Decryption time: {decryption_time:.6f} seconds')

        input('\nProcess completed. Press Enter to exit...')