Пример #1
0
def main():
    if _IN_IDA:
        # # get dyld_shared_cache path from IDA's openFile dialog
        print "[+] Please choose the original dyld_shared_cache_arm64"
        dsc_path = idc.AskFile(0, "*.*", "dyld shared cache file")
    else:
        dsc_path = sys.argv[1]

    if not dsc_path or not os.path.exists(dsc_path):
        raise RuntimeError("Couldn't find the dyld shared cache file..")

    print "[+] about to parse %s.." % (dsc_path)
    dsc_file = open(dsc_path, "rb")
    adrfind = AddrFinder(dsc_file, cache_symbols=False)
    map_shared_bridges(dsc_file, adrfind)
    if _IN_IDA:
        addresses = sorted(set(get_bad_addresses()))
    else:
        addresses = sorted(set(eval(open("addrs.txt", "rb").read())))

    segments, exports = get_segments_and_exports_for_addresses(
        addresses, adrfind)
    # segments = join_neighbors(segments, threshold=0x1000)
    if _IN_IDA:
        map_segments(segments, dsc_file)
        map_exports(exports)
        idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
Пример #2
0
    def _scanSegment(self, buffer, segment_start_address):
        offset = 0
        reanalyze_segment = False

        while offset < len(buffer):
            for pattern in self._pattern_list:
                if pattern.compareTo(buffer, offset) == False:
                    continue

                patch_address = segment_start_address + offset
                print "Patching following address: 0x%x" % patch_address
                self._patchDatabase(pattern.getPatch(), patch_address)

                offset += len(pattern.getPatch()) - 1
                reanalyze_segment = True

                break

            offset += 1

        if reanalyze_segment:
            print("Reanalyzing segment...")
            idaapi.analyze_area(segment_start_address,
                                segment_start_address + len(buffer))

        return
Пример #3
0
    def origFun(self, ip):
        print('look for fun having ip 0x%x' % ip)
        fun = self.getFun(ip)
        if fun is None:
            simicsString = gdbProt.Evalx(
                'SendGDBMonitor("@cgc.getSO(0x%x)");' % ip)
            print('No function found.  Check load for: %s' % simicsString)
            if ':' in simicsString:
                sofile, start_end = str(simicsString).rsplit(':', 1)
                print('sofile is %s start_end is %s' % (sofile, start_end))
                if '-' not in start_end:
                    print('Bad response from getSO: %s' % simicsString)
                    return
                full = os.path.join(self.root_prefix, sofile[1:])
                sopath = self.getFunPath(full)
                start, end = start_end.split('-')
                start = int(start, 16)
                end = int(end, 16)
                self.add(sopath, start)
                fun = self.getFun(ip)
                print('start 0x%x end 0x%x' % (start, end))
                idaapi.analyze_area(start, end)
                for fun in sorted(self.funs):
                    if fun >= start and fun <= end:
                        name = str(self.funs[fun]['name'])
                        nea = idaapi.get_name_ea(idaapi.BADADDR, name)
                        if nea != idaapi.BADADDR:
                            name = name + '_so'
                        idc.MakeName(int(fun), name)
                        print('made name for 0x%x  %s' % (int(fun), name))
                for fun in self.funs:
                    if fun >= start and fun <= end:
                        #print('fun 0x%x name <%s>' % (fun, name))
                        idc.MakeFunction(fun, idaapi.BADADDR)

        elif fun is not None:
            print('Do one fun 0x%x' % fun)
            for i in range(self.funs[fun]['start'], self.funs[fun]['end']):
                idc.MakeUnkn(i, 1)
            idaapi.auto_mark_range(self.funs[fun]['start'],
                                   self.funs[fun]['end'], 25)
            idaapi.autoWait()
            return fun
        return None
