"""Get direct callees of function at ea using CodeRefsFrom."""

callees = []

for head in idautils.FuncItems(ea):

for xref in idautils.CodeRefsFrom(head, False):

func = idaapi.get_func(xref)

if func and func.start_ea != ea:

if func.start_ea not in callees:

callees.append(func.start_ea)

if ea is None:

raise Exception("Address cannot be None")

disasm = idc.GetDisasm(ea)

try:

disasm = disasm[:disasm.index(';')]

except ValueError:

pass

if disasm == '':

return None

count_paths = 0

paths = nx.all_simple_paths(dg, start, end, cutoff)

DEBUG_PRINT((start, end, cutoff))

for x in paths:

global fcan

for f in fcan:

print(f)

print(fcan[f])

def __init__(self, msg):

ida_kernwin.Form.__init__(self, r"""STARTITEM 0

""", {

'note' : ida_kernwin.Form.StringLabel(msg)

path_count = 0

if cutoff < 1:

return

visited = [source]

stack = [iter(G[source])]

import time

start = time.time()

while stack:

children = stack[-1]

child = next(children, None)

if child is None:

stack.pop()

visited.pop()

elif len(visited) < cutoff:

if child == target:

path_count = path_count + 1

DEBUG_PRINT(path_count)

if path_count >= 100:

break

elif child not in visited:

visited.append(child)

stack.append(iter(G[child]))

else: #len(visited) == cutoff:

if child == target or target in children:

path_count = path_count + 1

break

runtime = time.time() - start

if runtime > 20:

return 100

"""Find simple cycles (elementary circuits) of a directed graph.

# Jon Olav Vik, 2010-08-09

def _unblock(thisnode):

"""Recursively unblock and remove nodes from B[thisnode]."""

if blocked[thisnode]:

blocked[thisnode] = False

while B[thisnode]:

_unblock(B[thisnode].pop())

def circuit(thisnode, startnode, component,result):

closed = False # set to True if elementary path is closed

path.append(thisnode)

blocked[thisnode] = True

for nextnode in component[thisnode]: # direct successors of thisnode

if nextnode == startnode:

result[startnode] = path

DEBUG_PRINT('get one result')

DEBUG_PRINT(result)

closed = True

return closed

elif not blocked[nextnode]:

if circuit(nextnode, startnode, component,result):

closed = True

if closed:

_unblock(thisnode)

else:

for nextnode in component[thisnode]:

if thisnode not in B[nextnode]: # TODO: use set for speedup?

B[nextnode].append(thisnode)

path.pop() # remove thisnode from path

return closed

path = [] # stack of nodes in current path

blocked = collections.defaultdict(bool) # vertex: blocked from search?

B = collections.defaultdict(list) # graph portions that yield no elementary circuit

result = {} # list to accumulate the circuits found

# Johnson's algorithm requires some ordering of the nodes.

# They might not be sortable so we assign an arbitrary ordering.

ordering=dict(zip(G,range(len(G))))

count = 0

for s in ordering:

DEBUG_PRINT(count)

count = count + 1

# Build the subgraph induced by s and following nodes in the ordering

subgraph = G.subgraph(node for node in G

if ordering[node] >= ordering[s])

# Find the strongly connected component in the subgraph

# that contains the least node according to the ordering

strongcomp = nx.strongly_connected_components(subgraph)

mincomp=min(strongcomp,

key=lambda nodes: min(ordering[n] for n in nodes))

component = G.subgraph(mincomp)

if component:

# smallest node in the component according to the ordering

startnode = min(component,key=ordering.__getitem__)

if startnode in result:

continue

for node in component:

blocked[node] = False

B[node][:] = []

dummy=circuit(startnode, startnode, component,result)

'''

xtree = {}

visited = []

towalk = [ea]

while towalk:

curr = towalk.pop()

if curr not in xtree: # the start point also will record in the tree

xtree[curr] = []

if curr not in visited:

visited.append(curr)

for x in idautils.XrefsTo(curr):

caller = idaapi.get_func(x.frm)

if caller:

caller = caller.start_ea

if caller not in visited:

towalk.append(caller)

xtree[curr].append(caller)

if maxlen and len(xtree) > maxlen:

return {}

'''

calleetree = {}

visited = []

towalk = [ea]

while towalk:

curr = towalk.pop()

if curr not in calleetree: # the start point also will record in the tree

calleetree[curr] = []

if curr not in visited:

visited.append(curr)

for x in GetCallees(curr):

if x not in visited:

towalk.append(x)

calleetree[curr].append(x)

if maxlen and len(calleetree) > maxlen:

return {}

syscalls = set()

try:

with open(filename,'rb') as fr:

syscalls = pickle.load(fr)

except:

print('oooooops, cannot load syscall list from ', filename)

raise

if not syscalls:

raise ValueError('syscalls list is empty')

def __init__(self):

'''

self.script_folder = ''

self.project_name = ''

print('---------------------', idc.ARGV[1])

arg = idc.ARGV[1]

self.script_folder = arg[arg.find('(')+2: arg.find(',')-1]

self.project_name = arg[arg.find(',')+2: arg.find(')')-1]

print('++++++++++project_name', self.project_name)

print('++++++++++script_folder',self.script_folder)

self.moduleName = ida_nalt.get_root_filename().replace('.','_') #name of current idb

if os.path.exists(self.moduleName):

#may need user's input to decide whether rewrite it or append it? this check shld be set as input in args

print('the db already exist')

clear = ConfirmDialog("Delete the current DB and create a new one?")

clear.Compile()

ok = clear.Execute()

if ok:

os.remove(self.moduleName)

else:

return

print('[Get_FunctionFeatures]moduleName: %s'%self.moduleName)

self.func_name_ea = {name:ea for ea, name in idautils.Names()} # all names within idb

self.ftable = collections.defaultdict(dict) # a dictionary stores the features of one function, will be refreshed for every function

self.exports = [] # all export functions

self.memop = {} #instructions with memory operation

self.syscalls = set()

self.priorMatrix = [('returnpoints', '_feature_returnpoints'), ('loopcount', '_feature_loopcount')]

#(ea, writemem, writetoglobal, cmpmem, loopcalc)

self.LoadExports()

print('table name: ' + self.moduleName)

def _CheckMemOp(self, ea):

'''

inst = idautils.DecodeInstruction(ea)

if inst == None:

return

if inst.itype in [160,159]:

# retn 159, retf 160

self.ftable["returnpoints"].append(ea)

elif inst.itype in [122,6,209]:

# mov 122 add 6 sub 209, write memory happened at first opr

if 2<= inst[0].type <=7:

#considered as memory write

if idc.get_segm_name(inst[0].addr) == '.idata':

self.ftable["memop"].append((ea,1,1,0,0))

else:

self.ftable["memop"].append((ea,1,0,0,0))

elif inst.itype in [27,210]:

#cmp 27 test 210

if (2<= inst[0].type <=7 and inst[0].type != 5) or (2<= inst[1].type <=7 and inst[1].type != 5):

#mem cmp

self.ftable["memop"].append((ea,0,0,1,0))

elif inst.itype in [44,34]:

#inc 44 dec 34;

self.ftable["memop"].append((ea,0,0,0,1))

elif inst.itype in [16]:

# call 13

if inst[0].type == 3 or inst[0].type == 4:

self.ftable["dynamiccall"].append(ea)

def _BuildBasicBlockInfo(self, f_ea):

f_start = f_ea

f_end = idc.find_func_end(f_start)

edges = set()

boundaries = set((f_start,))

self.ftable["exceptionhandlers"] = []

self.ftable["memop"] = []

self.ftable["returnpoints"] = []

self.ftable["dynamiccall"] = []

for head in idautils.Heads(f_start, f_end):

# If the element is an instruction

comm = idc.get_cmt(head, 1)

if comm != None:

if "Exception handler" in comm:

self.ftable["exceptionhandlers"].append(head)

self._CheckMemOp(head)

if ida_bytes.is_code(ida_bytes.get_full_flags(head)):

# Get the references made from the current instruction

# and keep only the ones local to the function.

refs = idautils.CodeRefsFrom(head, 0)

refs = set(filter(lambda x: x>=f_start and x<=f_end, refs))

if refs:

# If the flow continues also to the next (address-wise)

# instruction, we add a reference to it.

next_head = idc.next_head(head, f_end)

if ida_bytes.is_flow(ida_bytes.get_full_flags(next_head)):

refs.add(next_head)

# Update the boundaries found so far.

boundaries.update(refs)

# For each of the references found, and edge is

# created.

for r in refs:

# If the flow could also come from the address

# previous to the destination of the branching

# an edge is created.

if ida_bytes.is_flow(ida_bytes.get_full_flags(r)):

edges.add((idc.prev_head(r, f_start), r))

edges.add((head, r))

# Let's build the list of (startEA, startEA) couples

# for each basic block

sorted_boundaries = sorted(boundaries, reverse = True)

end_addr = idc.prev_head(f_end, f_start)

bb_addr = []

for begin_addr in sorted_boundaries:

#just feel this algorithm is great

bb_addr.append((begin_addr, end_addr))

# search the next end_addr which could be

# farther than just the previous head

# if data are interlaced in the code

# WARNING: it assumes it won't epicly fail ;)

end_addr = idc.prev_head(begin_addr, f_start)

while not ida_bytes.is_code(ida_bytes.get_full_flags(end_addr)):

end_addr = idc.prev_head(end_addr, f_start)

# And finally return the result

bb_addr.reverse()

bb_edges = set()

DEBUG_PRINT( 'BuildBasicBlockInfo' + hex(f_ea))

for (s,e) in edges:

sblock = [(x,y) for (x,y) in bb_addr if y==s] # x==s or

eblock = [(x,y) for (x,y) in bb_addr if x==e] # or y==e

if sblock!=[] and eblock!=[]:

bb_edges.add((sblock[0],eblock[0]))

return bb_addr, bb_edges

def _BuildDiGraph(self, f_ea):

self.ftable["dg"] = nx.DiGraph()

nodes,edges = self._BuildBasicBlockInfo(f_ea)

self.ftable["dg"].add_nodes_from(nodes)

self.ftable["dg"].add_edges_from(edges)

return

def LoadExports(self):

for i, ordinal, ea, name in idautils.Entries():

if not ordinal == ea:

# by observation: ordinal == address => not exported

self.exports.append(ea)

def _feature_name(self,f_ea):

return idaapi.get_func_name(f_ea)

def _feature_ea(self, f_ea):

return hex(f_ea)

def _feature_size(self,f_ea):

'''

fun = idaapi.get_func(f_ea)

return fun.end_ea-fun.start_ea+1

def _feature_export(self, f_ea):

'''

if f_ea in self.exports:

return 1

else:

return 0

def _featrue_callee(self, f_ea):

'''

callees = GetCallees(f_ea)

return len(callees),callees

def _feature_callers(self,f_ea):

'''

callers = get_callers(f_ea)

return len(callers)

def _feature_exceptionhandlers(self, f_ea):

'''

return len(self.ftable["exceptionhandlers"]),self.ftable["exceptionhandlers"]

def _feature_loopcount(self, f_ea):

'''

DEBUG_PRINT( 'feature_loopcount' + hex(f_ea))

loopcount = 0

mloops = set()

self.loops = {}

self.loops = simple_cycles(self.ftable["dg"])

loopcount = len(self.loops)

return loopcount

def _feature_returnpoints(self, f_ea):

'''

DEBUG_PRINT("in returnpoints")

fun = idaapi.get_func(f_ea)

visited = []

retVal = []

for ret in self.ftable["returnpoints"]:

towalk = [ret]

while towalk:

curr = towalk.pop()

if curr < fun.start_ea or curr > fun.end_ea + 1:

continue

if curr not in visited:

visited.append(curr)

inst = GetInstruction(curr)

if inst is None:

continue

elif 'eax' in inst:

retVal.append((ret,curr,inst))

continue

for xto in idautils.XrefsTo(curr, 0):

DEBUG_PRINT('xto')

if xto.frm not in visited:

DEBUG_PRINT(xto.frm)

towalk.append(xto.frm)

DEBUG_PRINT(retVal)

return len(self.ftable["returnpoints"]), retVal

def _feature_paths(self, f_ea):

'''

return 0

paths_count = 0

start = sorted(self.ftable["dg"].nodes())[0]

DEBUG_PRINT('start')

DEBUG_PRINT(start)

cutoff = len(self.ftable["dg"])//2

if cutoff > 70:

return 100

for ret in self.ftable["returnpoints"]:

tar = None

for (x,y) in self.ftable["dg"].nodes():

if y == ret:

tar = (x,y)

break

if tar != None:

DEBUG_PRINT((start, tar, cutoff))

paths_count = paths_count + simple_paths_count(self.ftable["dg"], start, tar, cutoff)

if paths_count > 100:

break

DEBUG_PRINT(paths_count)

return paths_count

def _feature_arg(self,ea):

'''

typeinfo = idc.get_type(ea)

if typeinfo == None:

return 0

else:

argstr = typeinfo[typeinfo.find('(')+1:typeinfo.rfind(')')].split(',')

return len(argstr)

def _feature_localvar(self, f_ea):

'''

lvarcount = 0

DEBUG_PRINT( 'feature_localvar' + hex(f_ea))

try:

id = idc.get_frame_id(f_ea)

offset = idc.get_last_member(id)

except:

DEBUG_PRINT('no stack info for function {ea}'.format(ea =f_ea)) #BF96206A

return 0

while offset >= idc.get_first_member(id):

msize = idc.get_member_size(id, offset)

mname = idc.get_member_name(id, offset)

if mname == None:

offset = offset - 1

else:

lvarcount = lvarcount +1

offset = offset - msize

if mname == ' s':

lvarcount = 0

return lvarcount

def _feature_writetoglobalvar(self, f_ea):

'''

lwritetog = list(filter(lambda x: x[2] == 1, self.ftable["memop"]))

for item in lwritetog:

DEBUG_PRINT(item[0])

writetog_count = len(lwritetog)

return writetog_count,lwritetog

def _feature_cmpmemory(self, f_ea):

'''

lcmpmem = list(filter(lambda x: x[3] == 1, self.ftable["memop"]))

DEBUG_PRINT('IN cmpmemory')

for item in lcmpmem:

DEBUG_PRINT(hex(item[0]))

cmpmem_count = len(lcmpmem)

return cmpmem_count,lcmpmem

def _feature_ClsorInstMethod(self,f_ea):

'''

func_name = idaapi.get_func_name(f_ea)

prototype = ida_name.demangle_name(func_name, 0)

if not prototype:

return 0

elif '__thiscall' in prototype:

return 1

else:

return 0

def _feature_dynamiccalls(self, default_arg):

'''

return len(self.ftable["dynamiccall"]),self.ftable["dynamiccall"]

#add new feature here

def _feature_syscalls(self,f_ea):

'''

calleetree = {}

syscallcount = []

calleetree[f_ea] = get_callees(f_ea)

for ea in calleetree[f_ea]:

fname = idc.get_func_name(ea)

if fname in self.syscalls:#

syscallcount.append(fname) #better record the syscalls name of address

return len(syscallcount), syscallcount

def _feature_functiontype(self, f_ea):

'''

imflag = 0

for loop in self.loops.values():

for block in loop:

for l_ea in idautils.Heads(block[0],block[1]):

inst = idautils.DecodeInstruction(l_ea)

if inst == None:

continue

if inst.itype in [122]: # mov

if 3 == inst[0].type or 4 == inst[0].type:

imflag = 1

elif inst.itype in [124,207,107]: #movs/movsd, stos lods

imflag = 1

elif inst.itype in [16]: # call library function

istr = GetInstruction(l_ea)

if 'strcpy' in istr or 'memcpy' in istr or 'alloc' in istr or 'free' in istr:

imflag = 1

if imflag:

return 1

else:

return 0

def BuildFeatureTable(self):

'''

DEBUG_PRINT("in BuildFeatureTable")

if self.script_folder:

filename = self.script_folder + r'\syscall.pkl'

else:

filename = r'syscall.pkl'

self.syscalls = getsyscalls(filename)

flist = self.priorMatrix

from FunctionMatrix import FunctionMatrix

# get the function list of the class

methods = inspect.getmembers(FunctionMatrix, predicate=inspect.isfunction)

for method in methods:

if '_feature_' in method[0]:

if method[0] in self.priorMatrix:

pass

flist.append((method[0][8:].strip('_'),method[0]))

DEBUG_PRINT('flist: ')

print(flist)

# connect to mongo db

client = MongoClient('localhost',24444)

db = client[self.moduleName] # can have a higher name, project name?

dbctrl = db[self.moduleName]

dbctrl.remove() # just during test/debug

for f in self.func_name_ea.values(): #functions' start addresses

fun = idaapi.get_func(f)

if fun: #

ffeatures = collections.OrderedDict()

self._BuildDiGraph(fun.start_ea)

for fl in flist:

ffeatures[fl[0]] = str(getattr(self, fl[1])(f)) # []?

DEBUG_PRINT('ffeatures')

DEBUG_PRINT(ffeatures)

dbctrl.insert(ffeatures)

else:

pass

db.collection_names()

dbctrl.count()