#!/usr/bin/env python3
"""Bounded, read-only PE structure inspection. Never execute the input."""

from __future__ import annotations

import argparse
import ctypes
import hashlib
import json
import os
from pathlib import Path
import stat
import struct
import sys
import zipfile
import zlib

MAX_BYTES = 128 * 1024 * 1024
MAX_SECTIONS = 96
MAX_ZIP_ENTRIES = 4096
MAX_CENTRAL_DIRECTORY_BYTES = 4 * 1024 * 1024
MAX_COMPRESSION_RATIO = 200
CHUNK_BYTES = 1024 * 1024
MACHINE_NAMES = {
    0x0000: "UNKNOWN", 0x014C: "I386", 0x01C0: "ARM", 0x01C2: "THUMB",
    0x01C4: "ARMNT", 0x0200: "IA64", 0x8664: "AMD64", 0xAA64: "ARM64",
    0xA641: "ARM64EC", 0xA64E: "ARM64X", 0x5032: "RISCV32",
    0x5064: "RISCV64", 0x5128: "RISCV128", 0x6232: "LOONGARCH32",
    0x6264: "LOONGARCH64", 0x0EBC: "EBC",
}


class InspectionError(ValueError):
    """Input is unsupported, exceeds bounds, or has invalid structure."""


def checked_range(offset: int, size: int, limit: int, label: str) -> int:
    """Validate a half-open range without relying on machine integer wraparound."""
    if offset < 0 or size < 0 or offset > limit or size > limit - offset:
        raise InspectionError(f"{label} is outside its enclosing range")
    return offset + size


def _limit(value: int) -> int:
    if not 1 <= value <= MAX_BYTES:
        raise InspectionError(f"byte limit must be between 1 and {MAX_BYTES}")
    return value


def _overlaps(start: int, end: int, other_start: int, other_end: int) -> bool:
    return start < other_end and other_start < end


