import argparse
import os
import base64
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes, padding


def genSecretKey(key, salt):
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        info=None
    )
    secret_key = hkdf.derive(key)
    return secret_key


def encrypt_file(file_path, encrypted_file_path=None):
    file_name = os.path.basename(file_path)
    key_Mes = "fb8MH7yIuVUeCp5W4S9cKFtsHaxXIP/zvkO36wNNcO4="
    key = base64.b64decode(key_Mes.encode("ascii"))

    with open(file_path, 'rb') as f:
        data = f.read()

        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(data)
        padded_data += padder.finalize()

        cipher = Cipher(
            algorithms.AES(key),
            modes.ECB()
        )

        encryptor = cipher.encryptor()
        ct = encryptor.update(padded_data) + encryptor.finalize()
        if encrypted_file_path is None:
            encrypted_file_path = f"{file_name}.enc"
        with open(encrypted_file_path, 'wb') as encrypted_file:
            encrypted_file.write(ct)


def main():
    # Parse Arguments
    parser = argparse.ArgumentParser(description='Encrypt file')
    parser.add_argument('-p', '--path', help='Path to file', required=True)
    parser.add_argument('-o', '--out_file', help='Write to file', default=None)
    args = parser.parse_args()

    return args


if __name__ == '__main__':
    args = main()
    file_path = args.path
    if args.out_file is not None:
        encrypted_file_path = args.out_file
        encrypt_file(file_path, encrypted_file_path)
    else:
        encrypt_file(file_path)