from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
import argparse
import os


def setup(key_file):
    key = ChaCha20Poly1305.generate_key()
    with open(key_file, "wb") as f:
        f.write(key)


def encrypt(input_file, key_file):
    with open(key_file, "rb") as f:
        key = f.read(32)

    with open(input_file, "rb") as f:
        plaintext = f.read()

    aad = b"authenticated but unencrypted data"
    nonce = os.urandom(12)

    cipher = ChaCha20Poly1305(key)
    ciphertext = cipher.encrypt(nonce, plaintext, aad)

    with open(f"{input_file}.enc", "wb") as f:
        f.write(nonce + ciphertext)



def decrypt(input_file, key_file):
    with open(key_file, "rb") as f:
        key = f.read(32)
    
    with open(input_file, "rb") as f:
        input_bytes = f.read()

    aad = b"authenticated but unencrypted data"
    nonce = input_bytes[:12]
    ciphertext = input_bytes[12:]


    cipher = ChaCha20Poly1305(key)
    try:
        plaintext = cipher.decrypt(nonce, ciphertext, aad)
    except Exception as e:
        print(f"Could not validate the authentication: {e}")
        return
    
    

    with open(f"{input_file}.dec", "wb") as f:
        f.write(plaintext)


def main():
    parser = argparse.ArgumentParser(
        description="Program to perform operations using Authenticated ChaCha20 cipher on files",
    )

    subparsers = parser.add_subparsers(dest="operation", help="Operation to perform")

    # Setup subcommand
    setup_parser = subparsers.add_parser("setup", help="Setup a key file")
    setup_parser.add_argument("fkey", help="File to contain the appropriate key for the ChaCha20 cipher")

    # Encrypt subcommand
    enc_parser = subparsers.add_parser("enc", help="Encrypt a file")
    enc_parser.add_argument("fich", help="File to be encrypted")
    enc_parser.add_argument("fkey", help="File containing the key for the ChaCha20 cipher")

    # Decrypt subcommand
    dec_parser = subparsers.add_parser("dec", help="Decrypt a file")
    dec_parser.add_argument("fich", help="File to be decrypted")
    dec_parser.add_argument("fkey", help="File containing the key for the ChaCha20 cipher")

    args = parser.parse_args()
    match args.operation:
        case "setup":
            key_file = args.fkey
            setup(key_file)
        case "enc":
            input_file = args.fich
            key_file = args.fkey
            encrypt(input_file,key_file)
        case "dec":
            input_file = args.fich
            key_file = args.fkey
            decrypt(input_file,key_file)
        case "help":
            parser.print_help()

if __name__ == "__main__":
    main()