import argparse
import io
from pathlib import Path
import zipfile

from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

def decrypt_profile(source: Path, destination: Path, key: bytes) -> None:
    ciphertext = source.read_bytes()
    decryptor = Cipher(algorithms.AES(key), modes.ECB()).decryptor()
    padded_zip = decryptor.update(ciphertext) + decryptor.finalize()

    unpadder = padding.PKCS7(128).unpadder()
    zip_data = unpadder.update(padded_zip) + unpadder.finalize()

    with zipfile.ZipFile(io.BytesIO(zip_data), "r") as archive:
        destination.write_bytes(archive.read("profile"))

def encrypt_profile(source: Path, destination: Path, key: bytes) -> None:
    json_data = source.read_bytes()
    zip_buffer = io.BytesIO()

    with zipfile.ZipFile(
        zip_buffer,
        "w",
        compression=zipfile.ZIP_DEFLATED,
        compresslevel=6,
    ) as archive:
        archive.writestr("profile", json_data)

    padder = padding.PKCS7(128).padder()
    padded_zip = padder.update(zip_buffer.getvalue()) + padder.finalize()

    encryptor = Cipher(algorithms.AES(key), modes.ECB()).encryptor()
    ciphertext = encryptor.update(padded_zip) + encryptor.finalize()

    destination.write_bytes(ciphertext)

def create_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)

    operation = parser.add_mutually_exclusive_group(required=True)
    operation.add_argument(
        "--decrypt",
        dest="handler",
        action="store_const",
        const=decrypt_profile,
        help="decrypt a save to JSON",
    )
    operation.add_argument(
        "--encrypt",
        dest="handler",
        action="store_const",
        const=encrypt_profile,
        help="encrypt JSON for use by the game",
    )
    parser.add_argument(
        "input_path",
        type=Path,
        metavar="INPUT",
        help="input file",
    )
    parser.add_argument(
        "output_path",
        type=Path,
        metavar="OUTPUT",
        help="output file",
    )
    return parser

def main() -> None:
    arguments = create_parser().parse_args()
    arguments.handler(
        arguments.input_path,
        arguments.output_path,
        b"b5qhh8saJ8UlDJUzTZXd2Tg6mbo8W8n5",
    )

if __name__ == "__main__":
    main()