def inspect_pe(data: bytes, max_bytes: int = MAX_BYTES) -> dict:
    """Inspect an already bounded byte string; no input code is executed."""
    _limit(max_bytes)
    size = len(data)
    checked_range(0, size, max_bytes, "input")
    checked_range(0, 64, size, "DOS header")
    if data[:2] != b"MZ":
        raise InspectionError("DOS signature is not MZ")
    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    if pe_offset < 64:
        raise InspectionError("PE header overlaps the fixed DOS header")
    checked_range(pe_offset, 24, size, "PE signature and COFF header")
    if data[pe_offset:pe_offset + 4] != b"PE\0\0":
        raise InspectionError("PE signature is not PE\\0\\0")
    (machine, section_count, timestamp, symbol_offset, symbol_count,
     optional_size, characteristics) = struct.unpack_from("<HHIIIHH", data, pe_offset + 4)
    if not 1 <= section_count <= MAX_SECTIONS:
        raise InspectionError(f"section count must be between 1 and {MAX_SECTIONS}")

    optional_offset = pe_offset + 24
    section_table = checked_range(optional_offset, optional_size, size, "optional header")
    checked_range(0, 2, optional_size, "optional header magic")
    magic = struct.unpack_from("<H", data, optional_offset)[0]
    if magic == 0x10B:
        pe_format, fixed_size, base_offset, base_format = "PE32", 96, 28, "<I"
    elif magic == 0x20B:
        pe_format, fixed_size, base_offset, base_format = "PE32+", 112, 24, "<Q"
    else:
        raise InspectionError("unsupported optional header magic (expected PE32 or PE32+)")
    checked_range(0, fixed_size, optional_size, "fixed optional header")
    directory_count = struct.unpack_from("<I", data, optional_offset + fixed_size - 4)[0]
    checked_range(fixed_size, directory_count * 8, optional_size, "data-directory array")
    table_end = checked_range(section_table, section_count * 40, size, "section table")
    image_base = struct.unpack_from(base_format, data, optional_offset + base_offset)[0]
    entry_rva = struct.unpack_from("<I", data, optional_offset + 16)[0]
    section_alignment, file_alignment = struct.unpack_from("<II", data, optional_offset + 32)
    image_size, header_size = struct.unpack_from("<II", data, optional_offset + 56)
    subsystem = struct.unpack_from("<H", data, optional_offset + 68)[0]
    if header_size < table_end:
        raise InspectionError("SizeOfHeaders does not contain the complete section table")
    checked_range(0, header_size, size, "declared headers")
    if image_size < header_size:
        raise InspectionError("SizeOfImage is smaller than SizeOfHeaders")
    if entry_rva and entry_rva >= image_size:
        raise InspectionError("entry-point RVA is outside SizeOfImage")

    sections = []
    raw_ranges = []
    for index in range(section_count):
        offset = section_table + index * 40
        name_bytes = data[offset:offset + 8].split(b"\0", 1)[0]
        # Make untrusted section-name controls visible, including terminal escapes.
        name = "".join(chr(c) if 32 <= c < 127 else f"\\x{c:02x}" for c in name_bytes)
        virtual_size, virtual_rva, raw_size, raw_offset = struct.unpack_from("<IIII", data, offset + 8)
        section_flags = struct.unpack_from("<I", data, offset + 36)[0]
        raw_end = checked_range(raw_offset, raw_size, size, f"section {index} raw data")
        virtual_end = checked_range(virtual_rva, virtual_size, 1 << 32, f"section {index} virtual data")
        mapped_size = max(virtual_size, raw_size)
        mapped_end = checked_range(virtual_rva, mapped_size, image_size, f"section {index} mapped extent")
        if mapped_size and virtual_rva < header_size:
            raise InspectionError(f"section {index} virtual extent overlaps headers")
        if raw_size:
            if raw_offset < header_size:
                raise InspectionError(f"section {index} raw data overlaps headers")
            for old_start, old_end in raw_ranges:
                if _overlaps(raw_offset, raw_end, old_start, old_end):
                    raise InspectionError(f"section {index} raw data overlaps another section")
            raw_ranges.append((raw_offset, raw_end))
        sections.append({
            "index": index, "name": name, "name_bytes_hex": data[offset:offset + 8].hex(),
            "raw_offset": raw_offset, "raw_size": raw_size,
            "raw_range": {"start": raw_offset, "end_exclusive": raw_end},
            "virtual_rva": virtual_rva, "virtual_size": virtual_size,
            "virtual_range": {"start_rva": virtual_rva, "end_rva_exclusive": virtual_end},
            "unrounded_mapped_extent": {"start_rva": virtual_rva, "end_rva_exclusive": mapped_end},
            "characteristics_hex": f"0x{section_flags:08x}",
        })

    security = {
        "declared": directory_count > 4, "present": False,
        "file_offset": None, "size": 0, "end_exclusive": None,
        "address_kind": "file_offset_not_rva", "signature_verification": "not_performed",
        "note": "The directory locates claimed certificate bytes; presence does not prove a valid signature or publisher.",
    }
    if directory_count > 4:
        cert_offset, cert_size = struct.unpack_from("<II", data, optional_offset + fixed_size + 4 * 8)
        if bool(cert_offset) != bool(cert_size):
            raise InspectionError("security-directory offset and size must both be zero or both nonzero")
        if cert_size:
            cert_end = checked_range(cert_offset, cert_size, size, "security directory")
            if cert_offset % 8:
                raise InspectionError("security-directory file offset is not aligned to 8 bytes")
            if cert_offset < header_size:
                raise InspectionError("security directory overlaps headers")
            if any(_overlaps(cert_offset, cert_end, a, b) for a, b in raw_ranges):
                raise InspectionError("security directory overlaps section raw data")
            security.update(present=True, file_offset=cert_offset, size=cert_size, end_exclusive=cert_end)

    last_raw_end = max((end for _, end in raw_ranges), default=None)
    trailing_offset = max(header_size, last_raw_end or 0)
    return {
        "schema_version": 1,
        "file": {"size_bytes": size, "sha256": hashlib.sha256(data).hexdigest()},
        "dos_signature": "MZ", "pe_signature": "PE\\0\\0", "pe_offset": pe_offset,
        "coff": {
            "machine": machine, "machine_hex": f"0x{machine:04x}",
            "machine_name": MACHINE_NAMES.get(machine, "unrecognized"),
            "number_of_sections": section_count, "timestamp_raw": timestamp,
            "optional_header_size": optional_size, "characteristics_hex": f"0x{characteristics:04x}",
            "symbol_table_file_offset_claim": symbol_offset, "symbol_count_claim": symbol_count,
        },
        "optional_header": {
            "format": pe_format, "magic_hex": f"0x{magic:04x}", "image_base": image_base,
            "entry_point_rva": entry_rva, "size_of_image": image_size,
            "size_of_headers": header_size, "section_alignment": section_alignment,
            "file_alignment": file_alignment, "subsystem": subsystem,
            "number_of_data_directories": directory_count,
        },
        "section_table": {"offset": section_table, "end_exclusive": table_end},
        "sections": sections, "last_claimed_section_raw_end": last_raw_end,
        "trailing_file_bytes": {
            "offset": trailing_offset, "size": size - trailing_offset,
            "note": "Bytes after headers and the last nonempty section raw range; may include certificate data, installer data, or other content. Not automatically malicious.",
        },
        "security_directory": security,
        "limitations": [
            "Static metadata only; no execution, disassembly, unpacking, or network access.",
            "SHA-256 identifies the observed bytes; it is not authentication or a malware verdict.",
            "The PE and DOS magic signatures are format markers, not cryptographic signatures.",
            "No Authenticode, certificate-chain, PE-checksum, publisher, or timestamp verification.",
            "Directory contents, COFF symbols, relocations, imports, virtual overlap, and loader compatibility are not validated.",
        ],
    }


