finalize_pe.py

#!/usr/bin/env python3
"""Set and verify the PE checksum without third-party packages."""

from __future__ import annotations

import argparse
import struct
from pathlib import Path

def checksum_offset(data: bytes) -> int:
    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    if data[pe_offset:pe_offset + 4] != b"PE\0\0":
        raise ValueError("invalid PE signature")
    optional_offset = pe_offset + 24
    if struct.unpack_from("<H", data, optional_offset)[0] != 0x20B:
        raise ValueError("expected PE32+")
    return optional_offset + 64


def timestamp_offset(data: bytes) -> int:
    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    return pe_offset + 8


def pe_checksum(data: bytes, field_offset: int) -> int:
    total = 0
    padded = data + (b"\0" if len(data) & 1 else b"")
    for offset in range(0, len(padded), 2):
        word = 0 if field_offset <= offset < field_offset + 4 else struct.unpack_from("<H", padded, offset)[0]
        total = (total + word) & 0xFFFFFFFF
        total = (total & 0xFFFF) + (total >> 16)
    total = (total & 0xFFFF) + (total >> 16)
    return (total + len(data)) & 0xFFFFFFFF


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("executable", type=Path)
    args = parser.parse_args()

    data = bytearray(args.executable.read_bytes())
    # Fixed UTC build epoch (2026-08-23) makes identical sources reproducible.
    struct.pack_into("<I", data, timestamp_offset(data), 1787443200)
    field_offset = checksum_offset(data)
    struct.pack_into("<I", data, field_offset, 0)
    checksum = pe_checksum(data, field_offset)
    struct.pack_into("<I", data, field_offset, checksum)
    args.executable.write_bytes(data)
    final = args.executable.read_bytes()
    if struct.unpack_from("<I", final, field_offset)[0] != pe_checksum(final, field_offset):
        raise SystemExit("checksum verification failed")
    print(f"checksum=0x{checksum:08x}")


if __name__ == "__main__":
    main()