ImpTune/tests/test_intunewin.py

"""
Byte-level validation tests for the .intunewin file format.

These tests serve as the format specification: if they pass, the byte layout is correct.
The only remaining validation is a real Intune upload (manual, Phase 10 gate).

Detection.xml format follows the IntuneWinAppUtil.exe reference exactly:
  - ToolVersion is an XML *attribute* on <ApplicationInfo> (not a child element)
  - No xmlns namespace (reference uses [XmlRoot("ApplicationInfo")] with no Namespace param)
  - No <?xml?> declaration header (reference uses OmitXmlDeclaration=true)
  - No MacAlgorithm element (not present in reference FileEncryptionInfo model)
"""
import base64
import hashlib
import hmac
import io
import os
import xml.etree.ElementTree as ET
import zipfile

import pytest
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad

from imptune.generators.intunewin_builder import _TOOL_VERSION, build_intunewin


@pytest.fixture
def source_dir(tmp_path):
    """Create a small test source directory with a few files."""
    src = tmp_path / "source"
    src.mkdir()
    (src / "install.ps1").write_text("Write-Host 'Installing printer...'")
    (src / "config.json").write_text('{"printer": "HP LaserJet"}')
    (src / "readme.txt").write_text("Printer deployment package")
    return str(src)


@pytest.fixture
def built_package(source_dir, tmp_path):
    """Build an .intunewin package and return the path."""
    output = str(tmp_path / "package.intunewin")
    build_intunewin(source_dir, "install.ps1", output)
    return output


@pytest.fixture
def package_contents(built_package):
    """Extract outer ZIP contents and parsed Detection.xml."""
    with zipfile.ZipFile(built_package, "r") as outer:
        names = outer.namelist()
        blob = outer.read("IntuneWinPackage/Contents/IntunePackage.intunewin")
        detection_xml_bytes = outer.read("IntuneWinPackage/Metadata/Detection.xml")

    tree = ET.fromstring(detection_xml_bytes.decode("utf-8"))
    return {
        "names": names,
        "blob": blob,
        "detection_xml_bytes": detection_xml_bytes,
        "tree": tree,
    }


def _get_xml_text(tree, tag):
    """Get text of a direct child element (no namespace — reference omits xmlns)."""
    elem = tree.find(tag)
    return elem.text if elem is not None else None


def _get_encryption_info(tree):
    """Get EncryptionInfo sub-element (no namespace — reference omits xmlns)."""
    return tree.find("EncryptionInfo")


def _get_enc_text(tree, tag):
    """Get text of a child of EncryptionInfo (no namespace)."""
    enc = _get_encryption_info(tree)
    if enc is None:
        return None
    elem = enc.find(tag)
    return elem.text if elem is not None else None


class TestOuterZipStructure:
    def test_output_is_valid_zip(self, built_package):
        """build_intunewin() output file is a valid ZIP archive."""
        assert zipfile.is_zipfile(built_package), "Output file must be a valid ZIP archive"

    def test_outer_zip_structure(self, package_contents):
        """Outer ZIP contains exactly the two required entries."""
        names = set(package_contents["names"])
        assert "IntuneWinPackage/Contents/IntunePackage.intunewin" in names
        assert "IntuneWinPackage/Metadata/Detection.xml" in names

    def test_outer_zip_stored(self, built_package):
        """Outer ZIP entries use ZIP_STORED compression (no extra compression on encrypted content)."""
        with zipfile.ZipFile(built_package, "r") as outer:
            for info in outer.infolist():
                assert info.compress_type == zipfile.ZIP_STORED, (
                    f"Entry {info.filename} must use ZIP_STORED, got compress_type={info.compress_type}"
                )


class TestDetectionXml:
    def test_detection_xml_valid(self, package_contents):
        """Detection.xml is valid XML with ApplicationInfo root element and ToolVersion attribute.

        Reference format: <ApplicationInfo ToolVersion="1.8.6.0"> with NO xmlns namespace.
        Presence of xmlns would change element identity for Intune's XML parser, causing
        silent metadata-parse failure in the upload wizard.
        """
        tree = package_contents["tree"]
        assert tree.tag == "ApplicationInfo", (
            f"Root element must be plain 'ApplicationInfo' (no xmlns namespace), got {tree.tag!r}"
        )
        assert tree.get("ToolVersion") == _TOOL_VERSION, (
            f"ApplicationInfo must have ToolVersion attribute = {_TOOL_VERSION!r}, "
            f"got {tree.get('ToolVersion')!r}"
        )

    def test_detection_xml_fields(self, package_contents):
        """Detection.xml contains all required elements matching the reference FileEncryptionInfo model.

        The reference model has 7 EncryptionInfo sub-elements (MacAlgorithm is NOT present).
        """
        tree = package_contents["tree"]
        # Direct children
        for field in ("Name", "UnencryptedContentSize", "FileName", "SetupFile"):
            assert _get_xml_text(tree, field) is not None, f"Missing field: {field}"

        # EncryptionInfo sub-elements — 7 required (MacAlgorithm absent per reference schema)
        for field in (
            "EncryptionKey",
            "MacKey",
            "InitializationVector",
            "Mac",
            "ProfileIdentifier",
            "FileDigest",
            "FileDigestAlgorithm",
        ):
            assert _get_enc_text(tree, field) is not None, f"Missing EncryptionInfo/{field}"

        # MacAlgorithm must NOT be present (not in reference FileEncryptionInfo model)
        assert _get_enc_text(tree, "MacAlgorithm") is None, (
            "EncryptionInfo/MacAlgorithm must NOT be present — not in reference schema"
        )

    def test_setup_file_in_detection_xml(self, package_contents):
        """SetupFile element matches the setup_file argument passed to build_intunewin."""
        tree = package_contents["tree"]
        assert _get_xml_text(tree, "SetupFile") == "install.ps1"


