diff --git a/setup.py b/setup.py index 510efd097..03fe64c42 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,7 @@ setup( "deepdiff", "numpy", "solc-select>=v1.0.0b1", + "openai", ] }, license="AGPL-3.0", diff --git a/slither/__main__.py b/slither/__main__.py index 70357586e..75707af06 100644 --- a/slither/__main__.py +++ b/slither/__main__.py @@ -166,7 +166,6 @@ def process_from_asts( def get_detectors_and_printers() -> Tuple[ List[Type[AbstractDetector]], List[Type[AbstractPrinter]] ]: - detectors_ = [getattr(all_detectors, name) for name in dir(all_detectors)] detectors = [d for d in detectors_ if inspect.isclass(d) and issubclass(d, AbstractDetector)] @@ -286,7 +285,6 @@ def parse_filter_paths(args: argparse.Namespace) -> List[str]: def parse_args( detector_classes: List[Type[AbstractDetector]], printer_classes: List[Type[AbstractPrinter]] ) -> argparse.Namespace: - usage = "slither target [flag]\n" usage += "\ntarget can be:\n" usage += "\t- file.sol // a Solidity file\n" @@ -316,6 +314,7 @@ 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", @@ -556,6 +555,48 @@ 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"], + ) + # debugger command parser.add_argument("--debug", help=argparse.SUPPRESS, action="store_true", default=False) diff --git a/slither/detectors/all_detectors.py b/slither/detectors/all_detectors.py index 2c8d24428..e8c8334b7 100644 --- a/slither/detectors/all_detectors.py +++ b/slither/detectors/all_detectors.py @@ -85,3 +85,4 @@ from .statements.msg_value_in_loop import MsgValueInLoop from .statements.delegatecall_in_loop import DelegatecallInLoop from .functions.protected_variable import ProtectedVariables from .functions.permit_domain_signature_collision import DomainSeparatorCollision +from .functions.codex import Codex diff --git a/slither/detectors/functions/codex.py b/slither/detectors/functions/codex.py new file mode 100644 index 000000000..fb00f64c0 --- /dev/null +++ b/slither/detectors/functions/codex.py @@ -0,0 +1,136 @@ +import logging +import uuid +from typing import List, Union + +from slither.detectors.abstract_detector import AbstractDetector, DetectorClassification +from slither.utils import codex +from slither.utils.output import Output, SupportedOutput + +logger = logging.getLogger("Slither") + +VULN_FOUND = "VULN_FOUND" + + +class Codex(AbstractDetector): + """ + Use codex to detect vulnerability + """ + + ARGUMENT = "codex" + HELP = "Use Codex to find vulnerabilities." + IMPACT = DetectorClassification.HIGH + CONFIDENCE = DetectorClassification.LOW + + WIKI = "https://github.com/crytic/slither/wiki/Detector-Documentation#codex" + + WIKI_TITLE = "Codex" + WIKI_DESCRIPTION = "Use [codex](https://openai.com/blog/openai-codex/) to find vulnerabilities" + + # region wiki_exploit_scenario + WIKI_EXPLOIT_SCENARIO = """N/A""" + # endregion wiki_exploit_scenario + + WIKI_RECOMMENDATION = "Review codex's message." + + def _run_codex(self, logging_file: str, prompt: str) -> str: + """ + Handle the codex logic + + Args: + logging_file (str): file where to log the queries + prompt (str): prompt to send to codex + + Returns: + codex answer (str) + """ + openai_module = codex.openai_module() # type: ignore + if openai_module is None: + return "" + + if self.slither.codex_log: + codex.log_codex(logging_file, "Q: " + prompt) + + answer = "" + res = {} + try: + res = openai_module.Completion.create( + prompt=prompt, + model=self.slither.codex_model, + temperature=self.slither.codex_temperature, + max_tokens=self.slither.codex_max_tokens, + ) + except Exception as e: # pylint: disable=broad-except + logger.info("OpenAI request failed: " + str(e)) + + # """ OpenAI completion response shape example: + # { + # "choices": [ + # { + # "finish_reason": "stop", + # "index": 0, + # "logprobs": null, + # "text": "VULNERABILITIES:. The withdraw() function does not check..." + # } + # ], + # "created": 1670357537, + # "id": "cmpl-6KYaXdA6QIisHlTMM7RCJ1nR5wTKx", + # "model": "text-davinci-003", + # "object": "text_completion", + # "usage": { + # "completion_tokens": 80, + # "prompt_tokens": 249, + # "total_tokens": 329 + # } + # } """ + + if res: + if self.slither.codex_log: + codex.log_codex(logging_file, "A: " + str(res)) + else: + codex.log_codex(logging_file, "A: Codex failed") + + if res.get("choices", []) and VULN_FOUND in res["choices"][0].get("text", ""): + # remove VULN_FOUND keyword and cleanup + answer = ( + res["choices"][0]["text"] + .replace(VULN_FOUND, "") + .replace("\n", "") + .replace(": ", "") + ) + return answer + + def _detect(self) -> List[Output]: + results: List[Output] = [] + + if not self.slither.codex_enabled: + return [] + + logging_file = str(uuid.uuid4()) + + for contract in self.compilation_unit.contracts: + if ( + self.slither.codex_contracts != "all" + and contract.name not in self.slither.codex_contracts.split(",") + ): + continue + prompt = f"Analyze this Solidity contract and find the vulnerabilities. If you find any vulnerabilities, begin the response with {VULN_FOUND}\n" + src_mapping = contract.source_mapping + content = contract.compilation_unit.core.source_code[src_mapping.filename.absolute] + start = src_mapping.start + end = src_mapping.start + src_mapping.length + prompt += content[start:end] + + answer = self._run_codex(logging_file, prompt) + + if answer: + info: List[Union[str, SupportedOutput]] = [ + "Codex detected a potential bug in ", + contract, + "\n", + answer, + "\n", + ] + + new_result = self.generate_result(info) + results.append(new_result) + return results diff --git a/slither/slither.py b/slither/slither.py index dcfc0ad7e..81e920d01 100644 --- a/slither/slither.py +++ b/slither/slither.py @@ -83,6 +83,14 @@ class Slither(SlitherCore): # pylint: disable=too-many-instance-attributes self.line_prefix = kwargs.get("change_line_prefix", "#") + # Indicate if Codex related features should be used + self.codex_enabled = kwargs.get("codex", False) + self.codex_contracts = kwargs.get("codex_contracts", "all") + self.codex_model = kwargs.get("codex_model", "text-davinci-003") + self.codex_temperature = kwargs.get("codex_temperature", 0) + self.codex_max_tokens = kwargs.get("codex_max_tokens", 300) + self.codex_log = kwargs.get("codex_log", False) + self._parsers: List[SlitherCompilationUnitSolc] = [] try: if isinstance(target, CryticCompile): diff --git a/slither/utils/codex.py b/slither/utils/codex.py new file mode 100644 index 000000000..0040fb03c --- /dev/null +++ b/slither/utils/codex.py @@ -0,0 +1,53 @@ +import logging +import os +from pathlib import Path + +logger = logging.getLogger("Slither") + + +# TODO: investigate how to set the correct return type +# So that the other modules can work with openai +def openai_module(): # type: ignore + """ + Return the openai module + Consider checking the usage of open (slither.codex_enabled) before using this function + + Returns: + Optional[the openai module] + """ + try: + # pylint: disable=import-outside-toplevel + import openai + + api_key = os.getenv("OPENAI_API_KEY") + if api_key is None: + logger.info( + "Please provide an Open API Key in OPENAI_API_KEY (https://beta.openai.com/account/api-keys)" + ) + return None + openai.api_key = api_key + except ImportError: + logger.info("OpenAI was not installed") # type: ignore + logger.info('run "pip install openai"') + return None + return openai + + +def log_codex(filename: str, prompt: str) -> None: + """ + Log the prompt in crytic/export/codex/filename + Append to the file + + Args: + filename: filename to write to + prompt: prompt to write + + Returns: + None + """ + + Path("crytic_export/codex").mkdir(parents=True, exist_ok=True) + + with open(Path("crytic_export/codex", filename), "a", encoding="utf8") as file: + file.write(prompt) + file.write("\n") diff --git a/slither/utils/command_line.py b/slither/utils/command_line.py index c2fef5eca..71305c56e 100644 --- a/slither/utils/command_line.py +++ b/slither/utils/command_line.py @@ -29,6 +29,12 @@ JSON_OUTPUT_TYPES = [ # Those are the flags shared by the command line and the config file defaults_flag_in_config = { + "codex": False, + "codex_contracts": "all", + "codex_model": "text-davinci-003", + "codex_temperature": 0, + "codex_max_tokens": 300, + "codex_log": False, "detectors_to_run": "all", "printers_to_run": None, "detectors_to_exclude": None,