#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import base64
import hashlib
import json
import os
import struct
import zlib
from pathlib import Path

DEFAULT_AAD = "Q05-Lidar init_boot resource v1"
FORMAT_MAGIC = {
    "q05": (b"Q05R2\x00", "q05-lidar-resource-v2", DEFAULT_AAD),
    "ez60": (b"EZ60R2\x00", "ez60-resource-v2", "Mazda-EZ60 init_boot resource v1"),
}


def b64u(data):
    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")


def decode_key_material(key_text):
    text = str(key_text).strip()
    if text.startswith("raw:"):
        key = text.split(":", 1)[1].encode("utf-8")
    elif len(text) == 64 and all(c in "0123456789abcdefABCDEF" for c in text):
        key = bytes.fromhex(text)
    elif len(text.encode("utf-8")) == 32:
        key = text.encode("utf-8")
    else:
        key = base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
    if len(key) != 32:
        raise ValueError(f"key must decode to 32 bytes, got {len(key)} bytes")
    return key


def main():
    parser = argparse.ArgumentParser(
        description="Encrypt a vehicle init_boot image into a resource dat file"
    )
    parser.add_argument("input_img", help="Magisk-patched init_boot image")
    parser.add_argument(
        "-o", "--output",
        default="resource.dat",
        help="Output encrypted resource file, default: resource.dat",
    )
    parser.add_argument(
        "--vehicle",
        choices=sorted(FORMAT_MAGIC.keys()),
        default="q05",
        help="Resource header/profile to write. Default: q05.",
    )
    parser.add_argument(
        "--aad",
        help="Optional AES-GCM AAD text. Defaults to the selected vehicle profile.",
    )
    parser.add_argument(
        "--key",
        help="Optional AES-256 key as base64url, 64-char hex, raw:TEXT, or 32-byte UTF-8 text. Omit to generate a random key.",
    )
    parser.add_argument(
        "--no-compress",
        action="store_true",
        help="Disable zlib compression before encryption.",
    )
    args = parser.parse_args()

    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError:
        raise SystemExit(
            "Missing dependency: cryptography. Install it with: "
            "python -m pip install cryptography"
        )

    input_path = Path(args.input_img)
    plain = input_path.read_bytes()
    sha256 = hashlib.sha256(plain).hexdigest()
    magic, resource_format, default_aad = FORMAT_MAGIC[args.vehicle]
    aad = (args.aad or default_aad).encode("utf-8")

    if args.key:
        try:
            key = decode_key_material(args.key)
        except Exception as e:
            raise SystemExit(str(e))
    else:
        key = os.urandom(32)

    compression = "none" if args.no_compress else "zlib"
    body = plain if args.no_compress else zlib.compress(plain, level=9)
    nonce = os.urandom(12)
    ciphertext = AESGCM(key).encrypt(nonce, body, aad)

    header = {
        "format": resource_format,
        "cipher": "AES-256-GCM",
        "kdf": "none",
        "compression": compression,
        "aad": aad.decode("utf-8"),
        "nonce": b64u(nonce),
        "sha256": sha256,
        "plainSize": len(plain),
        "packedSize": len(body),
    }

    output_path = Path(args.output)
    header_bytes = json.dumps(header, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    output_path.write_bytes(magic + struct.pack(">I", len(header_bytes)) + header_bytes + ciphertext)

    print("resource.dat created")
    print(f"input: {input_path}")
    print(f"output: {output_path}")
    print(f"plainSize: {len(plain)}")
    print(f"packedSize: {len(body)}")
    print(f"sha256: {sha256}")
    print(f"cloudKey: {b64u(key)}")


if __name__ == "__main__":
    main()