#!/usr/bin/env python3
"""add_dylib.py — Add an LC_LOAD_DYLIB load command to a Mach-O 64-bit binary.

Usage: python3 add_dylib.py <binary> <dylib_path> [--weak]

Inserts the new load command right after the existing load commands, before
the first section data.  Requires enough free space in the __TEXT header
region (checked automatically).

The binary is modified in-place; a .orig backup is created first.
"""
import struct, sys, shutil, os

LC_LOAD_DYLIB     = 0x0c
LC_LOAD_WEAK_DYLIB = 0x80000018  # LC_LOAD_WEAK_DYLIB with LC_REQ_DYLD

def main():
    args = sys.argv[1:]
    weak = False
    if '--weak' in args:
        weak = True
        args.remove('--weak')
    if len(args) != 2:
        sys.exit("Usage: add_dylib.py <binary> <dylib_path> [--weak]")
    path, dylib = args

    with open(path, 'rb') as f:
        data = bytearray(f.read())

    # Parse Mach-O 64-bit header
    magic = struct.unpack_from('<I', data, 0)[0]
    if magic != 0xfeedfacf:
        sys.exit(f"Not a 64-bit Mach-O (magic={hex(magic)})")

    cputype, cpusub, filetype, ncmds, sizeofcmds, flags, reserved = \
        struct.unpack_from('<i i I I I I I', data, 4)

    HEADER_SIZE = 32  # mach_header_64
    hdr_end = HEADER_SIZE + sizeofcmds

    # Find the earliest section offset (file offset) to know our free space
    off = HEADER_SIZE
    min_section_off = len(data)
    for _ in range(ncmds):
        cmd, cmdsize = struct.unpack_from('<II', data, off)
        if cmd == 0x19:  # LC_SEGMENT_64
            # segment_command_64: cmd(4) cmdsize(4) segname(16) vmaddr(8) vmsize(8) fileoff(8) filesize(8) maxprot(4) initprot(4) nsects(4) flags(4)
            fileoff = struct.unpack_from('<Q', data, off + 40)[0]   # fileoff at offset 40
            nsects  = struct.unpack_from('<I', data, off + 64)[0]    # nsects at offset 64
            sect_off = off + 72  # section_64 array starts at segment + 72
            for s in range(nsects):
                sect_fileoff = struct.unpack_from('<I', data, sect_off + s * 80 + 48)[0]
                if sect_fileoff > 0 and sect_fileoff < min_section_off:
                    min_section_off = sect_fileoff
        off += cmdsize

    # Build the LC_LOAD_DYLIB command
    name = dylib.encode() + b'\0'
    # name_offset = 24 (cmd + cmdsize + 4*4 for dylib struct)
    name_offset = 24
    cmdsize = name_offset + len(name)
    # align to 8 bytes
    cmdsize = (cmdsize + 7) & ~7

    needed = cmdsize
    free = min_section_off - hdr_end
    if free < needed:
        sys.exit(f"Not enough free space: need {needed}, have {free} "
                 f"(hdr_end={hdr_end}, first_section={min_section_off})")

    # Build the command bytes
    cmd_id = LC_LOAD_WEAK_DYLIB if weak else LC_LOAD_DYLIB
    cmd = struct.pack('<II', cmd_id, cmdsize)
    cmd += struct.pack('<IIII', name_offset, 2, 0x10000, 0x10000)  # dylib struct
    cmd += name
    cmd += b'\0' * (cmdsize - len(cmd))  # pad to cmdsize

    # Write the new command into existing free space (NO insertion —
    # the space between sizeofcmds and first section is zero padding).
    # Inserting bytes would shift all section file offsets and break the binary.
    data[hdr_end:hdr_end + cmdsize] = cmd

    # Update ncmds and sizeofcmds (in-place, no shift)
    struct.pack_into('<I', data, 16, ncmds + 1)
    struct.pack_into('<I', data, 20, sizeofcmds + cmdsize)

    # Backup and write
    shutil.copy2(path, path + '.orig')
    with open(path, 'wb') as f:
        f.write(data)

    print(f"Added {'weak ' if weak else ''}LC_LOAD_DYLIB: {dylib}")
    print(f"  cmdsize={cmdsize}, ncmds={ncmds}->{ncmds+1}, "
          f"sizeofcmds={sizeofcmds}->{sizeofcmds+cmdsize}")
    print(f"  free space was {free} bytes, backup saved as {path}.orig")

if __name__ == '__main__':
    main()