def add_enums(function):
""" Add standard enums from parsed MSDN documentation for all imported
library calls and their arguments.
Arguments:
function -- function object
"""
enum_count = 0
for argument in function.arguments:
# Add standard enums
if not argument.enums:
g_logger.debug(' No standard constants available for %s' %
argument.name)
else:
for enum in argument.enums:
g_logger.debug(' Importing enum %s for argument %s' %
(enum, argument.name))
if idc.import_type(-1, enum) != idaapi.BADADDR:
g_logger.debug(' ' + enum + ' ' +
hex(idc.get_enum(enum)) +
' added successfully')
enum_count = enum_count + 1
else:
g_logger.debug(' Could not add ' + enum)
if not argument.constants:
# No constants for this argument
continue
argument.name = argument.name
function.name = function.name
# Add constant descriptions
for constant in argument.constants:
if constant.name == 'NULL':
# Create unique name, so we can add descriptive comment to it
constant.name = 'NULL_{}_{}'.format(argument.name,
function.name)
# Add custom enum for NULL values if it does not exist yet
enumid = idc.get_enum(NULL_ENUM_NAME)
if enumid == idaapi.BADADDR:
enumid = idc.add_enum(-1, NULL_ENUM_NAME,
idaapi.hex_flag())
idc.add_enum_member(enumid, constant.name, 0, -1)
constid = idc.get_enum_member_by_name(constant.name)
idc.set_enum_member_cmt(constid,
format_comment(constant.description),
False)
else:
constid = idc.get_enum_member_by_name(constant.name)
if constid:
if idc.set_enum_member_cmt(
constid, format_comment(constant.description),
False):
g_logger.debug(' Description added for %s' %
constant.name)
else:
g_logger.debug(' No description added for %s' %
constant.name)
return enum_count
def set_enum_member(name, member, value):
ok = idc.add_enum_member(idc.get_enum(str(name)), str(member),
int(value), -1)
if not ok:
if idaapi.get_enum_member_by_name(str(member)) == idaapi.BADADDR:
print("Could not add enum member {member} at {name}".format(
name=name, member=member))
def markStructFromToolkit(func_ea, structOff, structName):
f = Function.Function(func_ea)
toolkitMovs = FuncAnalyzer.traceRegVar(f.func_ea, f.func_ea + f.getSize(withPool=False), 10, 0)
for trace_ea in toolkitMovs:
# lock into the register in question
insn = Instruction.Insn(trace_ea)
if insn.isComputationalInsn():
regWrites = FuncAnalyzer.traceRegWrites(f.func_ea, f.func_ea + f.getSize(withPool=False),
insn.ops[0].reg)
# TODO: traceRegWrites messes up insn
insn = Instruction.Insn(trace_ea)
# find the relevent register variable to trace
regVarIdx = FuncAnalyzer.getRegWriteIndex(trace_ea, regWrites)
insn = Instruction.Insn(trace_ea)
# trace all reads to this register variable in particular
# print('args', f.func_ea, f.func_ea + f.getSize(withPool=False),
# insn.ops[0].reg, regVarIdx)
toolkitAccesses = FuncAnalyzer.traceRegVar(f.func_ea, f.func_ea + f.getSize(withPool=False),
insn.ops[0].reg, regVarIdx)
# those now will represent all registers R10 moved to, and thus, all of those accesses
# are accesses to the R10 struct
for access in toolkitAccesses:
# now mark with Toolkit enum
accessInsn = Instruction.Insn(access)
if accessInsn.ops[1].type == ida_ua.o_displ:
if accessInsn.ops[1].addr == structOff:
# print(hex(access), idc.GetDisasm(access))
structAcccesses = FuncAnalyzer.traceRegVar(f.func_ea, f.func_ea + f.getSize(withPool=False),
accessInsn.ops[0].reg, regVarIdx+1)
for structAccess in structAcccesses:
# print(hex(structAccess), idc.GetDisasm(structAccess))
idc.op_enum(structAccess, 1, idc.get_enum(structName), 0)
return True
def rename_constant(arg_ea, fct_name, arg_name, arg_enums):
""" Rename constants to values from standard enumerations. """
instruction = idc.print_insn_mnem(arg_ea)
if instruction == 'push':
op_num = 0
elif instruction == 'mov':
op_num = 1
else:
raise RenamingException('Constant: unhandled instruction ' +
instruction)
op_val = idc.get_operand_value(arg_ea, op_num)
# NULL
if op_val == 0:
targetid = idc.get_enum_member_by_name('NULL_{}_{}'.format(
arg_name, fct_name))
serial = 0
enumid = idc.get_enum(NULL_ENUM_NAME)
constid = idc.get_enum_member(enumid, 0, serial, -1)
while constid != idaapi.BADADDR:
if constid == targetid:
idc.op_enum(arg_ea, op_num, enumid, serial)
return
serial = serial + 1
constid = idc.get_enum_member(enumid, 0, serial, -1)
# All other constants
op_type = idc.get_operand_type(arg_ea, op_num)
if op_type == idaapi.o_imm:
# only one choice
if len(arg_enums) == 1:
enumid = idc.get_enum(arg_enums[0])
idc.op_enum(arg_ea, op_num, enumid, 0)
return
for enum in arg_enums:
enumid = idc.get_enum(enum)
constid = get_constant_id(enumid, op_val)
if constid == idaapi.BADADDR:
# Not in this enum
continue
else:
# Found the right enum
idc.op_enum(arg_ea, op_num, enumid, 0)
return
def is_in_ida(symbol_type, symbol_name):
if symbol_type == "Enumeration":
return idc.get_enum(
IDAtools.clean_name(symbol_name)) != idaapi.BADADDR
else:
if str(IDAtools.clean_name(symbol_name)) in IDAtools.MAPPED_IDS:
return True
else:
IDAtools.MAPPED_IDS[str(
IDAtools.clean_name(symbol_name))] = "1"
return False
def get_enum_const_name(enum_name, const_value):
""" Name of the constant value of the enum or empty string if enum or const does not exist """
if not enum_name:
return ""
enum_id = idc.get_enum(enum_name)
if enum_id == BADADDR:
log.warn("Enum not found %s", enum_name)
return ""
const_id = idc.get_enum_member(enum_id, const_value, 0, ida_enum.DEFMASK)
if const_id == BADADDR:
log.warn("Enum const not found %s, 0x%X", enum_name, const_value)
return ""
return idc.get_enum_member_name(const_id)
def readStructMacro(path):
"""
Parses struct macros and updates a corresponding enum with their values
:param path: the path to the file containing the macros
:return:
"""
# parse macro file
macroFile = open(path)
members = []
structName = ''
for line in macroFile.readlines():
if line.lstrip().startswith('\struct_entry'):
if ', ' in line:
name = line[line.index(')')+1 : line.index(',')]
size = line[line.index(', ')+2 :].rstrip()
if '//' in size:
size = size[:size.index('//')].rstrip()
if size.startswith('0x'):
size = int(size, 16)
else:
size = int(size)
else:
name = line[line.index(')')+1 :].rstrip()
if '//' in name:
name = name[:name.index('//')].rstrip()
size = 0
members.append((name, size))
if line.startswith('def_struct_offsets'):
structName = line[line.index(', ')+2:].rstrip()
print('parsed struct "' + structName + '"')
# read into enum
enumId = idc.get_enum(structName)
if enumId == idaapi.BADADDR:
enumId = idc.add_enum(idaapi.BADADDR, structName, idaapi.decflag())
# parse all enum members, needed to know what member to replace
offset = 0x00
for member, size in members:
enumMember = idc.get_enum_member(enumId, offset, 0, ida_enum.DEFMASK)
if enumMember == idaapi.BADADDR:
print("why???")
idc.add_enum_member(enumId, structName + member, offset, idaapi.BADADDR)
elif idc.get_enum_member_name(enumMember) != structName + member:
# update member name, if value already exists
print('\tupdate %s // 0x%X' % (structName + member, offset))
idc.set_enum_member_name(enumMember, structName + member)
offset += size
return True