import os
import sys
import base64
import zipfile
import io
from pqcrypto.kem.ml_kem_1024 import generate_keypair, encrypt, decrypt
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


# =========================
# ZIP DOSSIER EN MEMOIRE
# =========================
def zip_folder_to_bytes(folder_path):
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                full_path = os.path.join(root, file)
                relative_path = os.path.relpath(full_path, folder_path)
                zipf.write(full_path, relative_path)
    return buffer.getvalue()


# =========================
# ENCRYPT FOLDER
# =========================
def encrypt_folder(folder_name):
    if not os.path.isdir(folder_name):
        print("Dossier introuvable.")
        return

    print("Compression du dossier...")
    zipped_data = zip_folder_to_bytes(folder_name)

    # Génération nouvelle paire Kyber
    public_key, private_key = generate_keypair()

    print("\n===== NEW KYBER 1024 KEYPAIR =====\n")

    pub_b64 = base64.b64encode(public_key).decode()
    priv_b64 = base64.b64encode(private_key).decode()

    print("PUBLIC KEY (base64):\n")
    print(pub_b64)
    print("\nPRIVATE KEY (base64):\n")
    print(priv_b64)
    print("\n⚠️ GARDE LA CLE PRIVEE POUR DECHIFFRER ⚠️\n")

    # Encapsulation
    ciphertext_kem, shared_secret = encrypt(public_key)

    aes_key = shared_secret[:32]
    aesgcm = AESGCM(aes_key)
    nonce = os.urandom(12)

    encrypted_data = aesgcm.encrypt(nonce, zipped_data, None)

    output_file = folder_name + ".bin"

    with open(output_file, "wb") as f:
        f.write(len(ciphertext_kem).to_bytes(4, "big"))
        f.write(ciphertext_kem)
        f.write(nonce)
        f.write(encrypted_data)

    print("Fichier généré :", output_file)


# =========================
# DECRYPT
# =========================
def decrypt_folder(input_bin, private_key_b64):
    private_key = base64.b64decode(private_key_b64)

    with open(input_bin, "rb") as f:
        kem_len = int.from_bytes(f.read(4), "big")
        ciphertext_kem = f.read(kem_len)
        nonce = f.read(12)
        encrypted_data = f.read()

    shared_secret = decrypt(ciphertext_kem, private_key)

    aes_key = shared_secret[:32]
    aesgcm = AESGCM(aes_key)

    decrypted_zip = aesgcm.decrypt(nonce, encrypted_data, None)

    folder_name = input_bin.replace(".bin", "_restored")

    print("Extraction du dossier...")

    with zipfile.ZipFile(io.BytesIO(decrypted_zip), "r") as zipf:
        zipf.extractall(folder_name)

    print("Dossier restauré :", folder_name)


# =========================
# MAIN
# =========================
if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("""
Usage:
  python kyber_encrypt_folder.py encrypt USEFULL
  python kyber_encrypt_folder.py decrypt USEFULL.bin <private_key_base64>
""")
        sys.exit(1)

    mode = sys.argv[1]

    if mode == "encrypt":
        encrypt_folder(sys.argv[2])

    elif mode == "decrypt":
        decrypt_folder(sys.argv[2], sys.argv[3])

    else:
        print("Mode inconnu")