def _local_regular_file(path: str | Path) -> Path:
    """Reject network/device paths and links before reading a regular local file."""
    raw = os.fspath(path)
    if not raw or raw.startswith(("\\\\", "//")) or "://" in raw or "\0" in raw:
        raise InspectionError("input must be a local filesystem path")
    absolute = Path(os.path.abspath(raw))
    if os.name == "nt":
        if ":" in str(absolute)[2:]:
            raise InspectionError("device paths and alternate data streams are unsupported")
        get_drive_type = ctypes.windll.kernel32.GetDriveTypeW
        get_drive_type.argtypes = [ctypes.c_wchar_p]
        get_drive_type.restype = ctypes.c_uint
        drive_type = get_drive_type(str(absolute.anchor))
        if drive_type not in (2, 3, 5, 6):
            raise InspectionError("input drive must be local")
    for component in reversed((absolute, *absolute.parents)):
        info = component.lstat()
        if stat.S_ISLNK(info.st_mode) or getattr(info, "st_file_attributes", 0) & 0x400:
            raise InspectionError("symlinks and reparse points are unsupported")
    if not stat.S_ISREG(absolute.stat().st_mode):
        raise InspectionError("input must be a regular file")
    return absolute


def _read_bounded(stream, limit: int) -> bytes:
    data = bytearray()
    while True:
        chunk = stream.read(min(CHUNK_BYTES, limit + 1 - len(data)))
        if not chunk:
            return bytes(data)
        data.extend(chunk)
        if len(data) > limit:
            raise InspectionError("input exceeds the byte limit")


def _member_name(name: str) -> None:
    if (not name or len(name) > 1024 or name.startswith("/") or "\\" in name
            or ":" in name or any(ord(c) < 32 or ord(c) == 127 for c in name)
            or any(part in ("", ".", "..") for part in name.split("/"))):
        raise InspectionError("archive member must be a safe, relative, exact file name")


