from os import urandom as urandom
from os import getcwd as getcwd
from os import _exit as exit
from os import path as ospath
# pip3 install cryptography
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import padding as sym_padding
from argparse import ArgumentParser

#   SYNTAX/USE GUIDE:
#   python3 encfile.py --{encrypt/decrypt} --input {filename} --output {filename, but optional if encrypting. defaults to current working directory/cwd.}
#   example:
#   python3 encfile.py --encrypt --input secret.js
#   python3 encfile.py -e -i secret.js
def banner():
    print(" ____  _____ ____ ____  _____ _____  _____   _  ___  _____ ___  ____  _   _ ")
    print("/ ___|| ____/ ___|  _ \| ____|_   _|/ / / \ | |/ _ \|  ___/ _ \|  _ \| \ | |")
    print("\___ \|  _|| |   | |_) |  _|   | | / / /|  \| | | | | |_ | | | | |_) |  \| |")
    print(" ___) | |__| |___|  _ <| |___  | |/ / / | |\  | |_| |  _|| |_| |  _ <| |\  |")
    print("|____/|_____\____|_| \_\_____| |_/_/_/  |_| \_|\___/|_|   \___/|_| \_\_| \_|\n")


    print("'[S] The Secrets of the Fighter.' ")
    print("Do not transmit this program over unsecure channels.")
    print("This program allows you to encrypt files to send them to your fellow operatives.")
    print("It is classified as: SECRET//NOFORN.\n")

parser = ArgumentParser()
parser.add_argument("-e", "--encrypt", action="store_true", help="Encrypt a file", dest="encrypt")
parser.add_argument("-d", "--decrypt", action="store_true", help="Decrypt a file", dest="decrypt")
parser.add_argument("-i", "-input", dest="input")
parser.add_argument("-o", "--output", dest="output")
args = parser.parse_args()

if not any([args.encrypt, args.decrypt]):
    print("Specify a mode with -e or -d.")
    exit(1)
if not args.input:
        print("Specify an input file with -i.")
        exit(1)
else:
    banner()


def derive_key(password:bytes, salt:bytes):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA512(),
        iterations=900000,
        salt=salt,
        length=32,
    )
    return kdf.derive(password)

def encrypt_file(input_file:str, output_file:str, password:bytes):
    salt = urandom(32)
    iv = urandom(16)
    key = derive_key(password, salt)

    cipher = Cipher(
        algorithms.AES(key),
        modes.CFB(iv),
    )
    encryptor = cipher.encryptor()

    with open(input_file, "rb") as f_in, open(output_file, "wb") as f_out:
        f_out.write(salt)  # Write the salt to the header
        f_out.write(iv)  # Write the IV to the header

        for block in iter(lambda: f_in.read(256), b""):
            encrypted_block = encryptor.update(block)
            f_out.write(encrypted_block)

def decrypt_file(input_file, output_file, password):
    with open(input_file, "rb") as f_in, open(output_file, "wb") as f_out:
        salt = f_in.read(32)  # Read the salt from the header
        iv = f_in.read(16)  # Read the IV from the header

        key = derive_key(password, salt)

        cipher = Cipher(
            algorithms.AES(key),
            modes.CFB(iv),
        )
        decryptor = cipher.decryptor()

        for block in iter(lambda: f_in.read(256), b""):
            decrypted_block = decryptor.update(block)
            f_out.write(decrypted_block)

    print("Decryption completed.")

if args.encrypt:
    args.output = getcwd()
    args.output = args.output + "/encryptedfile.enc"
    password = bytes(input("PASSWORD:\n"), 'UTF-8')
    encrypt_file(args.input, args.output, password)
    print("Done. File stored as ", args.output)

if args.decrypt:
    password = bytes(input("PASSWORD:\n"), 'UTF-8')
    decrypt_file(args.input, args.output, password)
    print("Done. File stored as ", args.output)