import sark
import idaapi
import logging
import DIE.Lib.DIE_Exceptions
from DIE.Lib.IDAConnector import get_native_size, regOffsetToName,\
get_function_name, get_function_start_address, get_function_end_address
#
# This file contains several wrappers for common IDA data type such as Functions, Function Argument,
# Structs, Struct Elements and arrays.
# Since retrieving this type information is not such a "trivial" task in IDA API, this wrappers are
# designed to act as convenience classes, making the data more easily accessible.
#
#######################################################################################################################
#
# IDA Array class wrapper
#
class Array():
"""
Array Class
"""
def __init__(self, type):
self.logger = logging.getLogger(__name__)
self.type_info = type
self.array_type_data = idaapi.array_type_data_t()
self.element_type = None
self.element_num = 0
self.element_size = 0
self.elements = []
# Extract array data
self.get_array_data()
def get_array_data(self):
"""
Extract the array data from tinfo_t object and populate all relevant class properties.
@return: True if successful, otherwise False
"""
try:
if self.type_info.is_array():
if self.type_info.get_array_details(self.array_type_data):
self.element_type = self.array_type_data.elem_type
self.element_num = self.array_type_data.nelems
self.element_size = self.element_type.get_size()
return True
return False
except Exception as ex:
self.logger.exception("Array: Error while getting array data: %s", ex)
return False
#######################################################################################################################
#
# IDA Function Argument class wrapper
#
class FuncArg():
"""
Function argument class
"""
def __init__(self, argType, argLoc, argNum, argName=None, isGuessed=False):
"""
Ctor
"""
self.logger = logging.getLogger(__name__)
self.isGussed = isGuessed # Is this argument known\guessed
self.argtype = argType # Argument type (type_info_t object)
self.argloc = argLoc # argloc object
self.argNum = argNum # Argument number (-1 for return argument)
self.argname = argName # Argument name
#self.inst_parser = InstructionParserX86()
def isReg(self):
"""
Is a register based argument
"""
if self.argloc.is_reg1():
return True
else:
return False
def isStack(self):
"""
Is a stack based argument
"""
if self.argloc.is_stkoff():
return True
else:
return False
def name(self):
"""
Argument name
"""
# If argument name was explicitly provided.
if self.argname:
return self.argname
# If this is a return argument.
if self.argNum is -1:
return "Ret_Arg"
# Otherwise, generate name according to offset.
#native_size = self.inst_parser.get_native_size()/8
native_size = get_native_size()/8
return "Arg_%s" % hex(self.argNum * native_size)
def getRegOffset(self):
"""
Get register offset (into ph.regnames)
"""
if self.argloc.is_reg1():
return self.argloc.reg1()
def offset(self):
"""
Stack Offset for stack args, or ph.regnames offset for register args
"""
if self.isStack():
return self.argloc.stkoff()
if self.isReg():
return self.getRegOffset()
self.logger.error("FuncArg: Failed to retrieve argument offset.")
return False
def registerName(self):
"""
Get register name for this arg
"""
if self.isReg():
return regOffsetToName(self.offset())
return None
def type_str(self):
"""
A string representation of the argument type
"""
typeStr = idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, self.argtype, '', '')
return typeStr
def isRetValue(self):
"""
Is this argument a return value?
"""
return self.argNum is -1
def getArgStr(self):
"""
Get a human readable argument description string
"""
guessedStr = ""
if self.isGussed:
guessedStr = "(Guessed)"
if self.isReg():
return "Arg: %d, Type: %s %s, Name: %s, Register: %s" % (self.argNum,
self.type_str(),
guessedStr,
self.name(),
self.registerName())
if self.isStack():
return "Arg: %s, Type: %s %s, Name: %s, StackOffset: %s" % (self.argNum,
self.type_str(),
guessedStr,
self.name(),
self.offset())
return
#######################################################################################################################
#
# IDA Function class wrapper
#
class Function():
"""
Function class
"""
def __init__(self, ea, iatEA=None, library_name=None):
"""
Ctor
"""
self.logger = logging.getLogger(__name__)
self.ea = ea # Effective Address of the function
self.iatEA = iatEA # If imported function, the address in the IAT
try:
function = sark.Function(ea)
except sark.exceptions.SarkNoFunction:
raise DIE.Lib.DIE_Exceptions.DieNoFunction("No Function at 0x%08X" % (ea, ))
self.funcName = get_function_name(function.ea)
self.func_start = function.startEA
self.func_end = function.endEA
self.proto_ea = self.getFuncProtoAdr() # Address of function prototype
self.typeInfo = idaapi.tinfo_t() # Function type info
self.funcInfo = idaapi.func_type_data_t() # Function info
self.argNum = 0 # Number of input arguments
self.args = [] # Function argument list
self.retArg = None # Return argument
self.library_name = library_name # If library function, name of containing library
self.isLibFunc = False
if self.iatEA:
self.isLibFunc = True # Is this a library function
elif sark.Function(ea).flags & (idaapi.FUNC_LIB | idaapi.FUNC_THUNK):
self.isLibFunc = True
try:
self.getArguments()
except Exception as ex:
self.logger.error("Failed to get function arguments for function %s: %s", self.funcName, ex)
def getFuncProtoAdr(self):
"""
Get the effective address of the function prototype definition.
In some cases will not be the same as the function ea (for example in library functions case)
"""
if self.iatEA:
return self.iatEA
return self.ea
def getArguments(self):
"""
Retrieve function arguments and populate the object`s args list.
"""
isGuessed = False # Is function prototype guessed
# Get function type info
if not idaapi.get_tinfo2(self.proto_ea, self.typeInfo):
idaapi.guess_tinfo2(self.proto_ea, self.typeInfo)
isGuessed = True
if self.typeInfo.empty():
raise RuntimeError("Failed to retrieve function type info for function %s at %s" % (self.funcName, hex(self.ea)))
# Get function detail
self.typeInfo.get_func_details(self.funcInfo)
# TODO: This seems to be creating false positives on 0 argument functions.
#if self.funcInfo.empty():
# errStr = "Failed to retrieve function info for function %s" % self.funcName
# raise RuntimeError(errStr)
self.argNum = len(self.funcInfo)
# Iterate function arguments
for argIndex in xrange(0, self.argNum):
argType = None # arg_type_info_t
argLoc = None # argloc_info
argName = None
#else: # Input Argument
argType = self.funcInfo.at(argIndex).type
argLoc = self.funcInfo.at(argIndex).argloc
argName = self.funcInfo.at(argIndex).name
curArg = FuncArg(argType, argLoc, argIndex, argName, isGuessed)
self.args.append(curArg)
# Set return argument
if not self.funcInfo.rettype.empty():
self.retArg = FuncArg(self.funcInfo.rettype,
self.funcInfo.retloc,
-1,
"Ret_Arg",
isGuessed)
def PrintFunction(self):
"""
Print Function Information
"""
for arg in self.args:
print arg.getArgStr()
#######################################################################################################################
#
# IDA Struct Element class wrapper
#
class StructElement():
"""
Struct Element
"""
def __init__(self, size, offset, type_, name=None, comment=None):
"""
Struct element class
@param size: Size of element
@param offset: Element offset within the struct
@param type_: Element type
@param name: Element name string
@param comment: Element comment (Optional)
"""
self.logger = logging.getLogger(__name__)
self.name = name
self.comment = comment
self.offset = offset
self.size = size
self.type = type_
def get_name(self):
"""
Get struct element`s name
"""
if self.name is None or self.name == "":
return "field_%d" % self.offset
return self.name
def type_name(self):
"""
Get type name (int, char, LPCSTR etc.)
"""
idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, self.type, '', '')
#######################################################################################################################
#
# IDA Struct class wrapper
#
class Struct():
"""
Struct class
"""
def __init__(self, type):
self.logger = logging.getLogger(__name__)
self.name = ""
self.size = 0
self.element_num = 0
self.is_union = False
self.elements = []
self.type_info = type
self.udt_type_data = idaapi.udt_type_data_t()
try:
if self.getStructData():
self.getElements()
except Exception as ex:
self.logger.exception("Error while extracting Struct data: %s",
idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, type, '', ''), ex)
return False
def getStructData(self):
"""
Extract the struct data from tinfo_t object and populate all relevant class properties.
@return: True if successful, otherwise False
"""
if self.type_info.is_udt():
if self.type_info.get_udt_details(self.udt_type_data):
self.name = idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, self.type_info, '', '')
self.size = self.udt_type_data.size
self.element_num = len(self.udt_type_data)
self.is_union = self.udt_type_data.is_union
return True
return False
def getElements(self):
"""
Get struct elements
"""
for element_index in xrange(0, self.element_num):
cur_element = self.udt_type_data[element_index]
name = None
comment = None
if cur_element.name is not None:
name = cur_element.name
if cur_element.cmt is not None:
comment = cur_element.cmt
strcut_elem = StructElement(cur_element.size,
cur_element.offset,
cur_element.type,
name,
comment)
self.elements.append(strcut_elem)
