#!/usr/bin/env python3
"""Properly add an LC_LOAD_DYLIB entry to a Mach-O arm64 binary."""

import struct, sys

def add_dylib(macho_path, dylib_name):
    with open(macho_path, 'r+b') as f:
        header = f.read(32)
        magic, cputype, cpusubtype, filetype, ncmds, sizeofcmds, flags, reserved = struct.unpack('<IIIIIIII', header)
        
        if magic != 0xFEEDFACF:
            print("Not arm64 Mach-O")
            return False
        
        print(f"ncmds={ncmds} sizeofcmds={sizeofcmds}")
        
        # Read all load commands
        f.seek(32)
        cmds_data = f.read(sizeofcmds)
        
        # Find LC_SYMTAB to get string table offset (for fixing LINKEDIT)
        pos = 0
        while pos < sizeofcmds:
            cmd, sz = struct.unpack('<II', cmds_data[pos:pos+8])
            if cmd == 0x2:  # LC_SYMTAB
                symoff, nsyms, stroff, strsize = struct.unpack('<IIII', cmds_data[pos+8:pos+24])
                print(f"LC_SYMTAB: symoff={symoff} stroff={stroff} strsize={strsize}")
                break
            pos += sz
        
        # Build the LC_LOAD_DYLIB entry (LC_LOAD_DYLIB = 0xC)
        # struct dylib_command:
        #   cmd(4) + cmdsize(4) + name_off(4) + timestamp(4) + cur_ver(4) + compat_ver(4) + name_string(...)
        name_bytes = dylib_name.encode() + b'\x00'
        
        # name_offset = offset from start of this load command to name string
        # It's at offset 24: 8 (cmd+cmdsize) + 16 (dylib struct) = 24
        name_off = 24
        
        total_size = 24 + len(name_bytes)
        total_size = (total_size + 7) & ~7  # align to 8
        
        entry = struct.pack('<IIIIIII', 
            0xC,           # cmd
            total_size,    # cmdsize  
            name_off,      # dylib.name (offset)
            0,             # timestamp
            0,             # current_version
            0,             # compat_version
        )
        entry += name_bytes
        while len(entry) < total_size:
            entry += b'\x00'
        
        # Update header
        new_size = sizeofcmds + entry_size
        f.seek(0)
        f.write(struct.pack('<IIIIIIII', magic, cputype, cpusubtype, filetype, ncmds + 1, new_size, flags, reserved))
        
        # Append to end of load commands
        f.seek(32 + sizeofcmds)
        f.write(entry)
        
        print(f"Added LC_LOAD_DYLIB for {dylib_name}")
        print(f"ncmds: {ncmds} -> {ncmds+1}, sizeofcmds: {sizeofcmds} -> {new_size}")
        return True

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <binary> <@rpath/dylib_name>")
        sys.exit(1)
    add_dylib(sys.argv[1], sys.argv[2])