"""One-shot: parse APK Signing Block v2/v3 and print SHA-256(hex) of signing cert.

Matches what AntiTamperModule.kt computes — SHA-256 of the DER-encoded X.509 cert
that the PackageManager would return for the APK.

Spec: https://source.android.com/docs/security/features/apksigning/v2
"""
from __future__ import annotations

import hashlib
import struct
import sys
import zipfile
from pathlib import Path


APK_SIG_BLOCK_MAGIC = b"APK Sig Block 42"
V2_BLOCK_ID = 0x7109871A
V3_BLOCK_ID = 0xF05368C0
V3_1_BLOCK_ID = 0x1B93AD61


def _find_eocd(data: bytes) -> int:
    sig = b"PK\x05\x06"
    # EOCD must be in last 65557 bytes
    start = max(0, len(data) - 65557)
    idx = data.rfind(sig, start)
    if idx < 0:
        raise RuntimeError("EOCD not found")
    return idx


def _read_uint32_le(b: bytes, off: int) -> int:
    return struct.unpack_from("<I", b, off)[0]


def _read_uint64_le(b: bytes, off: int) -> int:
    return struct.unpack_from("<Q", b, off)[0]


def extract_sig_block(path: Path) -> bytes:
    data = path.read_bytes()
    eocd = _find_eocd(data)
    cd_offset = _read_uint32_le(data, eocd + 16)
    # Magic ends at cd_offset; the 8 bytes before magic are block size (excluding self)
    magic_end = cd_offset
    magic_start = magic_end - len(APK_SIG_BLOCK_MAGIC)
    if data[magic_start:magic_end] != APK_SIG_BLOCK_MAGIC:
        raise RuntimeError("APK Signing Block magic not found before central directory")
    size_off = magic_start - 8
    block_size_excl = _read_uint64_le(data, size_off)
    # The block layout: size_of_block(8) | pairs | size_of_block(8) | magic(16)
    block_total = block_size_excl + 8
    block_start = magic_end - block_total
    # Block = leading_size(8) | pairs | trailing_size(8) | magic(16)
    # pairs region = between leading_size and trailing_size
    return data[block_start + 8 : magic_start - 8]


def iter_pairs(pairs: bytes):
    i = 0
    n = len(pairs)
    while i < n:
        length = _read_uint64_le(pairs, i)
        i += 8
        pair_id = _read_uint32_le(pairs, i)
        value = pairs[i + 4 : i + length]
        yield pair_id, value
        i += length


def extract_cert_der_v2_or_v3(block_value: bytes) -> bytes:
    # block_value = "signers" sequence
    #   signers = length-prefixed sequence of signer
    #     signer = signed_data || signatures || public_key (each length-prefixed)
    #     signed_data = digests || certificates || additional_attributes (each length-prefixed)
    #       certificates = sequence of length-prefixed DER X.509 certs

    def read_lp(buf: bytes, off: int) -> tuple[bytes, int]:
        length = _read_uint32_le(buf, off)
        return buf[off + 4 : off + 4 + length], off + 4 + length

    # outer = signers sequence (already length-prefixed by caller? actually block_value IS the value)
    # Per spec, the block "value" begins with sequence of length-prefixed signer structures
    # i.e. no outer length prefix here.
    off = 0
    # The very first uint32 in block_value is the length of the signers sequence
    signers_seq, _ = read_lp(block_value, 0)
    off = 0
    signer, _ = read_lp(signers_seq, off)
    # signer = signed_data || min_sdk(4) || max_sdk(4) || signatures || public_key (v3 has sdk fields)
    # Simplest: signed_data is the first length-prefixed blob in signer.
    signed_data, _ = read_lp(signer, 0)
    # signed_data = digests || certificates || ...
    inner_off = 0
    digests, inner_off = read_lp(signed_data, inner_off)
    certs_seq, inner_off = read_lp(signed_data, inner_off)
    # certs_seq = sequence of length-prefixed DER certs
    first_cert, _ = read_lp(certs_seq, 0)
    return first_cert


def main(argv: list[str]) -> int:
    if len(argv) < 2:
        print("usage: _extract_apk_sig_hash.py <path-to-apk>", file=sys.stderr)
        return 2
    apk = Path(argv[1])
    if not apk.is_file():
        print(f"not found: {apk}", file=sys.stderr)
        return 2

    pairs = extract_sig_block(apk)
    found_block: bytes | None = None
    chosen_id = None
    for pid, value in iter_pairs(pairs):
        if pid in (V2_BLOCK_ID, V3_BLOCK_ID, V3_1_BLOCK_ID):
            # prefer v3 over v2 if both present (matches what PM returns on modern Android)
            if chosen_id in (None, V2_BLOCK_ID) and pid in (V3_BLOCK_ID, V3_1_BLOCK_ID):
                found_block = value
                chosen_id = pid
            elif chosen_id is None:
                found_block = value
                chosen_id = pid
    if found_block is None:
        print("no v2/v3 signing block found", file=sys.stderr)
        return 3

    cert_der = extract_cert_der_v2_or_v3(found_block)
    sha = hashlib.sha256(cert_der).hexdigest()
    print(sha)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))