Пример #4
0
def resizeRegion(analyzer, start_ea, end_ea, new_start_ea, new_end_ea):
    """Resize a given code region, according to the new dimensions.

    Args:
        analyzer (instance): analyzer instance to be used
        start_ea (int): effective start address of the original region
        end_ea (int): effective end address of the original region
        new_start_ea (int): effective start address for the new region
        new_end_ea (int): effective end address for the new region
    """
    analyzer.logger.info(
        "Resizing code region of type %d: 0x%x (0x%x) - 0x%x (0x%x)",
        analyzer.codeType(start_ea), new_start_ea, start_ea, end_ea,
        new_end_ea)
    code_type_before = analyzer.codeType(min(start_ea, new_start_ea) - 1)
    code_type_middle = analyzer.codeType(start_ea)
    code_type_after = analyzer.codeType(max(end_ea, new_end_ea))
    # Make sure it will be treated as code
    fix_regions = []
    if new_start_ea < start_ea:
        fix_regions.append((new_start_ea, start_ea))
    elif new_start_ea != start_ea:
        fix_regions.append((start_ea, new_start_ea))
    if end_ea < new_end_ea:
        fix_regions.append((end_ea, new_end_ea))
    elif end_ea != new_end_ea:
        fix_regions.append((new_end_ea, end_ea))
    # Make the changed parts unknown, before re-analyzing them
    for region_start, region_end in fix_regions:
        idc.MakeUnknown(region_start, region_end - region_start, 0)
    # manually set the wanted value over the entire region
    if start_ea < new_start_ea:
        analyzer.setCodeType(start_ea, new_start_ea, code_type_before)
    elif start_ea != new_start_ea:
        analyzer.setCodeType(new_start_ea, start_ea, code_type_middle)
    if end_ea < new_end_ea:
        analyzer.setCodeType(end_ea, new_end_ea, code_type_middle)
    elif end_ea != new_end_ea:
        analyzer.setCodeType(new_end_ea, end_ea, code_type_after)
    # now reanalyze the new section
    for region_start, region_end in fix_regions:
        idaapi.analyze_area(region_start, region_end)
Пример #5
0
def convertRegion(analyzer, start_ea, end_ea):
    """Convert (Cancel) a given code region (change it's code type).

    Args:
        analyzer (instance): analyzer instance to be used
        start_ea (int): effective start address of the region
        end_ea (int): effective end address of the region
    """
    wanted_code_type = analyzer.codeType(end_ea)
    analyzer.logger.info(
        "Converting code region of type %d to %d: 0x%x - 0x%x (%d)",
        analyzer.codeType(start_ea), wanted_code_type, start_ea, end_ea,
        end_ea - start_ea)
    # Make sure it will be treated as code
    idc.MakeUnknown(start_ea, end_ea - start_ea, 0)
    # manually set the wanted value over the entire region
    analyzer.setCodeType(start_ea, end_ea, wanted_code_type)
    # now reanalyze the new section
    idaapi.analyze_area(
        analyzer.alignTransitionAddress(start_ea, wanted_code_type), end_ea)
Пример #6
0
def cleanStart(analyzer, scs, undef=False):
    """Clean the selected code segments, and re-analyzer them using the gathered metadata until now.

    Args:
        analyzer (instance): analyzer instance to be used
        scs (list): list of (sark) code segments to work on
        undef (bool, optional): True iff should undefine the code segments (False by default)
    """
    for sc in scs:
        if undef:
            analyzer.logger.info("Undefining code segment: 0x%x - 0x%x",
                                 sc.startEA, sc.endEA)
            sark.data.undefine(sc.startEA, sc.endEA)
            analyzer.logger.info(
                "Marking all known switch tables in the segment")
            analyzer.switch_identifier.markSwitchTables(sc)
    analyzer.logger.info("Marking all known fptr functions")
    analyzer.fptr_identifier.makePointedFunctions()
    for sc in scs:
        analyzer.logger.info("Re-Analyzing code segment: 0x%x - 0x%x",
                             sc.startEA, sc.endEA)
        idaapi.analyze_area(sc.startEA, sc.endEA)
        idaapi.autoWait()
