#!/usr/bin/env python3
"""Offline tests for check_zip_alignment.py.

Standard library only, no network, no nix, no external binary -- this is the
one piece of this repo that needs none of the pinned toolchain, so it is
tested at plain `python3` speed rather than through with-nixpkgs.sh.

Builds both a genuinely aligned fixture and a deliberately misaligned one
with nothing but `zipfile`/`struct`, and runs the checker as a *subprocess*
(not an import) against both, because that is how sign-apk.sh actually calls
it -- a test that only imports the function would miss an argv or exit-code
regression.

    python3 test_check_zip_alignment.py
"""

import os
import struct
import subprocess
import sys
import tempfile
import unittest
import zipfile

HERE = os.path.dirname(os.path.abspath(__file__))
CHECKER = os.path.join(HERE, "check_zip_alignment.py")


def add_stored(zf, name, data, align=None):
    """Write a STORED entry, padded to `align` bytes if given.

    Mirrors what a real zip-aligning tool does: pad the "extra field" of the
    local header so the entry's *data* (not its header) starts on the
    boundary. `align=None` writes the entry with no padding at all, which is
    what a plain zipfile.writestr call already does by default -- the
    negative-control case that must fail the check.
    """
    info = zipfile.ZipInfo(name)
    info.compress_type = zipfile.ZIP_STORED
    if align:
        header_len = 30 + len(name.encode())
        data_start_if_no_extra = zf.fp.tell() + header_len
        pad = (-data_start_if_no_extra) % align
        if pad:
            if pad < 4:
                pad += align
            extra_len = pad - 4
            info.extra = struct.pack("<HH", 0xD935, extra_len) + b"\x00" * extra_len
    zf.writestr(info, data)


class CheckZipAlignmentTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)

    def _run(self, path):
        return subprocess.run([sys.executable, CHECKER, path],
                              capture_output=True, text=True)

    def _build(self, name, entries):
        path = os.path.join(self.tmp.name, name)
        with zipfile.ZipFile(path, "w") as zf:
            for entry_name, data, align in entries:
                add_stored(zf, entry_name, data, align)
        return path

    def test_aligned_apk_passes(self):
        path = self._build("aligned.apk", [
            ("AndroidManifest.xml", b"0123456789", 4),
            ("lib/arm64-v8a/libfoo.so", b"X" * 50, 4096),
        ])
        result = self._run(path)
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("all uncompressed entries aligned", result.stdout)

    def test_misaligned_generic_entry_fails(self):
        # No padding at all -- the fault this check exists to catch.
        path = self._build("misaligned.apk", [
            ("x", b"pad", None),
            ("AndroidManifest.xml", b"0123456789", None),
        ])
        result = self._run(path)
        self.assertEqual(result.returncode, 1)
        self.assertIn("AndroidManifest.xml", result.stderr)
        self.assertIn("not a multiple of 4", result.stderr)

    def test_misaligned_native_library_fails_at_4096(self):
        # Aligned to 4 bytes (would pass a naive "mod 4" check) but not to
        # the 4096 a .so specifically needs.
        path = self._build("misaligned-so.apk", [
            ("lib/arm64-v8a/libfoo.so", b"X" * 50, 4),
        ])
        result = self._run(path)
        self.assertEqual(result.returncode, 1)
        self.assertIn("libfoo.so", result.stderr)
        self.assertIn("not a multiple of 4096", result.stderr)

    def test_compressed_entries_are_never_flagged(self):
        # Alignment only matters for STORED entries; a DEFLATEd one is
        # decompressed into a fresh buffer regardless of its file offset.
        path = os.path.join(self.tmp.name, "compressed.apk")
        with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as zf:
            zf.writestr("classes.dex", b"\x00" * 100)
        result = self._run(path)
        self.assertEqual(result.returncode, 0, result.stderr)

    def test_fixing_alignment_flips_fail_to_pass(self):
        """The fault-injection/revert pair the task asked for, in one test."""
        broken = self._build("broken.apk", [("AndroidManifest.xml", b"x" * 20, None)])
        self.assertEqual(self._run(broken).returncode, 1)

        fixed = self._build("fixed.apk", [("AndroidManifest.xml", b"x" * 20, 4)])
        self.assertEqual(self._run(fixed).returncode, 0)


if __name__ == "__main__":
    unittest.main()