#!/usr/bin/env python3
"""pad_for_optiplex.py — pad a UEFI bootloader PE for Optiplex 7010

The Optiplex 7010 AMI Aptio firmware rejects PEs smaller than ~160KB
("no bootable devices").  This pads the PE to match systemd-boot's size
and sets SizeOfImage accordingly.
"""

import struct, sys, os

def pad(path):
    try:
        with open(path, 'rb') as f:
            d = bytearray(f.read())
    except IOError as e:
        print(f"  ERROR: cannot read {path}: {e}")
        return False

    po = struct.unpack('<I', d[0x3C:0x40])[0]
    if d[po:po+4] != b'PE\x00\x00':
        print(f"  Not a PE file")
        return False

    opt = po + 24

    # Target size: 160256 bytes (matching systemd-boot)
    target = 160256
    if len(d) >= target:
        print(f"  Already >= {target} bytes ({len(d)}), nothing to do")
        return True

    # Compute SizeOfImage: round current SizeOfImage to match file growth
    cur_si = struct.unpack('<I', d[opt+56:opt+60])[0]
    cur_file_size = len(d)

    # The file is padded after all sections.  SizeOfImage is the highest
    # RVA + virtual size, rounded up to SectionAlignment.  After padding
    # the file, we need to increase SizeOfImage so the firmware sees a
    # consistent image.
    # 
    # Simple approach: set SizeOfImage = target - file_delta + old_si
    # Actually: SizeOfImage already covers our sections at up to ~0x4000.
    # We just need to make it large enough for the padded region.
    # Use systemd-boot's SizeOfImage = 0x2d000 as reference.
    file_growth = target - len(d)
    new_si = ((cur_si + file_growth + 0xfff) // 0x1000) * 0x1000

    # Pad the file
    d += b'\x00' * file_growth

    # Update SizeOfImage in optional header
    struct.pack_into('<I', d, opt+56, new_si)

    try:
        with open(path, 'wb') as f:
            f.write(d)
    except IOError as e:
        print(f"  ERROR: cannot write {path}: {e}")
        return False

    print(f"  Padded: {len(d)} bytes (was {cur_file_size}), SizeOfImage = {new_si:#x}")
    return True

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <pe-path>")
        sys.exit(1)
    if not pad(sys.argv[1]):
        sys.exit(1)