#!/usr/bin/env python3
import os
import sys
from macholib.MachO import MachO

class MachOExtender(MachO):
    """
    Adds a method which adjusts the Mach-O header to make it appear as
    if arbitrary data added to the end of the binary is a string within
    the SYMTAB section of the LINKEDIT segment.
    """

    def __init__(self, binary_file):
        MachO.__init__(self, binary_file)
        self.binary_file = binary_file
        self.filesize = os.path.getsize(binary_file)

    def update_header(self):
        # A fat binary has a header for each architecture, but we can only
        # operate on one architecture at a time.  Fat binaries must be split
        # appart with the lipo tool.
        if len(self.headers) > 1:
            raise RuntimeError(
                'This tool only works on single-architecture binaries')
        header = self.headers[0]
        # The Mach-O header is followed by a sequence of load commands, some of
        # which define segments.  For a 64-bit macOS system, the commands which
        # define a segment have command name LC_SEGMENT_64. The LINKEDIT segment
        # is always the last segment.  Our goal is to expand the LINKEDIT segment
        # to include the extra data appended to the mach-O binary.
        for command in header.commands:
            if command[0].get_cmd_name() == 'LC_SEGMENT_64':
                if command[1].segname.startswith(b'__LINKEDIT'):
                    linkedit_cmd = command
        linkedit_info = linkedit_cmd[1]
        data_size = self.filesize - linkedit_info.fileoff - linkedit_info.filesize
        assert data_size > 0 , "You must append your extra data first!"
        linkedit_info.filesize += data_size
        # Apple says that he virtual memory size of the segment must be rounded
        # up to a multiple of the page size.  But we round to a multiple of
        # 16384 like Apple does.
        linkedit_info.vmsize = (linkedit_info.filesize + 0x3fff) & 0x4000
        # Search for the LC_SYMTAB section.  The extra data is made to appear as
        # if it were a long string at the end of the string table within the
        # LC_SYMTAB section.  So we need to adjust the size of that section.
        for command in header.commands[:]:
            if command[0].get_cmd_name() == 'LC_SYMTAB':
                symtab_info = command[1]
                break
        symtab_info.strsize += data_size
        # Save the patched mach-O header.
        with open(self.binary_file, 'rb+') as the_binary:
            self.write(the_binary)

if __name__ == '__main__':
    if len(sys.argv) != 2 or not os.path.exists(sys.argv[1]):
        print("Usage: fix_macho path_to_binary")
        sys.exit(1)
    extender = MachOExtender(sys.argv[1])
    extender.update_header()