class TestEncryptedBlobLayout:
    def test_encrypted_blob_layout(self, package_contents):
        """Encrypted blob starts with 32 bytes (HMAC) + 16 bytes (IV) + remainder (ciphertext)."""
        blob = package_contents["blob"]
        # Must be at least 48 bytes (HMAC + IV) plus at least one AES block (16 bytes)
        assert len(blob) >= 64, f"Blob too short: {len(blob)} bytes"
        # Total length = 48 header + ciphertext length; ciphertext length is a multiple of 16
        ciphertext_len = len(blob) - 48
        assert ciphertext_len > 0, "Blob has no ciphertext after header"
        assert ciphertext_len % 16 == 0, (
            f"Ciphertext length {ciphertext_len} must be a multiple of AES block size 16"
        )

    def test_iv_is_16_bytes(self, package_contents):
        """IV extracted from Detection.xml base64-decodes to exactly 16 bytes (NOT 32 — critical per RESEARCH.md)."""
        tree = package_contents["tree"]
        iv_b64 = _get_enc_text(tree, "InitializationVector")
        assert iv_b64 is not None, "InitializationVector missing from Detection.xml"
        iv = base64.b64decode(iv_b64)
        assert len(iv) == 16, f"IV must be exactly 16 bytes, got {len(iv)}"

    def test_encryption_key_is_32_bytes(self, package_contents):
        """EncryptionKey from Detection.xml base64-decodes to exactly 32 bytes."""
        tree = package_contents["tree"]
        key_b64 = _get_enc_text(tree, "EncryptionKey")
        assert key_b64 is not None, "EncryptionKey missing from Detection.xml"
        key = base64.b64decode(key_b64)
        assert len(key) == 32, f"EncryptionKey must be exactly 32 bytes, got {len(key)}"

    def test_mac_key_is_32_bytes(self, package_contents):
        """MacKey from Detection.xml base64-decodes to exactly 32 bytes."""
        tree = package_contents["tree"]
        mac_key_b64 = _get_enc_text(tree, "MacKey")
        assert mac_key_b64 is not None, "MacKey missing from Detection.xml"
        mac_key = base64.b64decode(mac_key_b64)
        assert len(mac_key) == 32, f"MacKey must be exactly 32 bytes, got {len(mac_key)}"