def _zip_directory_bounds(stream, file_size: int) -> tuple[int, int]:
    """Preflight EOCD before ZipFile allocates objects for central-directory entries."""
    tail_size = min(file_size, 22 + 65535)
    stream.seek(file_size - tail_size)
    tail = stream.read(tail_size)
    position = tail.rfind(b"PK\x05\x06")
    if position < 0 or len(tail) - position < 22:
        raise InspectionError("ZIP end-of-central-directory record is missing or truncated")
    fields = struct.unpack_from("<4s4H2IH", tail, position)
    _, disk, start_disk, disk_entries, entries, directory_size, directory_offset, comment_size = fields
    if position + 22 + comment_size != len(tail):
        raise InspectionError("ZIP trailer or comment length is invalid")
    if disk or start_disk or disk_entries != entries:
        raise InspectionError("multi-disk ZIP archives are unsupported")
    if entries == 0xFFFF or directory_size == 0xFFFFFFFF or directory_offset == 0xFFFFFFFF:
        raise InspectionError("ZIP64 archives are unsupported")
    if entries > MAX_ZIP_ENTRIES or directory_size > MAX_CENTRAL_DIRECTORY_BYTES:
        raise InspectionError("ZIP central directory exceeds the metadata limit")
    end_position = file_size - tail_size + position
    if checked_range(directory_offset, directory_size, end_position, "ZIP central directory") != end_position:
        raise InspectionError("ZIP central directory must directly precede its end record")
    # Validate actual count too: an EOCD count is only an untrusted claim.
    stream.seek(directory_offset)
    directory = stream.read(directory_size)
    cursor = actual_entries = 0
    while cursor < directory_size:
        checked_range(cursor, 46, directory_size, "ZIP central entry")
        if directory[cursor:cursor + 4] != b"PK\x01\x02":
            raise InspectionError("ZIP central entry signature is invalid")
        if struct.unpack_from("<H", directory, cursor + 34)[0]:
            raise InspectionError("multi-disk ZIP central entries are unsupported")
        name_size, extra_size, entry_comment = struct.unpack_from("<HHH", directory, cursor + 28)
        cursor = checked_range(cursor, 46 + name_size + extra_size + entry_comment, directory_size, "ZIP central entry")
        actual_entries += 1
        if actual_entries > MAX_ZIP_ENTRIES:
            raise InspectionError("ZIP has too many central entries")
    if actual_entries != entries:
        raise InspectionError("ZIP entry count disagrees with its central directory")
    stream.seek(0)
    return entries, directory_offset


def _read_zip_member(stream, file_size: int, member: str, max_bytes: int) -> bytes:
    _member_name(member)
    count, directory_offset = _zip_directory_bounds(stream, file_size)
    with zipfile.ZipFile(stream, "r") as archive:
        infos = archive.infolist()
        if len(infos) != count:
            raise InspectionError("ZIP entry count changed during parsing")
        matches = [info for info in infos if info.filename == member]
        if len(matches) != 1:
            raise InspectionError("archive member must exist exactly once (duplicates are rejected)")
        info = matches[0]
        _member_name(info.orig_filename)
        if info.orig_filename != member or info.is_dir():
            raise InspectionError("archive member is not an exact regular file name")
        mode = info.external_attr >> 16
        if stat.S_IFMT(mode) not in (0, stat.S_IFREG):
            raise InspectionError("archive member is not a regular file")
        if info.flag_bits & (1 | 0x40):
            raise InspectionError("encrypted archive members are unsupported")
        if info.compress_type not in (zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED):
            raise InspectionError("only stored or deflated archive members are supported")
        if info.extract_version >= 45:
            raise InspectionError("ZIP64 and newer ZIP extensions are unsupported")
        if not 0 <= info.file_size <= max_bytes or not 0 <= info.compress_size <= max_bytes:
            raise InspectionError("archive member exceeds the byte limit")
        if info.file_size > MAX_COMPRESSION_RATIO * max(1, info.compress_size):
            raise InspectionError("archive member exceeds the compression-ratio limit")
        checked_range(info.header_offset, 30, directory_offset, "ZIP local header")
        stream.seek(info.header_offset)
        local = stream.read(30)
        if len(local) != 30 or local[:4] != b"PK\x03\x04":
            raise InspectionError("ZIP local header is invalid")
        flags, method = struct.unpack_from("<HH", local, 6)
        if flags != info.flag_bits or method != info.compress_type:
            raise InspectionError("ZIP local and central flags or methods disagree")
        if flags & ~(0x800 | 8 | 6):
            raise InspectionError("ZIP member uses unsupported flags")
        local_crc, local_compressed, local_size = struct.unpack_from("<III", local, 14)
        if not flags & 8 and (local_crc, local_compressed, local_size) != (info.CRC, info.compress_size, info.file_size):
            raise InspectionError("ZIP local and central CRC or sizes disagree")
        name_size, extra_size = struct.unpack_from("<HH", local, 26)
        payload_offset = checked_range(info.header_offset, 30 + name_size + extra_size, directory_offset, "ZIP local metadata")
        payload_end = checked_range(payload_offset, info.compress_size, directory_offset, "ZIP compressed member")
        raw_name = stream.read(name_size)
        try:
            decoded_name = raw_name.decode("utf-8" if flags & 0x800 else "cp437")
        except UnicodeError as error:
            raise InspectionError("ZIP local member name has invalid encoding") from error
        if decoded_name != member:
            raise InspectionError("ZIP local and central member names disagree")
        # Reject local records whose claimed payload includes another entry's header.
        next_header = min((other.header_offset for other in infos if other.header_offset > info.header_offset), default=directory_offset)
        if payload_end > next_header or any(other is not info and other.header_offset == info.header_offset for other in infos):
            raise InspectionError("ZIP member overlaps another local record")
        stream.seek(payload_offset)
        # Decompress independently: ZipExtFile may stop at an attacker-supplied
        # uncompressed size before reaching the actual end of a deflate stream.
        output = bytearray()
        remaining = info.compress_size
        inflater = zlib.decompressobj(-15) if method == zipfile.ZIP_DEFLATED else None
        while remaining:
            chunk = stream.read(min(CHUNK_BYTES, remaining))
            if not chunk:
                raise InspectionError("ZIP compressed member is truncated")
            remaining -= len(chunk)
            pending = chunk
            while pending:
                if inflater is None:
                    inflated, pending = pending, b""
                else:
                    inflated = inflater.decompress(pending, min(CHUNK_BYTES, info.file_size + 1 - len(output)))
                    pending = inflater.unconsumed_tail
                output.extend(inflated)
                if len(output) > info.file_size or len(output) > max_bytes:
                    raise InspectionError("ZIP inflated member exceeds its claimed size or byte limit")
                if inflater is not None and inflater.eof:
                    if inflater.unused_data or pending or remaining:
                        raise InspectionError("ZIP deflate stream has unclaimed trailing compressed bytes")
                    break
        if inflater is not None and not inflater.eof:
            raise InspectionError("ZIP deflate stream is truncated")
        data = bytes(output)
        if len(data) != info.file_size:
            raise InspectionError("archive member size disagrees with its directory")
        if zlib.crc32(data) & 0xFFFFFFFF != info.CRC:
            raise InspectionError("ZIP member CRC does not match")
        if flags & 8:
            checked_range(payload_end, 12, next_header, "ZIP data descriptor")
            stream.seek(payload_end)
            first = stream.read(4)
            if first == b"PK\x07\x08":
                checked_range(payload_end, 16, next_header, "ZIP data descriptor")
                descriptor = stream.read(12)
            else:
                descriptor = first + stream.read(8)
            if struct.unpack("<III", descriptor) != (info.CRC, info.compress_size, info.file_size):
                raise InspectionError("ZIP data descriptor disagrees with central directory")
        return data


