#!/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

AAD = b"Q05-Lidar init_boot resource v1"


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


def main():
    parser = argparse.ArgumentParser(
        description="Encrypt a Q05-Lidar init_boot image into resource.dat"
    )
    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(
        "--key",
        help="Optional AES-256 key as base64url or 64-char hex. 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()

    if args.key:
        text = args.key.strip()
        if len(text) == 64 and all(c in "0123456789abcdefABCDEF" for c in text):
            key = bytes.fromhex(text)
        else:
            key = base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
        if len(key) != 32:
            raise SystemExit("key must decode to 32 bytes")
    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": "q05-lidar-resource-v2",
        "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(b"Q05R2\x00" + 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()