Merge pull request #2413 from crytic/printer-cheatcode

Printer cheatcode

scripts/ci_test_printers.sh

@@ -5,7 +5,7 @@
 cd tests/e2e/solc_parsing/test_data/compile/ || exit
 
 # Do not test the evm printer,as it needs a refactoring
-ALL_PRINTERS="cfg,constructor-calls,contract-summary,data-dependency,echidna,function-id,function-summary,modifiers,call-graph,halstead,human-summary,inheritance,inheritance-graph,loc,martin,slithir,slithir-ssa,vars-and-auth,require,variable-order,declaration,ck"
+ALL_PRINTERS="cfg,constructor-calls,contract-summary,data-dependency,echidna,function-id,function-summary,modifiers,call-graph,halstead,human-summary,inheritance,inheritance-graph,loc,martin,slithir,slithir-ssa,vars-and-auth,require,variable-order,declaration,ck,cheatcode"
 
 # Only test 0.5.17 to limit test time
 for file in *0.5.17-compact.zip; do

slither/printers/all_printers.py

@@ -24,3 +24,4 @@
 from .summary.declaration import Declaration
 from .functions.dominator import Dominator
 from .summary.martin import Martin
+from .summary.cheatcodes import CheatcodePrinter

slither/printers/summary/cheatcodes.py

@@ -0,0 +1,74 @@
+"""
+    Cheatcode printer
+
+    This printer prints the usage of cheatcode in the code.
+"""
+from slither.printers.abstract_printer import AbstractPrinter
+from slither.slithir.operations import HighLevelCall
+from slither.utils import output
+
+
+class CheatcodePrinter(AbstractPrinter):
+
+    ARGUMENT = "cheatcode"
+
+    HELP = """
+        Print the usage of (Foundry) cheatcodes in the code.
+        For the complete list of Cheatcodes, see https://book.getfoundry.sh/cheatcodes/
+    """
+
+    WIKI = "https://github.com/trailofbits/slither/wiki/Printer-documentation#cheatcode"
+
+    def output(self, filename: str) -> output.Output:
+
+        info: str = ""
+
+        try:
+            vm = self.slither.get_contract_from_name("Vm").pop()
+        except IndexError:
+            return output.Output("No contract named VM found")
+
+        for contract in self.slither.contracts_derived:
+            # Check that the IS_TEST variable is set. (Only works for Foundry)
+            is_test_var = contract.variables_as_dict.get("IS_TEST", None)
+            is_test = False
+            if is_test_var is not None:
+                try:
+                    is_test = is_test_var.expression.value == "true"
+                except AttributeError:
+                    pass
+
+            if not is_test:
+                continue
+
+            found_contract: bool = False
+            contract_info: str = ""
+            for func in contract.functions_declared:
+                function_info = f"\t{func}\n"
+                found_function: bool = False
+                for node in func.nodes:
+                    for op in node.all_slithir_operations():
+                        if (
+                            isinstance(op, HighLevelCall)
+                            and op.function.contract == vm
+                            and op.function.visibility == "external"
+                        ):
+                            found_function = True
+                            function_info += (
+                                f"\t\t{op.function.name} - ({node.source_mapping.to_detailed_str()})\n"
+                                f"\t\t{node.expression}\n\n"
+                            )
+
+                if found_function:
+                    if found_contract is False:
+                        contract_info = f"{contract} ({contract.source_mapping.filename.short})\n"
+                        found_contract = True
+
+                    contract_info += function_info
+
+            if found_contract:
+                info += contract_info
+
+        self.info(info)
+        res = output.Output(info)
+        return res

tests/e2e/printers/test_data/test_printer_cheatcode/README.md

@@ -0,0 +1,7 @@
+# Counter
+
+Init using :
+
+```shell
+forge install foundry-rs/forge-std
+```

tests/e2e/printers/test_data/test_printer_cheatcode/foundry.toml