Пример #7
0
            idc.MakeCode(seg_ptr)
            idc.MakeFunction(seg_ptr, seg_ptr + 4)
            idc.MakeRptCmt(seg_ptr, new_name)
            offset = seg_ptr - ea
            seg_ptr += 4

        dw = offset - 5
        idc.PatchByte(ea, 0xE8)
        idc.PatchDword(ea + 1, dw)


def make_offsets(segname):
    segea = idc.SegByBase(idc.SegByName(segname))
    segend = idc.SegEnd(segea)

    while segea < segend:
        idc.OpOffset(segea, 0)
        ptr = idc.Dword(segea)
        idc.OpOffset(ptr, 0)
        segea += 4


if __name__ == '__main__':
    make_offsets('__cls_refs')
    make_offsets('__message_refs')
    idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
    fix_callgraph(idc.LocByName('_objc_msgSend'), 'msgSend', 4, 4)
    fix_callgraph(idc.LocByName('_objc_msgSendSuper'), 'msgSendSuper', 4, 4)
    idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
    print 'Done.'
Пример #8
0
                if (xr2.type == 19) and (xr2.to == next_next):

                    db = head + idaapi.get_item_size(head)

                    if idc.Byte(head) == 0x0F:
                        idaapi.patch_byte(head, 0x90)
                        idaapi.patch_byte(head + 1, 0xE9)
                    else:
                        idaapi.patch_byte(head, 0xEB)

                    idc.MakeUnknown(db, xr.to - db + 0x10, idaapi.DOUNK_SIMPLE)
                    idc.MakeCode(xr.to)

                    i = db
                    while i < xr.to:
                        if (i + 4) < xr.to:
                            idc.MakeDword(i)
                            i += 4
                        else:
                            idc.MakeByte(i)
                            i += 1

                    idaapi.analyze_area(head - 0x40, head + 0x40)
                    idaapi.analyze_area(xr.to - 0x40, xr.to + 0x40)

for head in idautils.Heads():
    if idc.Byte(head) == 0xE8:
        for xr in idautils.XrefsFrom(head, 0):
            # Find direct call targets
            if not (xr.type == 21):
                idc.MakeFunction(xr.to)
