diff --git a/setup.py b/setup.py index 86db4fa9a..89adf7f4a 100644 --- a/setup.py +++ b/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", ] }, ) diff --git a/slither/__main__.py b/slither/__main__.py index 75707af06..d8fce1ce5 100644 --- a/slither/__main__.py +++ b/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", - 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"], - ) + codex.init_parser(parser) # debugger command parser.add_argument("--debug", help=argparse.SUPPRESS, action="store_true", default=False) diff --git a/slither/core/declarations/function.py b/slither/core/declarations/function.py index a4624feec..2fdea7210 100644 --- a/slither/core/declarations/function.py +++ b/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 diff --git a/slither/solc_parsing/declarations/function.py b/slither/solc_parsing/declarations/function.py index 130375211..40ddba07e 100644 --- a/slither/solc_parsing/declarations/function.py +++ b/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 diff --git a/slither/tools/documentation/README.md b/slither/tools/documentation/README.md new file mode 100644 index 000000000..2ed90692c --- /dev/null +++ b/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) + diff --git a/slither/tools/documentation/__init__.py b/slither/tools/documentation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/slither/tools/documentation/__main__.py b/slither/tools/documentation/__main__.py new file mode 100644 index 000000000..ea45ed8e1 --- /dev/null +++ b/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() diff --git a/slither/utils/codex.py b/slither/utils/codex.py index 0040fb03c..3b06efe6f 100644 --- a/slither/utils/codex.py +++ b/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