"""
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 io
import os
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"


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.
    """
    # --- Step 1: Create inner ZIP (DEFLATE-compressed content) ---
    inner_zip_buf = io.BytesIO()
    with zipfile.ZipFile(inner_zip_buf, "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 = inner_zip_buf.getvalue()

    # --- 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)

    # --- Step 3: Encrypt with AES-256-CBC (PKCS7 padding) ---
    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
    ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))

    # --- Step 4: Compute HMAC-SHA256 over (IV + ciphertext) using mac_key ---
    # 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.
    mac_digest = hmac.new(mac_key, iv + ciphertext, hashlib.sha256).digest()

    # --- Step 5: Assemble encrypted blob: [HMAC(32)] + [IV(16)] + [ciphertext] ---
    encrypted_blob = mac_digest + iv + ciphertext

    # --- Step 6: Compute plaintext (inner ZIP) SHA256 digest for Detection.xml ---
    file_digest = hashlib.sha256(plaintext).digest()

    # --- 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(len(plaintext))
    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) ---
    with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_STORED) as outer:
        outer.writestr(
            "IntuneWinPackage/Contents/IntunePackage.intunewin",
            encrypted_blob,
        )
        outer.writestr(
            "IntuneWinPackage/Metadata/Detection.xml",
            detection_xml.encode("utf-8"),
        )