class TestCryptographicVerification:
    def test_hmac_matches(self, package_contents):
        """HMAC-SHA256 of (IV || ciphertext) matches first 32 bytes of blob AND Mac in Detection.xml.

        The reference implementation (svrooij/ContentPrep Zipper.cs DecryptFileAsync) reads
        the first 32 bytes as the stored HMAC, then computes the hash of the *remaining* bytes
        (bytes[32:] = IV || ciphertext) to verify integrity.  The IV MUST be included in the
        HMAC so that a forged IV cannot redirect decryption without being detected.
        """
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        blob_hmac = blob[:32]
        iv_and_ciphertext = blob[32:]   # IV (16 bytes) + ciphertext — what the reference hashes

        mac_key_b64 = _get_enc_text(tree, "MacKey")
        mac_key = base64.b64decode(mac_key_b64)

        computed_hmac = hmac.new(mac_key, iv_and_ciphertext, hashlib.sha256).digest()

        # Must match the blob header
        assert computed_hmac == blob_hmac, (
            "HMAC-SHA256 of (IV || ciphertext) does not match the first 32 bytes of the blob"
        )

        # Must also match Detection.xml Mac field
        xml_mac = base64.b64decode(_get_enc_text(tree, "Mac"))
        assert computed_hmac == xml_mac, (
            "HMAC-SHA256 of (IV || ciphertext) does not match the Mac value in Detection.xml"
        )

    def test_decryption_roundtrip(self, source_dir, package_contents):
        """Decrypt the ciphertext and verify it is a valid ZIP containing the original source files."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        ciphertext = blob[48:]

        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

        # Must be a valid ZIP
        assert zipfile.is_zipfile(io.BytesIO(plaintext)), (
            "Decrypted plaintext is not a valid ZIP file"
        )

        # Must contain the original source files
        with zipfile.ZipFile(io.BytesIO(plaintext), "r") as inner_zip:
            inner_names = set(inner_zip.namelist())

        for filename in ("install.ps1", "config.json", "readme.txt"):
            assert filename in inner_names, (
                f"Original file {filename} not found in decrypted inner ZIP. Found: {inner_names}"
            )

    def test_file_digest_matches(self, package_contents):
        """FileDigest in Detection.xml matches SHA256 of the decrypted plaintext ZIP."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        ciphertext = blob[48:]

        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

        computed_digest = hashlib.sha256(plaintext).digest()
        xml_digest = base64.b64decode(_get_enc_text(tree, "FileDigest"))

        assert computed_digest == xml_digest, (
            "FileDigest in Detection.xml does not match SHA256 of decrypted plaintext"
        )

    def test_unencrypted_content_size(self, package_contents):
        """UnencryptedContentSize in Detection.xml matches byte length of decrypted plaintext ZIP."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        ciphertext = blob[48:]

        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

        xml_size = int(_get_xml_text(tree, "UnencryptedContentSize"))
        assert xml_size == len(plaintext), (
            f"UnencryptedContentSize {xml_size} does not match actual plaintext size {len(plaintext)}"
        )


class TestStreamingChunkBoundaries:
    """The builder encrypts the inner ZIP in _CHUNK-sized pieces rather than in
    one buffer. Only the final (short) read carries PKCS7 padding, so a payload
    landing exactly on a chunk boundary is the case that breaks first."""

    @staticmethod
    def _decrypt(built_package):
        with zipfile.ZipFile(built_package, "r") as outer:
            blob = outer.read("IntuneWinPackage/Contents/IntunePackage.intunewin")
            tree = ET.fromstring(
                outer.read("IntuneWinPackage/Metadata/Detection.xml").decode("utf-8")
            )
        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(blob[48:]), AES.block_size)

        # MAC covers IV || ciphertext and must survive chunked hashing
        mac_key = base64.b64decode(_get_enc_text(tree, "MacKey"))
        assert hmac.new(mac_key, blob[32:], hashlib.sha256).digest() == blob[:32]
        # Metadata is computed while streaming, not from a buffered plaintext
        assert base64.b64decode(_get_enc_text(tree, "FileDigest")) == hashlib.sha256(plaintext).digest()
        assert int(_get_xml_text(tree, "UnencryptedContentSize")) == len(plaintext)
        return plaintext

    @pytest.mark.parametrize("payload_size", range(96, 128))
    def test_every_residue_of_chunk_size(self, tmp_path, monkeypatch, payload_size):
        """Sweep 32 consecutive sizes with a 16-byte chunk so every offset mod
        _CHUNK is exercised, including the exact-multiple case where the last
        read returns b"" and padding is a whole standalone block."""
        monkeypatch.setattr("imptune.generators.intunewin_builder._CHUNK", 16)

        src = tmp_path / "src"
        src.mkdir()
        # Incompressible, so inner-ZIP size tracks payload size 1:1
        (src / "install.ps1").write_bytes(os.urandom(payload_size))
        output = str(tmp_path / "package.intunewin")

        build_intunewin(str(src), "install.ps1", output)
        plaintext = self._decrypt(output)
        assert zipfile.is_zipfile(io.BytesIO(plaintext))

    def test_multi_megabyte_payload_roundtrips(self, tmp_path):
        """Several full 1 MB chunks through the default code path."""
        src = tmp_path / "src"
        src.mkdir()
        blob = os.urandom(2_500_000)
        (src / "install.ps1").write_text("Write-Host 'go'")
        (src / "driver.bin").write_bytes(blob)
        output = str(tmp_path / "package.intunewin")

        build_intunewin(str(src), "install.ps1", output)
        plaintext = self._decrypt(output)

        with zipfile.ZipFile(io.BytesIO(plaintext)) as inner:
            assert inner.read("driver.bin") == blob


class TestArgumentValidation:
    def test_missing_source_dir_raises(self, tmp_path):
        with pytest.raises(FileNotFoundError):
            build_intunewin(str(tmp_path / "nope"), "install.ps1", str(tmp_path / "o.intunewin"))

    def test_empty_setup_file_raises(self, source_dir, tmp_path):
        with pytest.raises(ValueError):
            build_intunewin(source_dir, "", str(tmp_path / "o.intunewin"))