Пример #9
0
def load_file(li, neflags, format):
    
    """
    Load the file into database

    @param li: a file-like object which can be used to access the input data
    @param neflags: options selected by the user, see loader.hpp
    @return: 0-failure, 1-ok
    """

    if format == _3DSX_FORMAT_NAME:
    
        idaapi.set_processor_type("arm", SETPROC_ALL|SETPROC_FATAL)
        
        li.seek(0)

        (magic, header_size, reloc_header_size, format_ver, flags, code_seg_size, rodata_seg_size, data_seg_size, bss_seg_size) = struct.unpack("<IHHIIIIII", li.read(4*8))
        
        #print(hex(header_size))
        #print(hex(reloc_header_size))
        #print(hex(code_seg_size))
        #print(hex(rodata_seg_size))
        #print(hex(data_seg_size))
        #print(hex(bss_seg_size))
        
        #CODE SEGMENT
        seg1 = CN_3DSX_LOADADR
        seg1_size = (code_seg_size + 0xFFF) &(~0xFFF)
        file_offset = header_size + reloc_header_size*3
        
        AddSeg(seg1, seg1 + seg1_size, 0, 1, idaapi.saRelPara, idaapi.scPub)
        SetSegmentType(seg1, idaapi.SEG_CODE)
        RenameSeg(seg1, "CODE")        
        li.file2base(file_offset, seg1, seg1 + code_seg_size, 0)

        
        #RO_DATA SEGMENT
        seg2 = seg1 + seg1_size
        seg2_size = (rodata_seg_size + 0xFFF) &(~0xFFF)
        file_offset += code_seg_size
        
        AddSeg(seg2, seg2 + seg2_size, 0, 1, idaapi.saRelPara, idaapi.scPub)
        SetSegmentType(seg2, idaapi.SEG_DATA)
        RenameSeg(seg2, "RODATA")
        li.file2base(file_offset, seg2, seg2 + rodata_seg_size, 0)
        
        #DATA SEGMENT
        seg3 = seg2 + seg2_size
        seg3_size = (data_seg_size + 0xFFF) &(~0xFFF)
        file_offset += rodata_seg_size
        
        AddSeg(seg3, seg3 + seg3_size, 0, 1, idaapi.saRelPara, idaapi.scPub)
        SetSegmentType(seg3, idaapi.SEG_DATA)
        RenameSeg(seg3, "DATA")
        li.file2base(file_offset, seg3, seg3 + (data_seg_size - bss_seg_size), 0)
        
        #relocations
        relocs_ptr = file_offset + (data_seg_size - bss_seg_size)
        
        segments = [seg1, seg2, seg3]
        
        for rel_table in range(3):
        
            li.seek(header_size + rel_table * 8)
            (abs_count, rel_count) = struct.unpack("<II", li.read(4*2))
        
            li.seek(relocs_ptr)
            relocs_ptr = relocs_ptr + (abs_count * 4 + rel_count * 4)

            pos = segments[rel_table]
            
            #absolute relocations
            for i in range(abs_count):
                (skip, patches) = struct.unpack("<HH", li.read(2*2))
                
                pos += skip * 4
                for x in range(patches):
                    
                    addr = Dword(pos)
                    if addr < seg1_size:
                        addr = seg1 + addr
                    elif addr < (seg1_size + seg2_size):
                        addr = seg2 + addr - seg1_size
                    else:
                        addr = seg3 + addr - (seg1_size + seg2_size)

                    PatchDword(pos, addr)
                    SetFixup(pos, idaapi.FIXUP_OFF32 or idaapi.FIXUP_CREATED, 0, addr, 0)
                    pos += 4
                    
            # cross-segment relative relocations
            for i in range(rel_count):
                (skip, patches) = struct.unpack("<HH", li.read(2*2))
                
                pos += skip * 4
                for x in range(patches):
                    
                    addr = Dword(pos)
                    if addr < seg1_size:
                        addr = seg1 + addr
                    elif addr < (seg1_size + seg2_size):
                        addr = seg2 + addr - seg1_size
                    else:
                        addr = seg3 + addr - (seg1_size + seg2_size)
                    
                    PatchDword(pos, addr - pos)
                    SetFixup(pos, idaapi.FIXUP_OFF32 or idaapi.FIXUP_CREATED, 0, addr - pos, 0)
                    pos += 4                    
        
        idaapi.add_entry(CN_3DSX_LOADADR, CN_3DSX_LOADADR, "start", 1)
        idaapi.analyze_area(seg1, seg1 + seg1_size)
        
        print "Load OK"
        return 1
    
    return 0
Пример #10
0
            idc.PatchWord(ea + 2, w2)
        else:
            if offset > 0 and offset & 0xFF000000:
                print 'Offset too far (%08x) Stopping [%08x]' % (offset, ea)
            dw = (0xFA000000 | (offset - 8 >> 2))
            if dw < 0:
                dw = dw & 0xFAFFFFFF
            idc.PatchDword(ea, dw)


def make_offsets(segname):
    segea = idc.SegByBase(idc.SegByName(segname))
    segend = idc.SegEnd(segea)

    while segea < segend:
        idc.OpOffset(segea, 0)
        ptr = idc.Dword(segea)
        idc.OpOffset(ptr, 0)
        segea += 4

if __name__ == '__main__':
    print 'Preparing class references segments'
    make_offsets('__objc_classrefs')
    make_offsets('__objc_superrefs')
    idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
    print 'Fixing callgraph'
    fix_callgraph(idc.LocByName('_objc_msgSend'), 'msgSend', 0, 1)
    fix_callgraph(idc.LocByName('_objc_msgSendSuper2'), 'msgSendSuper', 3, 1)
    idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
    print 'Done.'
