"""
Python-native .intunewin file assembler.

Generates valid .intunewin packages using AES-256-CBC encryption with HMAC-SHA256,
producing the exact byte layout that Microsoft Intune expects.

Encrypted blob layout (from svrooij.io reverse-engineering):
    [0:32]   HMAC-SHA256 of the ciphertext (32 bytes, mac_key)
    [32:48]  AES-256-CBC Initialization Vector (16 bytes — standard AES block size)
    [48:]    AES-256-CBC ciphertext (PKCS7-padded to 16-byte boundary)

IMPORTANT: IV is 16 bytes, NOT 32. STACK.md has a documentation error on this point.

Detection.xml format matches the reference IntuneWinAppUtil.exe output exactly:
    - ToolVersion is an XML *attribute* on <ApplicationInfo> (not a child element)
    - No xmlns namespace declaration (reference uses [XmlRoot("ApplicationInfo")] with no namespace)
    - No <?xml ...?> declaration header (reference uses OmitXmlDeclaration=true)
    - No <MacAlgorithm> element (not present in reference FileEncryptionInfo model)

Outer ZIP structure:
    IntuneWinPackage/
    ├── Contents/
    │   └── IntunePackage.intunewin    (the encrypted blob)
    └── Metadata/
        └── Detection.xml              (encryption metadata)
"""
import base64
import hashlib
import hmac
import os
import shutil
import tempfile
import zipfile
from xml.etree.ElementTree import Element, SubElement, indent, tostring

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad


# Version string that matches the reference IntuneWinAppUtil.exe tool.
# Intune's upload wizard validates or uses this field to confirm the package
# was produced by a compatible tool version.
_TOOL_VERSION = "1.8.6.0"

# Streaming chunk size. Must be a multiple of AES.block_size (16) so every
# chunk but the last is a whole number of CBC blocks.
_CHUNK = 1024 * 1024


