def yatest_enums(self):
values = []
for flag in flags:
for prefix, enum_width, is_bitfield, num_fields in tests:
name = '%s_%x_%d_%x' % (prefix, enum_width, is_bitfield, flag)
ea = None
eid = idc.AddEnum(-1, name, flag)
self.assertNotEqual(eid, idaapi.BADADDR)
if enum_width != 0:
idc.SetEnumWidth(eid, enum_width)
if is_bitfield:
self.assertTrue(idc.SetEnumBf(eid, True))
idc.SetEnumCmt(eid, prefix + 'cmt', False)
idc.SetEnumCmt(eid, prefix + 'rpt', True)
for n in range(0, num_fields):
field = '%s_%d' % (name, n)
cid = None
if is_bitfield:
self.assertEqual(
idc.AddConstEx(eid, field, 1 << n, 1 << n), 0)
else:
self.assertEqual(idc.AddConst(eid, field, n), 0)
if n == 0:
ea = get_ea()
self.assertNotEqual(idaapi.op_enum(ea, 1, eid, 0),
idaapi.BADADDR)
cid = idc.GetConstByName(field)
self.assertTrue(idc.SetConstCmt(cid, field + 'cmt', False))
self.assertTrue(idc.SetConstCmt(cid, field + 'rpt', True))
values.append((name, ea))
yaunit.save('enums', values)
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 make_enum_member(self, object_version, value):
name = object_version.get_name()
object_id = object_version.get_id()
enum_object_id = object_version.get_parent_object_id()
enum_id = 0
try:
enum_id = self.enum_ids[enum_object_id]
except:
return
self.hash_provider.put_hash_enum_member(idc.GetEnumName(enum_id), name,
value, object_id)
# create member
if not idc.IsBitfield(enum_id):
bmask = -1
else:
bmask = object_version.get_object_flags()
member_id = idc.GetConstEx(enum_id, value, 0, bmask)
if member_id == idc.BADADDR:
idc.AddConstEx(enum_id, name, value, bmask)
member_id = idc.GetConstEx(enum_id, value, 0, bmask)
else:
if idc.SetConstName(member_id, name) == 0:
logger.error("Failed to set const name for enum_id")
# apply comments
try:
repeatable_headercomment = self.sanitize_comment_to_ascii(
object_version.get_header_comment(True))
idc.SetConstCmt(member_id, repeatable_headercomment, 1)
except KeyError:
pass
try:
nonrepeatable_headercomment = self.sanitize_comment_to_ascii(
object_version.get_header_comment(True))
idc.SetConstCmt(member_id, nonrepeatable_headercomment, 0)
except KeyError:
pass
self.enum_member_ids[(idc.GetEnumName(enum_id), name,
value)] = (member_id, object_id)