Пример #11
0
 def fun(begin_addr, end_addr):
     idaapi.analyze_area(begin_addr, end_addr)
Пример #12
0
    end = func_data.end_ea
    # print(func_name, hex(start), hex(end))
    if "datadiv_decode" in func_name and not ("j_.datadiv_decode"
                                              in func_name):
        sim.emu_start(start, end)

sim.patch_segment('data')

for seg in sim.segments:
    if "data" in seg['name']:
        # 把data段全部undefined
        print("MakeUnknown %s" % seg['name'])
        idc.MakeUnknown(seg['start'], seg['end'] - seg['start'],
                        idaapi.DELIT_DELNAMES)
        # 调用ida重新解析data段
        print("analyze area: 0x%x - 0x%x" % (seg['start'], seg['end']))
        idaapi.analyze_area(seg['start'], seg['end'])
        # idaapi.clear_strlist()
        # idaapi.build_strlist()

# 查询string的交叉引用,在引用位置添加备注
s = idautils.Strings(False)
s.setup()
for i, str_info in enumerate(s):
    if str_info:
        # print("%x: len=%d  index=%d-> '%s'" % (str_info.ea, str_info.length, i, str(str_info)))
        str_cont = str(str_info)
        refs = idautils.DataRefsTo(str_info.ea)
        for ref in refs:
            idc.MakeComm(ref, str_cont)
                    # Undefine and redefine code borders
                    idc.MakeUnknown(db, xref.to - db + 0x10, idaapi.DOUNK_SIMPLE)
                    idc.MakeCode(xref.to)

                    # Convert the bad code into bytes
                    i = db
                    while i < xref.to:
                        if (i+4) < xref.to:
                            idc.MakeDword(i)
                            i += 4
                        else:
                            idc.MakeByte(i)
                            i += 1

                    # Analyze the area to fix xrefs
                    idaapi.analyze_area(head-0x40, head+0x40)
                    idaapi.analyze_area(xref.to-0x40, xref.to+0x40)

                    bad_jump_count += 1

print "Fixed %d anti-disassembly instructions." % bad_jump_count

# Make sure that all calls go to functions
for head in idautils.Heads():
    if idc.Byte(head) == 0xE8:
        for xref in idautils.XrefsFrom(head, 0):
            # Find direct call targets
            if not (xref.type == 21):
                idc.MakeFunction(xref.to)

