Example #1
0
    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
Example #2
0
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()
Example #3
0
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()