@@ -0,0 +1,6 @@
+[profile.default]
+src = 'src'
+out = 'out'
+libs = ['lib']
+
+# See more config options https://github.com/foundry-rs/foundry/tree/master/config

tests/e2e/printers/test_data/test_printer_cheatcode/src/Counter.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+contract Counter {
+    uint256 public number;
+
+    function setNumber(uint256 newNumber) public {
+        number = newNumber;
+    }
+
+    function increment() public {
+        number++;
+    }
+}

tests/e2e/printers/test_data/test_printer_cheatcode/test/Counter.t.sol

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import "forge-std/Test.sol";
+import "../src/Counter.sol";
+
+contract CounterTest is Test {
+    Counter public counter;
+
+    address public alice = address(0x42);
+    address public bob = address(0x43);
+
+    function difficulty(uint256 value) public {
+        // empty
+    }
+
+    function setUp() public {
+        counter = new Counter();
+        counter.setNumber(0);
+
+        vm.deal(alice, 1 ether);
+        vm.deal(bob, 2 ether);
+
+        difficulty(1);
+    }
+
+    function testIncrement() public {
+        vm.prank(alice);
+        counter.increment();
+        assertEq(counter.number(), 1);
+
+        vm.prank(bob);
+        counter.increment();
+        assertEq(counter.number(), 2);
+    }
+}

tests/e2e/printers/test_printers.py

@@ -1,17 +1,23 @@
 import re
+import shutil
 from collections import Counter
 from pathlib import Path
+import pytest
 
 from crytic_compile import CryticCompile
 from crytic_compile.platform.solc_standard_json import SolcStandardJson
 
 from slither import Slither
 from slither.printers.inheritance.inheritance_graph import PrinterInheritanceGraph
+from slither.printers.summary.cheatcodes import CheatcodePrinter
 from slither.printers.summary.slithir import PrinterSlithIR
 
 
 TEST_DATA_DIR = Path(__file__).resolve().parent / "test_data"
 
+foundry_available = shutil.which("forge") is not None
+project_ready = Path(TEST_DATA_DIR, "test_printer_cheatcode/lib/forge-std").exists()
+
 
 def test_inheritance_printer(solc_binary_path) -> None:
     solc_path = solc_binary_path("0.8.0")
@@ -48,6 +54,23 @@ def test_inheritance_printer(solc_binary_path) -> None:
     Path("test_printer.dot").unlink(missing_ok=True)
 
 
+@pytest.mark.skipif(
+    not foundry_available or not project_ready, reason="requires Foundry and project setup"
+)
+def test_printer_cheatcode():
+    slither = Slither(
+        Path(TEST_DATA_DIR, "test_printer_cheatcode").as_posix(), foundry_compile_all=True
+    )
+
+    printer = CheatcodePrinter(slither=slither, logger=None)
+    output = printer.output("")
+
+    assert (
+        output.data["description"]
+        == "CounterTest (test/Counter.t.sol)\n\tsetUp\n\t\tdeal - (test/Counter.t.sol#21 (9 - 32)\n\t\tvm.deal(alice,1000000000000000000)\n\n\t\tdeal - (test/Counter.t.sol#22 (9 - 30)\n\t\tvm.deal(bob,2000000000000000000)\n\n\ttestIncrement\n\t\tprank - (test/Counter.t.sol#28 (9 - 24)\n\t\tvm.prank(alice)\n\n\t\tassertEq - (test/Counter.t.sol#30 (9 - 38)\n\t\tassertEq(counter.number(),1)\n\n\t\tprank - (test/Counter.t.sol#32 (9 - 22)\n\t\tvm.prank(bob)\n\n\t\tassertEq - (test/Counter.t.sol#34 (9 - 38)\n\t\tassertEq(counter.number(),2)\n\n"
+    )
+
+
 def test_slithir_printer(solc_binary_path) -> None:
     solc_path = solc_binary_path("0.8.0")
     standard_json = SolcStandardJson()