Пример #14
0
def dataScan(analyzer, scs):
    """Scan the code segments for orphan data blobs that represent analysis errors.

    Args:
        analyzer (instance): analyzer instance to be used
        scs (list): list of (sark) code segments
    """
    # First Scan - unreffed data chunks inside functions ==> should be converted to code
    for sc in scs:
        first_line = None
        end_line = None
        for line in sc.lines:
            # After the first, the rest of the lines should have 0 crefs
            if first_line is not None and ((not line.is_data)
                                           or len(list(line.drefs_to)) > 0
                                           or len(list(line.crefs_to)) > 0):
                end_line = line
            # we only care about data lines with a single cref from the previous line
            elif first_line is None and (
                (not line.is_data) or len(list(line.drefs_to)) > 0
                    or len(list(line.crefs_to)) != 1
                    or sark.Line(list(line.crefs_to)[0]).next != line):
                end_line = line
            # don't mark switch entries
            elif analyzer.switch_identifier.isSwitchEntry(line.startEA):
                end_line = line
            # Finally, check if it could be a function of some type
            elif first_line is None:
                first_line = line
                continue
            # Found an adjacent suitable line
            else:
                continue
            # Now check if we found something (end_line is always != None at this point)
            if first_line is not None and end_line is not None:
                chunk_start = first_line.startEA
                chunk_end = end_line.startEA
                # check that the chunk before us is not the end of a function
                if analyzer.func_classifier.predictFunctionEnd(chunk_start):
                    # shouldn't really happen, do nothing in this case
                    pass
                # data chunk in the middle of a function, and not at it's end - convert it to code
                else:
                    analyzer.logger.debug(
                        "In-Function data chunk at: 0x%x - 0x%x (%d)",
                        chunk_start, chunk_end, chunk_end - chunk_start)
                    idc.MakeUnknown(chunk_start, chunk_end - chunk_start, 0)
                    idc.MakeCode(chunk_start)
                # reset the vars
                first_line = None
                end_line = None

    # Second scan - unreffed data chunks outside of functions ==> new functions, possibly of different code type
    size_limit = analyzer.func_classifier.functionStartSize()
    analyzer.logger.debug("Size limit for data scan is: %d", size_limit)
    conversion_candidates = []
    # recon pass
    for sc in scs:
        first_line = None
        end_line = None
        for line in sc.lines:
            # we only care about data lines without xrefs
            if (not line.is_data) or len(list(line.crefs_to)) > 0 or len(
                    list(line.drefs_to)) > 0:
                end_line = line
            # check if it's big enough for the classifier
            elif line.size < size_limit:
                end_line = line
            # check if it looks like a string
            elif analyzer.str_identifier.isLocalAsciiString(line.startEA,
                                                            check_refs=False):
                analyzer.str_identifier.defineAsciiString(line.startEA)
                end_line = line
            # make sure it isn't a switch entry
            elif analyzer.switch_identifier.isSwitchEntry(line.startEA):
                end_line = line
            # Finally, check if it could be a function of some type
            elif first_line is None:
                first_line = line
                continue
            # Found an adjacent suitable line
            else:
                continue
            # Now check if we found something (end_line is always != None at this point)
            if first_line is not None and end_line is not None:
                chunk_start = first_line.startEA
                chunk_end = end_line.startEA
                guess_code_type = analyzer.func_classifier.predictFunctionStartType(
                    chunk_start)
                original_code_type = analyzer.codeType(chunk_start)
                analyzer.logger.debug(
                    "Found a data chunk at: 0x%x - 0x%x (%d), (Type %d, Local type %d)",
                    chunk_start, chunk_end, chunk_end - chunk_start,
                    guess_code_type, original_code_type)
                # Check if this is the beginning of a function
                if analyzer.func_classifier.predictFunctionStart(
                        chunk_start, guess_code_type):
                    conversion_candidates.append(
                        (chunk_start, chunk_end, guess_code_type,
                         original_code_type))
                # reset the vars
                first_line = None
                end_line = None
    # conversion pass
    for chunk_start, chunk_end, guess_code_type, original_code_type in conversion_candidates:
        analyzer.logger.info(
            "Found an isolated data chunk at: 0x%x - 0x%x (%d), (Type %d, Local type %d)",
            chunk_start, chunk_end, chunk_end - chunk_start, guess_code_type,
            original_code_type)
        idc.MakeUnknown(chunk_start, chunk_end - chunk_start, 0)
        if original_code_type != guess_code_type:
            analyzer.setCodeType(chunk_start, chunk_end, guess_code_type)
        idaapi.analyze_area(chunk_start, chunk_end)
        idc.MakeFunction(chunk_start)
