From 833e390707e3542cd6ed907eddd8dad93523c30c Mon Sep 17 00:00:00 2001 From: Josselin Date: Fri, 1 Mar 2019 14:43:00 +0000 Subject: [PATCH] Clean call-graph printer (+ export every contract individually) Add call_graph.py script that is similar to the call-graph printer, but without the view/pure functions --- examples/scripts/call_graph.py | 54 +++++++++++++ slither/printers/call/call_graph.py | 116 ++++++++++++++++------------ 2 files changed, 120 insertions(+), 50 deletions(-) create mode 100644 examples/scripts/call_graph.py diff --git a/examples/scripts/call_graph.py b/examples/scripts/call_graph.py new file mode 100644 index 0000000000..7d1537d91f --- /dev/null +++ b/examples/scripts/call_graph.py @@ -0,0 +1,54 @@ +import os +import logging +import argparse +from slither import Slither +from slither.printers.all_printers import PrinterCallGraph +from slither.core.declarations.function import Function + +logging.basicConfig() +logging.getLogger("Slither").setLevel(logging.INFO) +logging.getLogger("Printers").setLevel(logging.INFO) + +class PrinterCallGraphStateChange(PrinterCallGraph): + + def _process_function(self, contract, function, contract_functions, contract_calls, solidity_functions, solidity_calls, external_calls, all_contracts): + if function.view or function.pure: + return + super()._process_function(contract, function, contract_functions, contract_calls, solidity_functions, solidity_calls, external_calls, all_contracts) + + def _process_internal_call(self, contract, function, internal_call, contract_calls, solidity_functions, solidity_calls): + if isinstance(internal_call, Function): + if internal_call.view or internal_call.pure: + return + super()._process_internal_call(contract, function, internal_call, contract_calls, solidity_functions, solidity_calls) + + def _process_external_call(self, contract, function, external_call, contract_functions, external_calls, all_contracts): + if isinstance(external_call[1], Function): + if external_call[1].view or external_call[1].pure: + return + super()._process_external_call(contract, function, external_call, contract_functions, external_calls, all_contracts) + +def parse_args(): + """ + """ + parser = argparse.ArgumentParser(description='Call graph printer. Similar to --print call-graph, but without printing the view/pure functions', + usage='call_graph.py filename') + + parser.add_argument('filename', + help='The filename of the contract or truffle directory to analyze.') + + parser.add_argument('--solc', help='solc path', default='solc') + + return parser.parse_args() + +def main(): + + args = parse_args() + slither = Slither(args.filename, is_truffle=os.path.isdir(args.filename), solc=args.solc) + + slither.register_printer(PrinterCallGraphStateChange) + + slither.run_printers() + +if __name__ == '__main__': + main() diff --git a/slither/printers/call/call_graph.py b/slither/printers/call/call_graph.py index ed5efdfaef..cf6776e699 100644 --- a/slither/printers/call/call_graph.py +++ b/slither/printers/call/call_graph.py @@ -5,7 +5,7 @@ what are the contracts/functions called. The output is a dot file named filename.dot """ - +from collections import defaultdict from slither.printers.abstract_printer import AbstractPrinter from slither.core.declarations.solidity_variables import SolidityFunction from slither.core.declarations.function import Function @@ -44,97 +44,113 @@ class PrinterCallGraph(AbstractPrinter): WIKI = 'https://github.com/trailofbits/slither/wiki/Printer-documentation#call-graph' - def __init__(self, slither, logger): - super(PrinterCallGraph, self).__init__(slither, logger) + def _process_functions(self, functions): + + contract_functions = defaultdict(set) # contract -> contract functions nodes + contract_calls = defaultdict(set) # contract -> contract calls edges + + solidity_functions = set() # solidity function nodes + solidity_calls = set() # solidity calls edges + external_calls = set() # external calls edges - self.contract_functions = {} # contract -> contract functions nodes - self.contract_calls = {} # contract -> contract calls edges + all_contracts = set() - for contract in slither.contracts: - self.contract_functions[contract] = set() - self.contract_calls[contract] = set() + for function in functions: + all_contracts.add(function.contract) + for function in functions: + self._process_function(function.contract, + function, + contract_functions, + contract_calls, + solidity_functions, + solidity_calls, + external_calls, + all_contracts) - self.solidity_functions = set() # solidity function nodes - self.solidity_calls = set() # solidity calls edges + render_internal_calls = '' + for contract in all_contracts: + render_internal_calls += self._render_internal_calls(contract, contract_functions, contract_calls) - self.external_calls = set() # external calls edges + render_solidity_calls = '' #self._render_solidity_calls(solidity_functions, solidity_calls) - self._process_contracts(slither.contracts) + render_external_calls = self._render_external_calls(external_calls) - def _process_contracts(self, contracts): - for contract in contracts: - for function in contract.functions: - self._process_function(contract, function) + return render_internal_calls + render_solidity_calls + render_external_calls - def _process_function(self, contract, function): - self.contract_functions[contract].add( + def _process_function(self, contract, function, contract_functions, contract_calls, solidity_functions, solidity_calls, external_calls, all_contracts): + contract_functions[contract].add( _node(_function_node(contract, function), function.name), ) for internal_call in function.internal_calls: - self._process_internal_call(contract, function, internal_call) + self._process_internal_call(contract, function, internal_call, contract_calls, solidity_functions, solidity_calls) for external_call in function.high_level_calls: - self._process_external_call(contract, function, external_call) + self._process_external_call(contract, function, external_call, contract_functions, external_calls, all_contracts) - def _process_internal_call(self, contract, function, internal_call): + def _process_internal_call(self, contract, function, internal_call, contract_calls, solidity_functions, solidity_calls): if isinstance(internal_call, (Function)): - self.contract_calls[contract].add(_edge( + contract_calls[contract].add(_edge( _function_node(contract, function), _function_node(contract, internal_call), )) elif isinstance(internal_call, (SolidityFunction)): - self.solidity_functions.add( + solidity_functions.add( _node(_solidity_function_node(internal_call)), ) - self.solidity_calls.add(_edge( + solidity_calls.add(_edge( _function_node(contract, function), _solidity_function_node(internal_call), )) - def _process_external_call(self, contract, function, external_call): + def _process_external_call(self, contract, function, external_call, contract_functions, external_calls, all_contracts): external_contract, external_function = external_call + if not external_contract in all_contracts: + return + # add variable as node to respective contract if isinstance(external_function, (Variable)): - self.contract_functions[external_contract].add(_node( + return + contract_functions[external_contract].add(_node( _function_node(external_contract, external_function), external_function.name )) - self.external_calls.add(_edge( + external_calls.add(_edge( _function_node(contract, function), _function_node(external_contract, external_function), )) - def _render_internal_calls(self): + def _render_internal_calls(self, contract, contract_functions, contract_calls): lines = [] - for contract in self.contract_functions: - lines.append(f'subgraph {_contract_subgraph(contract)} {{') - lines.append(f'label = "{contract.name}"') + lines.append(f'subgraph {_contract_subgraph(contract)} {{') + lines.append(f'label = "{contract.name}"') - lines.extend(self.contract_functions[contract]) - lines.extend(self.contract_calls[contract]) + lines.extend(contract_functions[contract]) + lines.extend(contract_calls[contract]) - lines.append('}') + lines.append('}') return '\n'.join(lines) - def _render_solidity_calls(self): + def _render_solidity_calls(self, solidity_functions, solidity_calls): lines = [] lines.append('subgraph cluster_solidity {') lines.append('label = "[Solidity]"') - lines.extend(self.solidity_functions) - lines.extend(self.solidity_calls) + lines.extend(solidity_functions) + lines.extend(solidity_calls) lines.append('}') return '\n'.join(lines) - def _render_external_calls(self): - return '\n'.join(self.external_calls) + def _render_external_calls(self, external_calls): + return '\n'.join(external_calls) + + def output(self, filename): """ @@ -142,19 +158,19 @@ def output(self, filename): Args: filename(string) """ - if not filename: - filename = "contracts.dot" if not filename.endswith('.dot'): filename += '.dot' - - self.info(f'Call Graph: {filename}') + if filename == ".dot": + filename = "all_contracts.dot" with open(filename, 'w', encoding='utf8') as f: - f.write('\n'.join([ - 'strict digraph {', - self._render_internal_calls(), - self._render_solidity_calls(), - self._render_external_calls(), - '}', - ])) + self.info(f'Call Graph: {filename}') + f.write('\n'.join(['strict digraph {'] + [self._process_functions(self.slither.functions)] + ['}'])) + + + for derived_contract in self.slither.contracts_derived: + with open(f'{derived_contract.name}.dot', 'w', encoding='utf8') as f: + self.info(f'Call Graph: {derived_contract.name}.dot') + f.write('\n'.join(['strict digraph {'] + [self._process_functions(derived_contract.functions)] + ['}'])) +