call-table-enum.py

import ida_segment
import ida_bytes
import idautils
import idaapi
import idc
import re

__author__     = 'zznop'
__copyright__  = 'Copyright 2019, zznop'
__email__      = '[email protected]'

def get_call_table_instrs():
    '''Return all jsr instructions with pc in the operand
    '''
    instrs = []

    # Iterate instructions in functions
    for funcea in idautils.Functions():
        for (startea, endea) in idautils.Chunks(funcea):
            for head in idautils.Heads(startea, endea):
                instr = idc.GetDisasm(head).split()
                if instr[0] == 'jsr' or instr[0] == 'jmp':
                    if 'pc' in instr[1]:
                        instrs.append(instr)

    # Iterate instructions not in a function
    addr = idaapi.find_not_func(0, 1)
    while addr != idc.BADADDR:
        instr = idc.GetDisasm(addr).split()
        if instr[0] == 'jsr' or instr[0] == 'jmp':
            if 'pc' in instr[1]:
                instrs.append(instr)
        addr = idaapi.find_not_func(addr, 1)

    return instrs

def get_call_table_addrs_from_instrs(instrs):
    '''Get the base addr of the call tables from the instrs
    '''
    call_table_addrs = []
    for mnem, opnd in instrs:
        print(opnd)
        match = re.match(r"^.*_(.*)\(pc,.*\)$", opnd)
        if match:
            addr = int(match.group(1), 16)
            call_table_addrs.append(addr)
    return call_table_addrs

def disas_call_tables(base_table_addrs):
    for base_addr in base_table_addrs:
        i = base_addr
        while True:
            instr_dword = ida_bytes.get_32bit(i)
            if not ((instr_dword >> 24) == 0x60):
                break
            idc.MakeCode(i)
            i += 4

def main():
    '''
    '''
    call_table_instrs = get_call_table_instrs()
    if call_table_instrs is []:
        print("Failed to find JSR instructions that modify PC")
        return

    call_table_addrs = get_call_table_addrs_from_instrs(call_table_instrs)
    if call_table_addrs is []:
        print("Failed to enumerate call table base addresses from instructions")
        return

    disas_call_tables(call_table_addrs)

if __name__ == '__main__':
    main()