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