def build_intunewin(source_dir: str, setup_file: str, output_path: str) -> None:
    """Build a .intunewin file from source_dir, with setup_file as entry point.

    Args:
        source_dir: Path to the directory containing files to package.
        setup_file: Name of the setup/entry-point file (e.g., "install.ps1").
                    Must be present in source_dir. Used in Detection.xml metadata.
        output_path: Destination path for the generated .intunewin file.

    Raises:
        FileNotFoundError: If source_dir does not exist.
        ValueError: If setup_file is empty.

    Driver payloads run to ~100 MB, so every step streams through 1 MB chunks
    against temp files: buffering the inner ZIP, its ciphertext and the outer
    ZIP in memory held 3-4 full copies of the package per concurrent download.
    """
    if not setup_file:
        raise ValueError("setup_file must not be empty")
    if not os.path.isdir(source_dir):
        raise FileNotFoundError(f"source_dir does not exist: {source_dir}")

    with tempfile.TemporaryDirectory(prefix="intunewin_") as staging:
        inner_zip_path = os.path.join(staging, "inner.zip")
        ciphertext_path = os.path.join(staging, "inner.enc")

        # --- Step 1: Create inner ZIP (DEFLATE-compressed content) ---
        with zipfile.ZipFile(inner_zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
            for root, dirs, files in os.walk(source_dir):
                dirs.sort()   # deterministic ordering
                for filename in sorted(files):
                    abs_path = os.path.join(root, filename)
                    arc_name = os.path.relpath(abs_path, source_dir)
                    # Normalise to forward slashes for cross-platform consistency
                    arc_name = arc_name.replace("\\", "/")
                    zf.write(abs_path, arc_name)
        plaintext_size = os.path.getsize(inner_zip_path)

        # --- Step 2: Generate random keys and IV ---
        aes_key = os.urandom(32)   # 256-bit AES key
        mac_key = os.urandom(32)   # 256-bit HMAC key (same size as AES key)
        iv = os.urandom(16)        # 128-bit IV — standard AES-CBC block size (NOT 32 bytes)

        # --- Steps 3-4: Encrypt (AES-256-CBC, PKCS7) while MAC-ing and digesting ---
        # The reference (svrooij/ContentPrep Zipper.cs DecryptFileAsync) reads the first
        # 32 bytes as the stored HMAC, then computes the hash of the *remaining* bytes
        # (= IV || ciphertext) to verify integrity.  Authenticated-encryption best
        # practice (Encrypt-then-MAC) also requires the IV to be covered by the MAC so
        # that a forged IV cannot redirect decryption.
        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        mac = hmac.new(mac_key, iv, hashlib.sha256)
        plaintext_hash = hashlib.sha256()

        with open(inner_zip_path, "rb") as src, open(ciphertext_path, "wb") as dst:
            while True:
                # BufferedReader.read(n) returns n bytes unless EOF, so a short
                # read means "last chunk" — the only one that carries padding.
                chunk = src.read(_CHUNK)
                plaintext_hash.update(chunk)
                if len(chunk) < _CHUNK:
                    block = cipher.encrypt(pad(chunk, AES.block_size))
                    mac.update(block)
                    dst.write(block)
                    break
                block = cipher.encrypt(chunk)
                mac.update(block)
                dst.write(block)

        mac_digest = mac.digest()

        # --- Step 6: Plaintext (inner ZIP) SHA256 digest for Detection.xml ---
        file_digest = plaintext_hash.digest()

        return _finalize(
            output_path=output_path,
            ciphertext_path=ciphertext_path,
            setup_file=setup_file,
            plaintext_size=plaintext_size,
            aes_key=aes_key,
            mac_key=mac_key,
            iv=iv,
            mac_digest=mac_digest,
            file_digest=file_digest,
        )


def _finalize(
    *,
    output_path: str,
    ciphertext_path: str,
    setup_file: str,
    plaintext_size: int,
    aes_key: bytes,
    mac_key: bytes,
    iv: bytes,
    mac_digest: bytes,
    file_digest: bytes,
) -> None:
    """Write Detection.xml plus the [HMAC(32)][IV(16)][ciphertext] blob into the outer ZIP."""
    # --- Step 7: Build Detection.xml ---
    # Format MUST match IntuneWinAppUtil.exe reference output exactly:
    #   - ToolVersion is an XML attribute on ApplicationInfo (not a child element)
    #   - No xmlns namespace (reference omits it)
    #   - No <?xml?> declaration header
    #   - No MacAlgorithm element (not in reference FileEncryptionInfo model)
    app_info = Element(
        "ApplicationInfo",
        attrib={"ToolVersion": _TOOL_VERSION},
    )
    SubElement(app_info, "Name").text = setup_file
    SubElement(app_info, "UnencryptedContentSize").text = str(plaintext_size)
    SubElement(app_info, "FileName").text = "IntunePackage.intunewin"
    SubElement(app_info, "SetupFile").text = setup_file

    enc_info = SubElement(app_info, "EncryptionInfo")
    SubElement(enc_info, "EncryptionKey").text = base64.b64encode(aes_key).decode()
    SubElement(enc_info, "MacKey").text = base64.b64encode(mac_key).decode()
    SubElement(enc_info, "InitializationVector").text = base64.b64encode(iv).decode()
    SubElement(enc_info, "Mac").text = base64.b64encode(mac_digest).decode()
    SubElement(enc_info, "ProfileIdentifier").text = "ProfileVersion1"
    SubElement(enc_info, "FileDigest").text = base64.b64encode(file_digest).decode()
    SubElement(enc_info, "FileDigestAlgorithm").text = "SHA256"

    # indent() adds pretty-print whitespace in-place (Python 3.9+).
    # tostring with xml_declaration=False omits the <?xml?> header.
    indent(app_info, space="  ")
    detection_xml = tostring(app_info, encoding="unicode", xml_declaration=False)

    # --- Step 8: Build outer ZIP (STORED — no extra compression on encrypted content) ---
    # The blob is streamed in as [HMAC(32)] + [IV(16)] + [ciphertext] so the
    # encrypted payload is never materialised as a second in-memory copy.
    with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_STORED) as outer:
        with outer.open("IntuneWinPackage/Contents/IntunePackage.intunewin", "w") as dest:
            dest.write(mac_digest)
            dest.write(iv)
            with open(ciphertext_path, "rb") as ct:
                shutil.copyfileobj(ct, dest, _CHUNK)
        outer.writestr(
            "IntuneWinPackage/Metadata/Detection.xml",
            detection_xml.encode("utf-8"),
        )