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 add_enum(name=None, index=None, flags=idaapi.hexflag(), bitfield=False):
"""Create a new enum.
Args:
name: Name of the enum to create.
index: The index of the enum. Leave at default to append the enum as the last enum.
flags: Enum type flags.
bitfield: Is the enum a bitfield.
Returns:
An `Enum` object.
"""
if name is not None:
with ignored(exceptions.EnumNotFound):
_get_enum(name)
raise exceptions.EnumAlreadyExists()
if index is None or index < 0:
index = idaapi.get_enum_qty()
eid = idaapi.add_enum(index, name, flags)
if eid == idaapi.BADADDR:
raise exceptions.EnumCreationFailed('Failed creating enum "{}"'.format(name))
if bitfield:
idaapi.set_enum_bf(eid, bitfield)
return Enum(eid=eid)
def createenum(self, symbols):
"""
Given full symbols and addresses create an enum name with the library name (the string before !)
Some constants will fail due to weird characters in symbols used by MS. eg( `$)
symbols: (dict) A set of symbols and addresses that have been cleaned.
"""
enum_name = symbols.keys()[0].split('!')[0]
enum = idc.AddEnum(0, enum_name, idaapi.hexflag())
if enum == idaapi.BADADDR:
print "[!] Failed to create enum: %s\n" % enum_name
return
for symbol, address in symbols.iteritems():
# "ADVAPI32!RegCreateKeyExWStub": "0xffff8007be2f89f0"
org_symb = symbol
symbol = str(symbol.split('!')[1].encode('utf-8'))
symbol = symbol.strip()
symbol = 's_'+symbol
address = int(address,16)
ret = idc.AddConstEx(enum, symbol, address, -1)
if ret !=0:
print "[!] Failed to create constant for symbol %s - (%s). %s" % (org_symb,symbol,ENUM_ERRORS[ret])
continue
self.enums[address] = enum
print "[+] Finished adding enum %s\n" % enum_name
def add_enum(name=None, index=None, flags=idaapi.hexflag(), bitfield=False):
"""Create a new enum.
Args:
name: Name of the enum to create.
index: The index of the enum. Leave at default to append the enum as the last enum.
flags: Enum type flags.
bitfield: Is the enum a bitfield.
Returns:
An `Enum` object.
"""
if name is not None:
with ignored(exceptions.EnumNotFound):
_get_enum(name)
raise exceptions.EnumAlreadyExists()
if index is None or index < 0:
index = idaapi.get_enum_qty()
eid = idaapi.add_enum(index, name, flags)
if eid == idaapi.BADADDR:
raise exceptions.EnumCreationFailed(
'Failed creating enum "{}"'.format(name))
if bitfield:
idaapi.set_enum_bf(eid, bitfield)
return Enum(eid=eid)
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 add_enum(name=None, index=idaapi.BADADDR, flags=idaapi.hexflag(), bitfield=False):
"""Create a new enum."""
if name is not None:
with ignored(exceptions.EnumNotFound):
_get_enum(name)
raise exceptions.EnumAlreadyExists()
eid = idaapi.add_enum(index, name, flags)
if eid == idaapi.BADADDR:
raise exceptions.EnumCreationFailed('Failed creating enum "{}"'.format(name))
if bitfield:
idaapi.set_enum_bf(eid, bitfield)
return Enum(eid=eid)
def define_consts():
enum = AddEnum(-1, 'YARA_CONST', idaapi.hexflag())
if enum == BADADDR:
print 'Unable to create enum YARA_CONST'
return
if idc.AddConst(enum, 'UNDEFINED_32', UNDEFINED_MAGIC & 0xFFFFffff):
print 'Unable to create UNDEFINED_32 value'
return
const_id = GetConstByName('UNDEFINED_32')
if const_id == -1:
print 'Unable to get id of UNDEFINED_32'
return
if not SetConstCmt(const_id, 'internal UNDEFINED value for YARA VM', 1):
print 'failed setting comment for UNDEFINED_32'
return
return True
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import idaapi
import idautils
import idc
import unittest
import yaunit
flags = [
idaapi.hexflag(),
idaapi.charflag(),
idaapi.decflag(),
idaapi.octflag(),
idaapi.binflag(),
]
# name, enum_width 0->default, is bitfield, num_fields
tests = [
('std', 0, False, 0),
('std', 4, False, 0),
('bit', 0, True, 0),
('std_fields', 0, False, 0x20),
('bit_fields', 0, True, 0x20),
]
