def _getChildNumbers(self, parentId, childId):
parentNode = node_list_find(self.graph.graph.nodes, parentId)
numbers = []
if parentNode.has_child1 and parentNode.child1.node_id == childId:
numbers.append(1)
if parentNode.has_child2 and parentNode.child2.node_id == childId:
numbers.append(2)
return numbers
def setRootNode(self, rootNodeAddress):
rootNode = node_list_find(self.graph.graph.nodes, rootNodeAddress)
if not rootNode:
raise Exception("Invalid root node")
if rootNode in self.targetNodes:
raise Exception("Node already used")
self.rootNode = rootNode
self.resetColored()
print "Set root node to {}".format(hex(rootNode.node_id))
def addTargetNode(self, targetNodeAddress):
targetNode = node_list_find(self.graph.graph.nodes, targetNodeAddress)
if not targetNode:
print "WARNING: Node not found in CFG (is it reachable from a function or from the entrypoint ?)"
return
if targetNode in self.targetNodes:
print "WARNING: Target node already in use"
return
self.targetNodes.append(targetNode)
self.coloredNodes.add(targetNodeAddress)
self.colorNode(targetNodeAddress, self.targetColor)
print "Added target node {}".format(hex(targetNode.node_id))
def preview_match(self, address, base_text, match_type):
appended_text = ""
node = node_list_find(self.graph.graph.nodes, address)
if node:
if match_type in ["match_full", "match_opcode", "match_opcode_arg1", "match_opcode_arg2", "match_opcode_arg3"]:
appended_text = ": " + node.info.opcode
if match_type in ["match_full", "match_opcode_arg1", "match_opcode_arg2", "match_opcode_arg3"]:
if match_type in ["match_full", "match_opcode_arg1"] and node.info.nargs >= 1:
appended_text += " " + node.info.arg1
elif node.info.nargs >= 1:
appended_text += " *"
if match_type in ["match_full", "match_opcode_arg2"] and node.info.nargs >= 2:
appended_text += ", " + node.info.arg2
elif node.info.nargs >= 2:
appended_text += ", *"
if match_type in ["match_full", "match_opcode_arg3"] and node.info.nargs >= 3:
appended_text += ", " + node.info.arg3
elif node.info.nargs >= 3:
appended_text += ", *"
return base_text + appended_text
def preview_match(self, address, base_text, type):
appended_text = ""
node = node_list_find(self.graph.graph.nodes, address)
if node:
if type in [
"match_full", "match_opcode", "match_opcode_arg1",
"match_opcode_arg2"
]:
appended_text = ": " + node.info.opcode
added_arg1 = False
if type in ["match_full", "match_opcode_arg1"
] and node.info.nargs >= 1:
appended_text += " " + node.info.arg1
added_arg1 = True
if type in ["match_full", "match_opcode_arg2"
] and node.info.nargs >= 2:
if not added_arg1:
appended_text += " *"
appended_text += ", " + node.info.arg2
elif type in ["match_opcode_arg1"] and node.info.nargs >= 2:
appended_text += ", *"
return base_text + appended_text
def dis(self, ea, is_child1 = False, ifrom=None):
"""Disassemble the current address and fill the nodes list.
Args:
ea (ea_t): Effective address.
ifrom (node_t*): Predecessor node.
is_child1 (bool)
"""
node_list = self.graph.nodes
args_queue = []
args_queue.append((ea, is_child1, ifrom))
while args_queue != []:
ea, is_child1, ifrom = args_queue.pop(0)
try:
n = Node(ea, self.info, self.capstone)
except CodeException as e:
continue
#except Exception as e:
# print("WARNING:", e)
# continue
# If the node exists
if node_list_find(node_list, n.getid()):
if ifrom and is_child1 is not None:
# Link the father and the child
node_link(node_list_find(node_list, ifrom.getid()),
node_list_find(node_list, n.getid()), False,
is_child1)
continue
# Get the instruction
try:
inst = DecodeInstruction(ea)
except Exception as e:
print("WARNING:", e)
continue
if not inst:
continue
# Add the node
node_list_add(node_list, node_copy(node_alloc(), n))
if ifrom and is_child1 is not None:
node_link(node_list_find(node_list, ifrom.getid()),
node_list_find(node_list, n.getid()), False,
is_child1)
# No child
if inst.itype in RETS:
pass
# 1 remote child
elif inst.itype in JMPS:
try:
op = inst.ops[0]
except:
op = inst.Operands[0]
try:
args_queue.insert(0, (op.addr, False, n))
except Exception as e:
print("WARNING:", e)
pass
# 2 children (next, then remote) - except call
elif inst.itype in CJMPS:
try:
op = inst.ops[0]
except:
op = inst.Operands[0]
# Next
args_queue.insert(0, (inst.ea + inst.size, True, n))
# Remote
args_queue.insert(0, (op.addr, False, n))
# 2 children (next, then remote) - call
elif inst.itype in CALLS:
try:
op = inst.ops[0]
except:
op = inst.Operands[0]
# Next
# Catch the end of a noret function
if not is_noret(inst.ea):
args_queue.insert(0, (inst.ea + inst.size, True, n))
# Remote
if op.type in OP_MEM:
args_queue.insert(0, (op.addr, False, n))
# 1 child (next) - basic instruction
else:
args_queue.insert(0, (inst.ea + inst.size, True, n))
return
def fromNodeId(graph, nodeId):
node = node_list_find(graph, nodeId)
return PatternGeneratorNode(nodeId, node.info.opcode, node.info.arg1, node.info.arg2, node.info.arg3,
None, None, node.children_nb)
def generate(self, auto=False):
if self.rootNode is None:
print "WARNING: Missing the root node. Make sure to first \"Load the CFG\", then define the root node and target nodes (right click in IDA View) before you \"Generate a pattern\"."
return
if len(self.targetNodes) == 0:
if not auto:
print "WARNING: Missing target node(s). Make sure to first \"Load the CFG\", then define the root node and target nodes (right click in IDA View) before you \"Generate a pattern\"."
return
queue = deque()
targetNodes = set([i.node_id for i in self.targetNodes])
foundNodes = set()
coloring = set()
previous = {}
queue.append(self.rootNode)
coloring.add(self.rootNode.node_id)
while queue and (len(foundNodes) != len(targetNodes)):
node = queue.pop()
for child in self._getChildren(node):
if child.node_id not in coloring:
if child.node_id in targetNodes:
foundNodes.add(child.node_id)
coloring.add(child.node_id)
queue.appendleft(child)
previous[child.node_id] = node.node_id
if len(foundNodes) != len(targetNodes):
print "WARNING: Can not reach all target nodes from the root."
# Generate the nodes and edges of the pattern
patternNodes = OrderedDict()
patternEdges = []
for finalNode in reversed([s for s in self.targetNodes if s.node_id in foundNodes]):
node_id = finalNode.node_id
while node_id != self.rootNode.node_id:
if node_id not in patternNodes:
self.coloredNodes.add(node_id)
if node_id != self.rootNode.node_id and node_id not in foundNodes:
self.colorNode(node_id, self.nodeColor)
patternNodes[node_id] = PatternGeneratorNode.fromNodeId(self.graph.graph.nodes, node_id)
previous_id = previous[node_id]
patternEdges.append((previous_id, node_id))
node_id = previous_id
else:
break
patternNodes[self.rootNode.node_id] = PatternGeneratorNode.fromNodeId(self.graph.graph.nodes,
self.rootNode.node_id)
# Add edges between adjacent colored nodes
for node_id in self.coloredNodes:
node = node_list_find(self.graph.graph.nodes, node_id)
children = self._getChildren(node)
for c in children:
child_id = c.node_id
if child_id in self.coloredNodes:
patternEdges.append((node_id, child_id))
# Create the pattern graph's edges from patternEdges
for patternEdge in patternEdges:
numbers = self._getChildNumbers(patternEdge[0], patternEdge[1])
if 1 in numbers:
patternNodes[patternEdge[0]].child1 = patternEdge[1]
if 2 in numbers:
patternNodes[patternEdge[0]].child2 = patternEdge[1]
# Transformations
for _, patternNode in patternNodes.items():
if self.generic_arguments_option:
patternNode.arg1 = None
patternNode.arg2 = None
patternNode.arg3 = None
if self.lighten_memory_ops_option:
if patternNode.opcode in ["lea", "push", "pop"] or patternNode.opcode.startswith("mov"):
patternNode.opcode = None
if self.factorize_option:
for patternNodeId, patternNode in reversed(patternNodes.items()):
if patternNodeId not in patternNodes:
continue
while True:
if patternNode.child2 is not None or patternNode.child1 is None:
break
child = patternNodes[patternNode.child1]
if (patternNode.opcode is not None or child.opcode is not None) and (patternNode.opcode != child.opcode or patternNode.arg1 != child.arg1 or patternNode.arg2 != child.arg2 or patternNode.arg3 != child.arg3):
break
del patternNodes[patternNode.child1]
if child.child1 is not None:
patternNode.child1 = patternNodes[child.child1].node_id
else:
patternNode.child1 = None
patternNode.repeat = '*'
# End of transformations
patternStr = "digraph G {\n"
for patternNodeId, patternNode in reversed(patternNodes.items()):
patternStr += " {} [cond={}".format(hex(patternNodeId), self._getConditionString(patternNode))
if patternNode.repeat is not None:
patternStr += ', repeat=' + str(patternNode.repeat) + ', lazyrepeat=true'
patternStr += "]\n"
patternStr += "\n"
for patternNodeId, patternNode in reversed(patternNodes.items()):
if patternNode.child1 is not None:
patternStr += " {} -> {} [childnumber={}]\n".format(hex(patternNode.node_id),
hex(patternNode.child1), 1)
if patternNode.child2 is not None:
patternStr += " {} -> {} [childnumber={}]\n".format(hex(patternNode.node_id),
hex(patternNode.child2), 2)
patternStr += "}\n"
return patternStr