def read_input(path: str | Path, archive_member: str | None = None, max_bytes: int = MAX_BYTES) -> bytes:
    """Read one file or one named ZIP member, with no extraction or output files."""
    _limit(max_bytes)
    local = _local_regular_file(path)
    with local.open("rb") as stream:
        before = os.fstat(stream.fileno())
        if not stat.S_ISREG(before.st_mode) or before.st_size > max_bytes:
            raise InspectionError("input file or archive exceeds the byte limit or is not regular")
        if archive_member is None:
            data = _read_bounded(stream, max_bytes)
            if len(data) != before.st_size:
                raise InspectionError("input size changed while reading")
        else:
            data = _read_zip_member(stream, before.st_size, archive_member, max_bytes)
        after = os.fstat(stream.fileno())
        if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns):
            raise InspectionError("input changed while reading")
        return data


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("file", help="local PE file, or ZIP when --archive-member is used")
    parser.add_argument("--archive-member", metavar="EXACT_NAME", help="read only this ZIP member, without extraction")
    parser.add_argument("--max-bytes", type=int, default=MAX_BYTES, help=f"input/member limit, 1..{MAX_BYTES} (default: %(default)s)")
    args = parser.parse_args(argv)
    try:
        data = read_input(args.file, args.archive_member, args.max_bytes)
        report = inspect_pe(data, args.max_bytes)
        report["input_kind"] = "zip_member" if args.archive_member is not None else "local_file"
        print(json.dumps(report, indent=2, ensure_ascii=True))
    except (InspectionError, OSError, zipfile.BadZipFile, NotImplementedError, RuntimeError,
            EOFError, zlib.error, UnicodeError, struct.error) as error:
        # Avoid printing untrusted paths, archive strings, or terminal controls.
        message = str(error) if isinstance(error, InspectionError) else "input could not be read or ZIP integrity validation failed"
        print(json.dumps({"error": message}, ensure_ascii=True), file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())