def trace_param(ea, min_ea, op_type, op_val):
    '''
    trace_param: ea, min_ea, op_type, op_val

    Taking ea as start, this function does basic backtrace of
    an operand (defined by op_type and op_val) until it finds
    a data reference which we consider the "source". It stops
    when ea < min_ea (usually the function start).

    It does not support arithmetic or complex modifications of
    the source. This will be improved on future versions.
    '''
    global displ_re, msgsend, var_re

    ea_call = ea
    while ea != idc.BADADDR and ea != min_ea:
        ea = idc.PrevHead(ea, min_ea)

        if op_type == idaapi.o_reg and op_val == 0 and idaapi.is_call_insn(ea):
            # We have a BL/BLX that will modify the R0
            # we're tracking
            #TODO: resolve more situation
            return None

        if idc.GetMnem(ea) in ['LDR', 'MOV']:
            src_op = 1
            dest_op = 0
        elif idc.GetMnem(ea) == 'STR':
            src_op = 0
            dest_op = 1
        else:
            continue

        if idc.GetOpType(ea, dest_op) == op_type and idc.GetOperandValue(ea, dest_op) == op_val:
            # Found, see where it comes from
            if idc.GetOpType(ea, src_op) == idc.o_mem or idc.GetOpType(ea, src_op) == idc.o_imm: #add o_imm support
                # Got the final reference
                refs = list(idautils.DataRefsFrom(ea))
                if not refs:
                    local_ref = idc.GetOperandValue(ea, src_op)
                    far_ref = idc.Dword(local_ref)
                else:
                    while len(refs) > 0:
                        far_ref = refs[0]
                        refs = list(idautils.DataRefsFrom(refs[0]))
                #patch by lc
                if far_ref:

                return far_ref
            elif idc.GetOpType(ea, src_op) == idc.o_displ:
                if ', [SP' in idc.GetDisasm(ea):
                    if 'arg_' in idc.GetDisasm(ea):
                        # We don't track function arguments
                        return None

                    # We're tracking an stack variable
                    try:
                        var_name = var_re.search(idc.GetDisasm(ea)).group('varname')
                    except:
                        print '%08x: Unable to recognize variable' % ea
                        return None

                    while ea != idc.BADADDR and ea > min_ea:
                        if idc.GetMnem(ea) == 'STR' and var_name in idc.GetDisasm(ea):
                            # New reg to track
                            op_val = idc.GetOperandValue(ea, dest_op)
                            break
                        ea = idc.PrevHead(ea, min_ea)
                else:
                    # New reg to track
                    if '[LR]' in idc.GetDisasm(ea):
                        # Optimizations use LR as general reg
                        op_val = 14
                    else:
                        try:
                            op_val = int(displ_re.search(idc.GetDisasm(ea)).group('regnum'))
                        except:
                            print '%08x: Unable to recognize register' % ea
                            return None
            elif idc.GetOpType(ea, src_op) == idc.o_reg:
                # Direct reg-reg assignment
                op_val = idc.GetOperandValue(ea, src_op)
            else:
                # We don't track o_phrase or other complex source operands :(
                return None
    #register R0-R3 assigned by function parameter
    if ea <= min_ea and op_type == idc.o_reg and op_val in range(4):
        f_info = get_func_info(ea)
        return ['pself', 'selector', f_info['fparam_type'], f_info['sparam_name']][op_val]#fix: error
    return None



