Merge branch 'dev-doc' of github.com:crytic/slither into dev

setup.py

@@ -46,6 +46,7 @@ setup(
             "slither-mutate = slither.tools.mutator.__main__:main",
             "slither-read-storage = slither.tools.read_storage.__main__:main",
             "slither-doctor = slither.tools.doctor.__main__:main",
+            "slither-documentation = slither.tools.documentation.__main__:main",
         ]
     },
 )

slither/__main__.py

@@ -24,6 +24,7 @@ from slither.detectors.abstract_detector import AbstractDetector, DetectorClassi
 from slither.printers import all_printers
 from slither.printers.abstract_printer import AbstractPrinter
 from slither.slither import Slither
+from slither.utils import codex
 from slither.utils.output import output_to_json, output_to_zip, output_to_sarif, ZIP_TYPES_ACCEPTED
 from slither.utils.output_capture import StandardOutputCapture
 from slither.utils.colors import red, set_colorization_enabled
@@ -314,7 +315,6 @@ def parse_args(
         "Checklist (consider using https://github.com/crytic/slither-action)"
     )
     group_misc = parser.add_argument_group("Additional options")
-    group_codex = parser.add_argument_group("Codex (https://beta.openai.com/docs/guides/code)")
 
     group_detector.add_argument(
         "--detect",
@@ -555,47 +555,7 @@ def parse_args(
         default=False,
     )
 
-    group_codex.add_argument(
+    codex.init_parser(parser)
-        "--codex",
-        help="Enable codex (require an OpenAI API Key)",
-        action="store_true",
-        default=defaults_flag_in_config["codex"],
-    )
-
-    group_codex.add_argument(
-        "--codex-log",
-        help="Log codex queries (in crytic_export/codex/)",
-        action="store_true",
-        default=False,
-    )
-
-    group_codex.add_argument(
-        "--codex-contracts",
-        help="Comma separated list of contracts to submit to OpenAI Codex",
-        action="store",
-        default=defaults_flag_in_config["codex_contracts"],
-    )
-
-    group_codex.add_argument(
-        "--codex-model",
-        help="Name of the Codex model to use (affects pricing).  Defaults to 'text-davinci-003'",
-        action="store",
-        default=defaults_flag_in_config["codex_model"],
-    )
-
-    group_codex.add_argument(
-        "--codex-temperature",
-        help="Temperature to use with Codex.  Lower number indicates a more precise answer while higher numbers return more creative answers.  Defaults to 0",
-        action="store",
-        default=defaults_flag_in_config["codex_temperature"],
-    )
-
-    group_codex.add_argument(
-        "--codex-max-tokens",
-        help="Maximum amount of tokens to use on the response.  This number plus the size of the prompt can be no larger than the limit (4097 for text-davinci-003)",
-        action="store",
-        default=defaults_flag_in_config["codex_max_tokens"],
-    )
 
     # debugger command
     parser.add_argument("--debug", help=argparse.SUPPRESS, action="store_true", default=False)

slither/core/declarations/function.py

@@ -220,6 +220,9 @@ class Function(SourceMapping, metaclass=ABCMeta):  # pylint: disable=too-many-pu
 
         self._id: Optional[str] = None
 
+        # To be improved with a parsing of the documentation
+        self.has_documentation: bool = False
+
     ###################################################################################
     ###################################################################################
     # region General properties

slither/solc_parsing/declarations/function.py

@@ -91,6 +91,9 @@ class FunctionSolc(CallerContextExpression):
             Union[LocalVariableSolc, LocalVariableInitFromTupleSolc]
         ] = []
 
+        if "documentation" in function_data:
+            function.has_documentation = True
+
     @property
     def underlying_function(self) -> Function:
         return self._function

slither/tools/documentation/README.md

