def to_dot(self, name: str, rpo: Dict[str, List[int]] = {}, showintervals=False) -> DotGraph: dotgraph = DotGraph(name) for n in self.nodes: if n in rpo: index = ":".join(str(i) for i in rpo[n]) labeltxt = index + ":" + n else: labeltxt = n dotgraph.add_node(n, labeltxt=labeltxt) if showintervals: clusters: Dict[str, Set[Tuple[str, str]]] = {} for src in self.edges: for tgt in self.edges[src]: if self.rev_intervals[src] == self.rev_intervals[tgt]: h = self.rev_intervals[src] clusters.setdefault(h, set([])) clusters[h].add((src, tgt)) else: dotgraph.add_edge(src, tgt) for (h, edges) in clusters.items(): dotgraph.add_cluster(h, edges) else: for src in self.edges: for tgt in self.edges[src]: dotgraph.add_edge(src, tgt) return dotgraph
class DotCfg(object): def __init__( self, graphname, fn, looplevelcolors=[], # [ color numbers ] showpredicates=False, # show branch predicates on edges showcalls=False, # show call instrs on nodes mips=False, # for mips subtract 4 from block end addr sink=None, # restrict paths to basic block destination segments=[], # restrict paths to include these basic blocks replacements={}): # replacement text for node and edge labels self.fn = fn self.graphname = graphname self.looplevelcolors = looplevelcolors self.showpredicates = showpredicates self.showcalls = showcalls self.mips = mips self.sink = sink self.segments = segments self.replacements = replacements self.pathnodes = set([]) self.dotgraph = DotGraph(graphname) def build(self): if not self.sink is None: self.restrict_nodes(self.sink) elif len(self.segments) > 0: self.restrict_paths(self.segments) else: self.pathnodes = self.fn.cfg.blocks for n in self.fn.cfg.blocks: self.add_cfg_node(n) for e in self.fn.cfg.edges: self.add_cfg_edge(e) return self.dotgraph def restrict_nodes(self, sink): nodes = self.fn.cfg.blocks edges = self.fn.cfg.edges # adjacency list n -> [ n ] if not sink in nodes: print('Sink ' + sink + ' not found in nodes') self.pathnodes = nodes return g = UG.DirectedGraph(nodes, edges) g.find_paths(self.fn.faddr, sink) for p in g.paths: print('Path: ' + str(p)) self.pathnodes = self.pathnodes.union(p) if len(self.pathnodes) == 0: self.pathnodes = nodes def restrict_paths(self, segments): nodes = self.fn.cfg.blocks edges = self.fn.cfg.edges for b in segments: if not b in nodes: print('Segment ' + b + ' not found in nodes') self.pathnodes = nodes return segments = [self.fn.faddr] + segments g = UG.DirectedGraph(nodes, edges) for i in range(len(segments) - 1): src = segments[i] dst = segments[i + 1] g.find_paths(src, dst) for p in g.paths: print('Path: ' + str(p)) self.pathnodes = self.pathnodes.union(p) if len(self.pathnodes) == 0: self.pathnodes = nodes def get_branch_instruction(self, edge): srcblock = self.fn.cfg.blocks[edge] instraddr = srcblock.lastaddr if instraddr.startswith('B'): ctxtaddr = instraddr[2:].split('_') iaddr = int(ctxtaddr[1], 16) if self.mips: iaddr -= 4 # delay slot instraddr = 'B:' + ctxtaddr[0] + '_' + hex(iaddr) else: instraddr = int(instraddr, 16) if self.mips: instraddr -= 4 # take into account delay slot instraddr = hex(instraddr) return self.fn.get_instruction(instraddr) def to_json(self): d = {} d['nodes'] = [] d['edges'] = {} for n in self.fn.cfg.blocks: d['nodes'].append(str(n)) for e in self.fn.cfg.edges: d['edges'][str(e)] = {} def default(): for tgt in self.fn.cfg.edges[e]: d['edges'][str(e)][str(tgt)] = 'none' if len(self.fn.cfg.edges[e]) > 1: branchinstr = self.get_branch_instruction(e) if branchinstr.is_branch_instruction(): ftconditions = branchinstr.get_ft_conditions() if len(ftconditions) > 1: for i, tgt in enumerate(self.fn.cfg.edges[e]): d['edges'][str(e)][str(tgt)] = ftconditions[i] else: default() else: default() else: default() return d def replace_text(self, txt): result = txt for src in sorted(self.replacements, key=lambda x: len(x), reverse=True): result = result.replace(src, self.replacements[src]) return result def add_cfg_node(self, n): if not n in self.pathnodes: return basicblock = self.fn.get_block(str(n)) blocktxt = str(n) color = 'lightblue' if self.showcalls: callinstrs = basicblock.get_call_instructions() callinstrs = [str(i.get_annotation()) for i in callinstrs] print(' \n'.join([str(a) for a in callinstrs])) if len(callinstrs) > 0: blocktxt = (blocktxt + '\\n' + '\\n'.join([str(a) for a in callinstrs])) if len(self.looplevelcolors) > 0: looplevels = self.fn.cfg.get_loop_levels(n) if len(looplevels) > 0: level = len(looplevels) if level > len(self.looplevelcolors): color = self.looplevelcolors[-1] else: color = self.looplevelcolors[level - 1] if n == self.fn.faddr: color = 'purple' blocktxt = self.replace_text(blocktxt) self.dotgraph.add_node(str(n), labeltxt=str(blocktxt), color=color) def add_cfg_edge(self, e): if not e in self.pathnodes: return def default(): for tgt in self.fn.cfg.edges[e]: if tgt in self.pathnodes: self.dotgraph.add_edge(str(e), str(tgt), labeltxt=None) labeltxt = None if len(self.fn.cfg.edges[e]) > 1: branchinstr = self.get_branch_instruction(e) if branchinstr and branchinstr.is_branch_instruction(): if self.showpredicates: ftconditions = branchinstr.get_ft_conditions() if len(ftconditions) == 2: for i, tgt in enumerate(self.fn.cfg.edges[e]): if tgt in self.pathnodes: labeltxt = str(ftconditions[i]) labeltxt = self.replace_text(labeltxt) self.dotgraph.add_edge(str(e), str(tgt), labeltxt=labeltxt) else: default() else: default() else: default() else: default()
class DotCallgraph(object): def __init__(self, graphname, callgraph, sinks=[], startaddr=None, getname=lambda x: x): self.graphname = graphname self.callgraph = callgraph # address -> address/name -> count self.dotgraph = DotGraph(graphname) self.dotgraph.rankdir = 'LR' self.sinks = sinks self.startaddr = startaddr self.pathnodes = set([]) self.getname = getname if self.startaddr and not self.sinks: self.restrict_nodes_from(self.startaddr) else: self.restrict_nodes() def build(self, coloring=lambda n: 'purple'): # name -> color / None if len(self.sinks) > 0: self.restrict_nodes() for n in self.callgraph: if coloring(n) is None: continue self.add_cg_node(n, coloring(n)) for d in self.callgraph[n]: self.add_cg_edge(n, d, self.callgraph[n][d], coloring) return self.dotgraph def restrict_nodes(self): nodes = set([]) edges = {} for n in self.callgraph: nodes.add(n) for d in self.callgraph[n]: nodes.add(d) edges.setdefault(n, []) edges[n].append(d) if self.startaddr is None: self.pathnodes = nodes return g = UG.DirectedGraph(nodes, edges) if len(self.sinks) > 0: g.find_paths(self.startaddr, self.sinks[0]) for p in g.paths: print('Path: ' + str(p)) self.pathnodes = self.pathnodes.union(p) if len(self.pathnodes) == 0: self.pathnodes = nodes else: self.pathnodes = nodes def restrict_nodes_from(self, startaddr): nodes = set([]) edges = {} nodes.add(startaddr) for d in self.callgraph[startaddr]: nodes.add(d) edges.setdefault(startaddr, []) edges[startaddr].append(d) nodecount = len(nodes) while True: for n in self.callgraph: if n in nodes: for d in self.callgraph[n]: nodes.add(d) edges.setdefault(n, []) edges[n].append(d) if len(nodes) == nodecount: break nodecount = len(nodes) self.pathnodes = nodes def add_cg_node(self, n, color): blocktxt = self.getname(str(n)) if str(n) in self.pathnodes: self.dotgraph.add_node(str(n), labeltxt=blocktxt, color=color) def add_cg_edge(self, n, d, count, coloring=lambda n: 'purple'): labeltxt = str(count) if coloring(d) is None: return if str(n) in self.pathnodes and str(d) in self.pathnodes: blocktxt = self.getname(str(d)) self.dotgraph.add_node(str(d), labeltxt=blocktxt, color=coloring(d)) self.dotgraph.add_edge(str(n), str(d))
class DotCfg: def __init__( self, graphname: str, fn: "chb.app.Function.Function", looplevelcolors: List[str] = [], # [ color numbers ] showpredicates: bool = False, # show branch predicates on edges showcalls: bool = False, # show call instrs on nodes showinstr_opcodes: bool = False, # show all instrs on nodes showinstr_text: bool = False, # show all instr annotations on nodes showstores: bool = False, # show all STR and STRB and STRH instr annotations mips: bool = False, # for mips subtract 4 from block end addr sink: str = None, # restrict paths to basic block destination segments: List[ str] = [], # restrict paths to include these basic blocks # replacement text for node and edge labels replacements: Dict[str, str] = {} ) -> None: self.fn = fn self.graphname = graphname self.looplevelcolors = looplevelcolors self.showpredicates = showpredicates self.showcalls = showcalls self.showinstr_opcodes = showinstr_opcodes self.showinstr_text = showinstr_text self.showstores = showstores self.mips = mips self.sink = sink self.segments = segments self.replacements = replacements self.pathnodes: Set[str] = set([]) self.dotgraph = DotGraph(graphname) def build(self) -> DotGraph: if self.sink is not None: self.restrict_nodes(self.sink) elif len(self.segments) > 0: self.restrict_paths(self.segments) else: self.pathnodes = set(self.fn.cfg.blocks.keys()) for n in self.fn.cfg.blocks: self.add_cfg_node(n) for e in self.fn.cfg.edges: self.add_cfg_edge(e) return self.dotgraph def restrict_nodes(self, sink: str) -> None: nodes = self.fn.cfg.blocks edges = self.fn.cfg.edges # adjacency list n -> [ n ] if sink not in nodes: print('Sink ' + sink + ' not found in nodes') self.pathnodes = set(nodes.keys()) return g = UG.DirectedGraph(list(nodes.keys()), edges) g.find_paths(self.fn.faddr, sink) for p in g.paths: print('Path: ' + str(p)) self.pathnodes = self.pathnodes.union(p) if len(self.pathnodes) == 0: self.pathnodes = set(nodes.keys()) def restrict_paths(self, segments: List[str]) -> None: nodes = self.fn.cfg.blocks edges = self.fn.cfg.edges for b in segments: if b not in list(nodes.keys()): print('Segment ' + b + ' not found in nodes') self.pathnodes = set(nodes.keys()) return segments = [self.fn.faddr] + segments g = UG.DirectedGraph(list(nodes.keys()), edges) for i in range(len(segments) - 1): src = segments[i] dst = segments[i + 1] g.find_paths(src, dst) for p in g.paths: print('Path: ' + str(p)) self.pathnodes = self.pathnodes.union(p) if len(self.pathnodes) == 0: self.pathnodes = set(nodes.keys()) def get_branch_instruction(self, edge: str) -> "chb.app.Instruction.Instruction": srcblock = self.fn.cfg.blocks[edge] instraddr = srcblock.lastaddr if instraddr.startswith('B'): ctxtaddr = instraddr[2:].split('_') iaddr_i = int(ctxtaddr[1], 16) if self.mips: iaddr_i -= 4 # delay slot instraddr = 'B:' + ctxtaddr[0] + '_' + hex(iaddr_i) else: instraddr_i = int(instraddr, 16) if self.mips: instraddr_i -= 4 # take into account delay slot instraddr = hex(instraddr_i) return self.fn.instruction(instraddr) def to_json(self) -> Dict[str, Any]: d: Dict[str, Any] = {} d['nodes'] = [] d['edges'] = {} for n in self.fn.cfg.blocks: d['nodes'].append(str(n)) for e in self.fn.cfg.edges: d['edges'][str(e)] = {} def default() -> None: for tgt in self.fn.cfg.edges[e]: d['edges'][str(e)][str(tgt)] = 'none' if len(self.fn.cfg.edges[e]) > 1: branchinstr = self.get_branch_instruction(e) if branchinstr.is_branch_instruction: ftconditions = branchinstr.ft_conditions if len(ftconditions) > 1: for i, tgt in enumerate(self.fn.cfg.edges[e]): d['edges'][str(e)][str(tgt)] = ftconditions[i] else: default() else: default() else: default() return d def replace_text(self, txt: str) -> str: result = txt for src in sorted(self.replacements, key=lambda x: len(x), reverse=True): result = result.replace(src, self.replacements[src]) return result def add_cfg_node(self, n: str) -> None: if n not in self.pathnodes: return basicblock = self.fn.block(str(n)) blocktxt = str(n) color = 'lightblue' if self.showinstr_opcodes: instrs = basicblock.instructions.values() pinstrs = [i.opcodetext for i in instrs] blocktxt = (blocktxt + "\\n" + "\\n".join(pinstrs)) elif self.showinstr_text: instrs = basicblock.instructions.values() pinstrs = [i.annotation for i in instrs] blocktxt = (blocktxt + "\\n" + "\\n".join(pinstrs)) elif self.showcalls or self.showstores: if self.showcalls: callinstrs = basicblock.call_instructions pcallinstrs = [i.annotation for i in callinstrs] print(' \n'.join([str(a) for a in pcallinstrs])) if len(callinstrs) > 0: blocktxt = (blocktxt + '\\n' + '\\n'.join(pcallinstrs)) if self.showstores: storeinstrs = basicblock.store_instructions pstoreinstrs = [i.annotation for i in storeinstrs] print(' \n'.join([str(a) for a in pstoreinstrs])) if len(storeinstrs) > 0: blocktxt = (blocktxt + "\\n" + "\\n".join(pstoreinstrs)) if len(self.looplevelcolors) > 0: looplevels = self.fn.cfg.loop_levels(n) if len(looplevels) > 0: level = len(looplevels) if level > len(self.looplevelcolors): color = self.looplevelcolors[-1] else: color = self.looplevelcolors[level - 1] # if n == self.fn.faddr: # color = 'purple' blocktxt = self.replace_text(blocktxt) self.dotgraph.add_node(str(n), labeltxt=str(blocktxt), color=color) def add_cfg_edge(self, e: str) -> None: if e not in self.pathnodes: return def default() -> None: for tgt in self.fn.cfg.edges[e]: if tgt in self.pathnodes: self.dotgraph.add_edge(str(e), str(tgt), labeltxt=None) labeltxt: Optional[str] = None if len(self.fn.cfg.edges[e]) > 1: if self.showpredicates: branchinstr = self.get_branch_instruction(e) if branchinstr and branchinstr.is_branch_instruction: ftconditions = branchinstr.ft_conditions if len(ftconditions) == 2: for i, tgt in enumerate(self.fn.cfg.edges[e]): if tgt in self.pathnodes: labeltxt = str(ftconditions[i]) labeltxt = self.replace_text(labeltxt) self.dotgraph.add_edge(str(e), str(tgt), labeltxt=labeltxt) else: default() else: default() else: default() else: default()