def fix_callgraph(msgsend, segname, class_param, sel_param): #class_param == 0, sel_param == 1
    '''
    fix_callgraph: msgsend, segname, class_param, sel_param

    Given the msgsend flavour address as a parameter, looks
    for the parameters (class and selector, identified by
    class_param and sel_param) and creates a new segment where
    it places a set of dummy calls named as classname_methodname
    (we use method instead of selector most of the time).
    '''

    t1 = time.time()
    if not msgsend:
        print 'ERROR: msgSend not found'
        return

    total = 0
    resolved = 0
    call_table = dict()

    for xref in idautils.XrefsTo(msgsend, idaapi.XREF_ALL):
        total += 1
        ea_call = xref.frm
        func_start = idc.GetFunctionAttr(ea_call, idc.FUNCATTR_START)
        if not func_start or func_start == idc.BADADDR:
            continue
        ea = ea_call
        method_name_ea = trace_param(ea, func_start, idc.o_reg, sel_param)#sel_param == 1
        if method_name_ea and idc.isASCII(idc.GetFlags(method_name_ea)):
            method_name = idc.GetString(method_name_ea, -1, idc.ASCSTR_C)
            if not method_name:
                method_name = '_unk_method'
        else:
            method_name = '_unk_method'

        class_name_ea = trace_param(ea, func_start, idc.o_reg, class_param)#class_param == 0
        if class_name_ea:
            class_name = idc.Name(class_name_ea)
            if not class_name:
                class_name = '_unk_class'
        else:
            class_name = '_unk_class'

        if method_name == '_unk_method' and class_name == '_unk_class':
            continue

        # Using this name convention, if the class and method
        # are identified by IDA, the patched call will point to
        # the REAL call and not one of our dummy functions
        #
        class_name = class_name.replace('_OBJC_CLASS_$_', '')
        class_name = class_name.replace('_OBJC_METACLASS_$_', '')
        new_name = '_[' + class_name + '_' + method_name + ']'
        print '%08x: %s' % (ea_call, new_name)
        call_table[ea_call] = new_name
        resolved += 1

    print '\nFinal stats:\n\t%d total calls, %d resolved' % (total, resolved)
    print '\tAnalysis took %.2f seconds' % (time.time() - t1)

    if resolved == 0:
        print 'Nothing to patch.'
        return

    print 'Adding new segment to store new nullsubs'

    # segment size = opcode ret (4 bytes) * num_calls
    seg_size = resolved * 4
    seg_start = idc.MaxEA() + 4
    idaapi.add_segm(0, seg_start, seg_start + seg_size, segname, 'CODE')

    print 'Patching database...'
    seg_ptr = seg_start
    for ea, new_name in call_table.items():
        if idc.LocByName(new_name) != idc.BADADDR:
            offset = idc.LocByName(new_name) - ea
        else:
            # create code and name it
            idc.PatchDword(seg_ptr, 0xE12FFF1E) # BX LR
            idc.MakeName(seg_ptr, new_name)
            idc.MakeCode(seg_ptr)
            idc.MakeFunction(seg_ptr, seg_ptr + 4)
            idc.MakeRptCmt(seg_ptr, new_name)
            offset = seg_ptr - ea
            seg_ptr += 4

        # patch the msgsend call
        if idc.GetReg(ea, "T") == 1:
            if offset > 0 and offset & 0xFF800000:
                print 'Offset too far for Thumb (%08x) Stopping [%08x]' % (offset, ea)
                return

            off1 = (offset & 0x7FF000) >> 12
            off2 = (offset & 0xFFF) / 2
            w1 = (0xF000 | off1)
            w2 = (0xE800 | off2) - 1
            idc.PatchWord(ea, w1)
            idc.PatchWord(ea + 2, w2)
        else:
            if offset > 0 and offset & 0xFF000000:
                print 'Offset too far (%08x) Stopping [%08x]' % (offset, ea)
            dw = (0xFA000000 | (offset - 8 >> 2))
            if dw < 0:
                dw = dw & 0xFAFFFFFF
            idc.PatchDword(ea, dw)


def make_offsets(segname):
    '''
    change the segment's data value into offset by class name
    '''
    segea = idc.SegByBase(idc.SegByName(segname))
    segend = idc.SegEnd(segea)

    while segea < segend:
        idc.OpOffset(segea, 0)
        ptr = idc.Dword(segea)
        idc.OpOffset(ptr, 0)
        segea += 4

if __name__ == '__main__':
    print 'Preparing class references segments'
    make_offsets('__objc_classrefs') #TODO: what's these two segment means?
    make_offsets('__objc_superrefs')
    idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
    print 'Fixing callgraph'
    fix_callgraph(idc.LocByName('_objc_msgSend'), 'msgSend', 0, 1)
    fix_callgraph(idc.LocByName('_objc_msgSendSuper2'), 'msgSendSuper', 3, 1)
    idaapi.analyze_area(idc.MinEA(), idc.MaxEA())
    print 'Done.'