"""zipalign replacement: rewrites an APK so every STORED entry's data starts on a 4-byte boundary.

Re-compressing deflated entries doesn't break the v1 (JAR) signature, which covers uncompressed
contents only.
"""
import struct
import sys
import zipfile

ALIGN = 4


def align(src, dst):
    zin = zipfile.ZipFile(src)
    with zipfile.ZipFile(dst, "w") as zout:
        for info in zin.infolist():
            data = zin.read(info.filename)
            zi = zipfile.ZipInfo(info.filename, date_time=info.date_time)
            zi.compress_type = info.compress_type
            zi.external_attr = info.external_attr
            zi.create_system = info.create_system
            extra = b""
            if info.compress_type == zipfile.ZIP_STORED:
                offset = zout.fp.tell()
                name_len = len(zi.filename.encode("utf-8"))
                base = offset + 30 + name_len + 6  # 6 = our extra-field header + alignment value
                pad = (ALIGN - base % ALIGN) % ALIGN
                # 0xD935 is the extra-field id Android's own zipalign uses for padding
                extra = struct.pack("<HHH", 0xD935, 2 + pad, ALIGN) + b"\0" * pad
            zi.extra = extra
            zout.writestr(zi, data)
    verify(dst)


def verify(path):
    z = zipfile.ZipFile(path)
    raw = open(path, "rb").read()
    for info in z.infolist():
        if info.compress_type != zipfile.ZIP_STORED:
            continue
        o = info.header_offset
        n, e = struct.unpack("<HH", raw[o + 26:o + 30])
        start = o + 30 + n + e
        if start % ALIGN:
            raise SystemExit("misaligned: %s at %d" % (info.filename, start))
    if z.testzip() is not None:
        raise SystemExit("corrupt zip")


if __name__ == "__main__":
    align(sys.argv[1], sys.argv[2])