@@ -0,0 +1,6 @@
+# Demo
+
+This directory contains an example of Slither utility.
+
+See the [utility documentation](https://github.com/crytic/slither/wiki/Adding-a-new-utility)
+

slither/tools/documentation/__main__.py

@@ -0,0 +1,252 @@
+import argparse
+import logging
+import uuid
+from typing import Optional, Dict, List
+from crytic_compile import cryticparser
+from slither import Slither
+from slither.core.compilation_unit import SlitherCompilationUnit
+from slither.core.declarations import Function
+
+from slither.formatters.utils.patches import create_patch, apply_patch, create_diff
+from slither.utils import codex
+
+logging.basicConfig()
+logging.getLogger("Slither").setLevel(logging.INFO)
+
+logger = logging.getLogger("Slither")
+
+
+def parse_args() -> argparse.Namespace:
+    """
+    Parse the underlying arguments for the program.
+    :return: Returns the arguments for the program.
+    """
+    parser = argparse.ArgumentParser(description="Demo", usage="slither-documentation filename")
+
+    parser.add_argument("project", help="The target directory/Solidity file.")
+
+    parser.add_argument(
+        "--overwrite", help="Overwrite the files (be careful).", action="store_true", default=False
+    )
+
+    parser.add_argument(
+        "--force-answer-parsing",
+        help="Apply heuristics to better parse codex output (might lead to incorrect results)",
+        action="store_true",
+        default=False,
+    )
+
+    parser.add_argument("--retry", help="Retry failed query (default 1). Each retry increases the temperature by 0.1", action="store", default=1)
+
+    # Add default arguments from crytic-compile
+    cryticparser.init(parser)
+
+    codex.init_parser(parser, always_enable_codex=True)
+
+    return parser.parse_args()
+
+
+def _use_tab(char: str) -> Optional[bool]:
+    """
+    Check if the char is a tab
+
+    Args:
+        char:
+
+    Returns:
+
+    """
+    if char == " ":
+        return False
+    if char == "\t":
+        return True
+    return None
+
+
+def _post_processesing(
+    answer: str, starting_column: int, use_tab: Optional[bool], force_and_stopped: bool
+) -> Optional[str]:
+    """
+    Clean answers from codex
+
+    Args:
+        answer:
+        starting_column:
+
+    Returns:
+
+    """
+    if answer.count("/**") != 1:
+        return None
+    # Sometimes codex will miss the */, even if it finished properly the request
+    # In this case, we allow slither-documentation to force the */
+    if answer.count("*/") != 1:
+        if force_and_stopped:
+            answer += "*/"
+        else:
+            return None
+    if answer.find("/**") > answer.find("*/"):
+        return None
+    answer = answer[answer.find("/**") : answer.find("*/") + 2]
+    answer_lines = answer.splitlines()
+    # Add indentation to all the lines, aside the first one
+
+    space_char = "\t" if use_tab else " "
+
+    if len(answer_lines) > 0:
+        answer = (
+            answer_lines[0]
+            + "\n"
+            + "\n".join(
+                [space_char * (starting_column - 1) + line for line in answer_lines[1:] if line]
+            )
+        )
+        answer += "\n" + space_char * (starting_column - 1)
+        return answer
+    return answer_lines[0]
+
+
+def _handle_codex(
+    answer: Dict, starting_column: int, use_tab: Optional[bool], force: bool
+) -> Optional[str]:
+    if "choices" in answer:
+        if answer["choices"]:
+            if "text" in answer["choices"][0]:
+                has_stopped = answer["choices"][0].get("finish_reason", "") == "stop"
+                answer_processed = _post_processesing(
+                    answer["choices"][0]["text"], starting_column, use_tab, force and has_stopped
+                )
+                if answer_processed is None:
+                    return None
+                return answer_processed
+    return None
+
+
+# pylint: disable=too-many-locals
+def _handle_compilation_unit(
+    slither: Slither,
+    compilation_unit: SlitherCompilationUnit,
+    overwrite: bool,
+    force: bool,
+    retry: int,
+) -> None:
+
+    logging_file = str(uuid.uuid4())
+
+    for scope in compilation_unit.scopes.values():
+
+        # TODO remove hardcoded filtering
+        if (
+            ".t.sol" in scope.filename.absolute
+            or "mock" in scope.filename.absolute.lower()
+            or "test" in scope.filename.absolute.lower()
+        ):
+            continue
+
+        functions_target: List[Function] = []
+
+        for contract in scope.contracts.values():
+            functions_target += contract.functions_declared
+
+        functions_target += list(scope.functions)
+
+        all_patches: Dict = {}
+
+        for function in functions_target:
+
+            if function.source_mapping.is_dependency or function.has_documentation or function.is_constructor_variables:
+                continue
+            prompt = (
+                "Create a natpsec documentation for this solidity code with only notice and dev.\n"
+            )
+            src_mapping = function.source_mapping
+            content = compilation_unit.core.source_code[src_mapping.filename.absolute]
+            start = src_mapping.start
+            end = src_mapping.start + src_mapping.length
+            prompt += content[start:end]
+
+            use_tab = _use_tab(content[start - 1])
+            if use_tab is None and src_mapping.starting_column > 1:
+                logger.info(f"Non standard space indentation found {content[start-1:end]}")
+                if overwrite:
+                    logger.info("Disable overwrite to avoid mistakes")
+                    overwrite = False
+
+            openai = codex.openai_module()  # type: ignore
+            if openai is None:
+                return
+
+            if slither.codex_log:
+                codex.log_codex(logging_file, "Q: " + prompt)
+
+            tentative = 0
+            answer_processed: Optional[str] = None
+            while tentative < retry:
+                tentative += 1
+
+                answer = openai.Completion.create(  # type: ignore
+                    prompt=prompt,
+                    model=slither.codex_model,
+                    temperature=min(slither.codex_temperature + tentative*0.1, 1),
+                    max_tokens=slither.codex_max_tokens,
+                )
+
+                if slither.codex_log:
+                    codex.log_codex(logging_file, "A: " + str(answer))
+
+                answer_processed = _handle_codex(
+                    answer, src_mapping.starting_column, use_tab, force
+                )
+                if answer_processed:
+                    break
+
+                logger.info(
+                    f"Codex could not generate a well formatted answer for {function.canonical_name}"
+                )
+                logger.info(answer)
+
+            if not answer_processed:
+                continue
+
+            create_patch(
+                all_patches, src_mapping.filename.absolute, start, start, "", answer_processed
+            )
+
+        # all_patches["patches"] should have only 1 file
+        if "patches" not in all_patches:
+            continue
+        for file in all_patches["patches"]:
+            original_txt = compilation_unit.core.source_code[file].encode("utf8")
+            patched_txt = original_txt
+
+            patches = all_patches["patches"][file]
+            offset = 0
+            patches.sort(key=lambda x: x["start"])
+
+            for patch in patches:
+                patched_txt, offset = apply_patch(patched_txt, patch, offset)
+
+            if overwrite:
+                with open(file, "w", encoding="utf8") as f:
+                    f.write(patched_txt.decode("utf8"))
+            else:
+                diff = create_diff(compilation_unit, original_txt, patched_txt, file)
+                with open(f"{file}.patch", "w", encoding="utf8") as f:
+                    f.write(diff)
+
+
+def main() -> None:
+    args = parse_args()
+
+    logger.info("This tool is a WIP, use it with cautious")
+    logger.info("Be aware of OpenAI ToS: https://openai.com/api/policies/terms/")
+    slither = Slither(args.project, **vars(args))
+
+    for compilation_unit in slither.compilation_units:
+        _handle_compilation_unit(
+            slither, compilation_unit, args.overwrite, args.force_answer_parsing, int(args.retry)
+        )
+
+
+if __name__ == "__main__":
+    main()

slither/utils/codex.py

@@ -1,10 +1,70 @@
 import logging
 import os
+from argparse import ArgumentParser
 from pathlib import Path
 
+from slither.utils.command_line import defaults_flag_in_config
+
 logger = logging.getLogger("Slither")
 
 
+def init_parser(parser: ArgumentParser, always_enable_codex: bool = False) -> None:
+    """
+    Init the cli arg with codex features
+
+    Args:
+        parser:
+        always_enable_codex (Optional(bool)): if true, --codex is not enabled
+
+    Returns:
+
+    """
+    group_codex = parser.add_argument_group("Codex (https://beta.openai.com/docs/guides/code)")
+
+    if not always_enable_codex:
+        group_codex.add_argument(
+            "--codex",
+            help="Enable codex (require an OpenAI API Key)",
+            action="store_true",
+            default=defaults_flag_in_config["codex"],
+        )
+
+    group_codex.add_argument(
+        "--codex-log",
+        help="Log codex queries (in crytic_export/codex/)",
+        action="store_true",
+        default=False,
+    )
+
+    group_codex.add_argument(
+        "--codex-contracts",
+        help="Comma separated list of contracts to submit to OpenAI Codex",
+        action="store",
+        default=defaults_flag_in_config["codex_contracts"],
+    )
+
+    group_codex.add_argument(
+        "--codex-model",
+        help="Name of the Codex model to use (affects pricing).  Defaults to 'text-davinci-003'",
+        action="store",
+        default=defaults_flag_in_config["codex_model"],
+    )
+
+    group_codex.add_argument(
+        "--codex-temperature",
+        help="Temperature to use with Codex.  Lower number indicates a more precise answer while higher numbers return more creative answers.  Defaults to 0",
+        action="store",
+        default=defaults_flag_in_config["codex_temperature"],
+    )
+
+    group_codex.add_argument(
+        "--codex-max-tokens",
+        help="Maximum amount of tokens to use on the response.  This number plus the size of the prompt can be no larger than the limit (4097 for text-davinci-003)",
+        action="store",
+        default=defaults_flag_in_config["codex_max_tokens"],
+    )
+
+
 # TODO: investigate how to set the correct return type
 # So that the other modules can work with openai
 def openai_module():  # type: ignore