#!/usr/bin/env python3

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms
import sys
import os
import argparse

def setup(key_file):
    f = open(key_file,"wb")
    key_bytes = os.urandom(32)
    f.write(key_bytes)
    f.close()

def encrypt(input_file, key_file):
    fkey = open(key_file,"rb")
    inp = open(input_file,"rb")
    out = open(f"{input_file}.enc","wb")

    key = fkey.read()
    plaintext = inp.read()
    nonce = os.urandom(16)
    cipher = Cipher(algorithms.ChaCha20(key=key,nonce=nonce),mode=None)
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(plaintext)
    ciphertext = nonce + ciphertext

    out.write(ciphertext)

    fkey.close()
    inp.close()
    out.close()

def decrypt(input_file,key_file):
    fkey = open(key_file,"rb")
    inp = open(f"{input_file}","rb")
    out = open(f"{input_file}.dec","wb")

    key = fkey.read()
    input_bytes = inp.read()
    nonce = input_bytes[:16]
    ciphertext = input_bytes[16:]

    cipher = Cipher(algorithms.ChaCha20(key=key,nonce=nonce),mode=None)
    decryptor = cipher.decryptor()
    plaintext = decryptor.update(ciphertext)

    out.write(plaintext)

    fkey.close()
    inp.close()
    out.close()

def main():
    parser = argparse.ArgumentParser(
        description="Program to perform operations using 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()