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.Til2Idb(-1, enum) != idaapi.BADADDR:
g_logger.debug(' ' + enum + ' ' + hex(idc.GetEnum(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.encode('utf-8')
function.name = function.name.encode('utf-8')
# Add constant descriptions
for constant in argument.constants:
constant.name = constant.name.encode('utf-8')
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.GetEnum(NULL_ENUM_NAME)
if enumid == idaapi.BADADDR:
enumid = idc.AddEnum(-1, NULL_ENUM_NAME, idaapi.hexflag())
idc.AddConstEx(enumid, constant.name, 0, -1)
constid = idc.GetConstByName(constant.name)
idc.SetConstCmt(constid, format_comment(constant.description),
False)
else:
constid = idc.GetConstByName(constant.name)
if constid:
if idc.SetConstCmt(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 yacheck_enums(self):
values = yaunit.load('enums')
for flag in flags:
for prefix, enum_width, is_bitfield, num_fields in tests:
name, ea, values = str(values[0][0]), values[0][1], values[1:]
eid = idc.GetEnum(name)
self.assertNotEqual(eid, idaapi.BADADDR)
self.assertEqual(idc.GetEnumFlag(eid), flag)
self.assertEqual(idc.GetEnumName(eid), name)
self.assertEqual(idc.IsBitfield(eid), is_bitfield)
self.assertEqual(idc.GetEnumCmt(eid, False), prefix + 'cmt')
self.assertEqual(idc.GetEnumCmt(eid, True), prefix + 'rpt')
if enum_width != 0:
self.assertEqual(idc.GetEnumWidth(eid), enum_width)
n = 0
for value, bmask in walk_enum(eid):
self.assertLessEqual(n, num_fields)
v = 1 << n if is_bitfield else n
self.assertEqual(value, value, v)
cid = idc.GetConstEx(eid, v, 0, bmask)
self.assertNotEqual(cid, idaapi.BADADDR)
field = '%s_%d' % (name, n)
self.assertEqual(idc.GetConstName(cid), field)
# FIXME comments are not working
#self.assertEqual(idc.GetConstCmt(cid, False), field + 'cmt')
#self.assertEqual(idc.GetConstCmt(cid, True), field + 'rpt')
n += 1
self.assertEqual(n, num_fields)
if ea != None:
self.assertEqual(idaapi.get_enum_id(ea, 1)[0], eid)
def get_struc_enum_id_for_name(self, name):
item_id = idc.GetStrucIdByName(name)
if item_id == idc.BADADDR:
item_id = idc.GetEnum(name)
if item_id == idc.BADADDR:
logger.error("no struc or enum id for name : %s", name)
return None
return self.get_struc_enum_object_id(item_id, name)
def rename_constant(arg_ea, fct_name, arg_name, arg_enums):
""" Rename constants to values from standard enumerations. """
instruction = idc.GetMnem(arg_ea)
if instruction == 'push':
op_num = 0
elif instruction == 'mov':
op_num = 1
else:
raise RenamingException('Constant: unhandled instruction ' +
instruction)
op_val = idc.GetOperandValue(arg_ea, op_num)
# NULL
if op_val == 0:
targetid = idc.GetConstByName('NULL_{}_{}'.format(arg_name, fct_name))
serial = 0
enumid = idc.GetEnum(NULL_ENUM_NAME)
constid = idc.GetConstEx(enumid, 0, serial, -1)
while constid != idaapi.BADADDR:
if constid == targetid:
idc.OpEnumEx(arg_ea, op_num, enumid, serial)
return
serial = serial + 1
constid = idc.GetConstEx(enumid, 0, serial, -1)
# All other constants
op_type = idc.GetOpType(arg_ea, op_num)
if op_type == idaapi.o_imm:
# only one choice
if len(arg_enums) == 1:
enumid = idc.GetEnum(arg_enums[0])
idc.OpEnumEx(arg_ea, op_num, enumid, 0)
return
for enum in arg_enums:
enumid = idc.GetEnum(enum)
constid = get_constant_id(enumid, op_val)
if constid == idaapi.BADADDR:
# Not in this enum
continue
else:
# Found the right enum
idc.OpEnumEx(arg_ea, op_num, enumid, 0)
return
def __init__(self, name, canCreate=False):
if isinstance(name, (int, long)):
name = str(name)
self.name = name
self.id = idc.GetEnum(name)
self._vals = None
self._names = None
if self.id == idc.BADADDR:
try:
self.id = int(name)
self.name = idc.GetEnumName(self.id)
except Exception:
self.name = None
self.id = None
if not self.name:
if not canCreate:
raise Enum.EnumNotFoundError(name)
self.id = idc.AddEnum(idc.GetEnumQty(), name, 0)
def make_enum(self, object_version, address):
enum_name = object_version.get_name()
object_id = object_version.get_id()
# build flags
bitfield = False
flags = object_version.get_object_flags()
if flags & 0x1 == 0x1:
bitfield = True
flags = flags & ~0x1
try:
enum_width = object_version.get_size()
except KeyError:
enum_width = None
# check if enum already exists
enum_id = idc.GetEnum(enum_name)
enum_xrefs_ids = object_version.get_xrefed_ids()
logger.debug(
"%s:%d : Check here that enum_xrefs_ids is a set(YaToolObjectId) and correctly used"
% (__file__, inspect.currentframe().f_lineno))
if enum_id != idc.BADADDR:
# enum already exists, deleting all members
for (const_id, const_value,
bmask) in YaToolIDATools.enum_member_iterate_all(enum_id):
if bmask == idc.BADADDR:
bmask = -1
const_name = idc.GetConstName(const_id)
if (enum_name, const_name,
const_value) not in self.enum_member_ids:
idc.DelConstEx(enum_id, const_value, 0, bmask)
elif self.hash_provider.get_enum_member_id(
enum_id, enum_name, const_id, const_name,
const_value) not in enum_xrefs_ids:
logger.debug("deleting not found constant : %s/%s" %
(enum_name, const_name))
idc.DelConstEx(enum_id, const_value, 0, bmask)
"""
if the enum "bitfield" state has changed : change it!
We can't access the functions
"""
if idc.IsBitfield(enum_id) != bitfield:
idaapi.set_enum_bf(enum_id, bitfield)
else:
# create enum
# SetEnumFlag return "'module' object has no attribute 'set_enum_flag'".
enum_id = idc.AddEnum(address, enum_name, flags)
if enum_width is not None:
idc.SetEnumWidth(enum_id, enum_width)
if bitfield:
idc.SetEnumBf(enum_id, 1)
# process comment
try:
repeatable_headercomment = self.sanitize_comment_to_ascii(
object_version.get_header_comment(True))
idc.SetEnumCmt(enum_id, repeatable_headercomment, 1)
except KeyError:
pass
try:
nonrepeatable_headercomment = self.sanitize_comment_to_ascii(
object_version.get_header_comment(False))
idc.SetEnumCmt(enum_id, nonrepeatable_headercomment, 0)
except KeyError:
pass
self.enum_ids[object_id] = enum_id
self.hash_provider.put_hash_struc_or_enum(enum_id, object_id)
logger.debug("adding enum id %s : '0x%.016X'" %
(self.hash_provider.hash_to_string(object_id), enum_id))