#!/usr/bin/python3
import argparse
import hashlib
import os
import sys
import io
import shutil
import tempfile
import base64
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

CHUNK_SIZE = 64 * 1024  # 64KB
FILE_TYPE_FLAG = b"\x00"
FOLDER_TYPE_FLAG = b"\x01"


def get_cipher_key(key_source):
    hasher = hashlib.sha256()
    if key_source.startswith("#"):
        hasher.update(key_source[1:].encode("utf-8"))
        return hasher.digest()

    if not os.path.exists(key_source):
        raise FileNotFoundError(f"Key file '{key_source}' missing.")

    with open(key_source, "rb") as f:
        while chunk := f.read(CHUNK_SIZE):
            hasher.update(chunk)
    return hasher.digest()


def encrypt_data(source, output, key, use_base64=False):
    temp_zip = None
    is_folder = False
    aesgcm = AESGCM(key)
    nonce = os.urandom(12)  # GCM standard nonce length

    try:
        if not source.startswith("@") and os.path.isdir(source):
            is_folder = True
            temp_dir = tempfile.mkdtemp()
            temp_zip = os.path.join(temp_dir, "archive.zip")
            shutil.make_archive(temp_zip.replace(".zip", ""), "zip", source)
            input_source = temp_zip
        else:
            input_source = source

        if isinstance(input_source, str) and input_source.startswith("@"):
            input_stream = io.BytesIO(input_source[1:].encode("utf-8"))
        else:
            input_stream = open(input_source, "rb")

        data = input_stream.read()
        input_stream.close()

        ciphertext = aesgcm.encrypt(nonce, data, None)
        payload = (
            nonce + (FOLDER_TYPE_FLAG if is_folder else FILE_TYPE_FLAG) + ciphertext
        )

        if use_base64:
            output.write(base64.b64encode(payload))
        else:
            output.write(payload)

    finally:
        if temp_zip and os.path.exists(os.path.dirname(temp_zip)):
            shutil.rmtree(os.path.dirname(temp_zip))


def decrypt_data(source, output, key, extract_path=None, use_base64=False):
    aesgcm = AESGCM(key)
    if source.startswith("@"):
        raw_payload = source[1:].encode("utf-8")
    else:
        with open(source, "rb") as f:
            raw_payload = f.read()

    if use_base64:
        raw_payload = base64.b64decode(raw_payload)

    nonce = raw_payload[:12]
    flag = raw_payload[12:13]
    ciphertext = raw_payload[13:]

    try:
        decrypted_data = aesgcm.decrypt(nonce, ciphertext, None)
    except Exception:
        raise ValueError("Decryption failed! Wrong key or data has been tampered with.")

    if flag == FOLDER_TYPE_FLAG and extract_path:
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_zip = os.path.join(tmp_dir, "decrypted.zip")
            with open(tmp_zip, "wb") as f:
                f.write(decrypted_data)
            shutil.unpack_archive(tmp_zip, extract_path, "zip")
    else:
        output.write(decrypted_data)


def main():
    parser = argparse.ArgumentParser(description="Safenc Tool")
    parser.add_argument(
        "-i", "--input", required=True, help="Path to file/folder or @text"
    )
    parser.add_argument(
        "-k", "--key", required=True, help="Path to key file or #passphrase"
    )
    parser.add_argument(
        "-o", "--output", help="Output file path or directory for extraction"
    )
    parser.add_argument("-d", "--decrypt", action="store_true")
    parser.add_argument("-b", "--base64", action="store_true")

    args = parser.parse_args()

    try:
        key = get_cipher_key(args.key)

        if args.decrypt:
            if args.output and (
                os.path.isdir(args.output) or not os.path.exists(args.output)
            ):
                out_stream = io.BytesIO()  # Temp buffer for logic
                decrypt_data(args.input, out_stream, key, args.output, args.base64)
            else:
                decrypt_data(
                    args.input, sys.stdout.buffer, key, args.output, args.base64
                )
        else:
            if args.output:
                with open(args.output, "wb") as f:
                    encrypt_data(args.input, f, key, args.base64)
            else:
                encrypt_data(args.input, sys.stdout.buffer, key, args.base64)

        if args.output:
            print("\nOperation successful.")

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
