From c02231ff1c42261ac573f06e85e78042ecda9e66 Mon Sep 17 00:00:00 2001 From: Josselin Feist Date: Tue, 6 Dec 2022 13:39:51 +0100 Subject: [PATCH 1/4] Add Codex vuln detector The detector requires: - The user to use the flag `--codex` (meaning that codex is not ran by default) - `openai` must be installed, and `OPENAI_API_KEY` set The detector works at the contract level, and send the whole contract body to codex --- setup.py | 1 + slither/__main__.py | 4 ++ slither/detectors/all_detectors.py | 1 + slither/detectors/functions/codex.py | 80 ++++++++++++++++++++++++++++ slither/slither.py | 3 ++ 5 files changed, 89 insertions(+) create mode 100644 slither/detectors/functions/codex.py diff --git a/setup.py b/setup.py index 510efd0971..b29244713c 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,7 @@ "deepdiff", "numpy", "solc-select>=v1.0.0b1", + "openai" ] }, license="AGPL-3.0", diff --git a/slither/__main__.py b/slither/__main__.py index 70357586e9..00ea308d2a 100644 --- a/slither/__main__.py +++ b/slither/__main__.py @@ -301,6 +301,10 @@ def parse_args( parser.add_argument("filename", help=argparse.SUPPRESS) + parser.add_argument( + "--codex", help="Enable codex (require an OpenAI API Key)", action="store_true", default=False + ) + cryticparser.init(parser) parser.add_argument( diff --git a/slither/detectors/all_detectors.py b/slither/detectors/all_detectors.py index 2c8d244281..e8c8334b72 100644 --- a/slither/detectors/all_detectors.py +++ b/slither/detectors/all_detectors.py @@ -85,3 +85,4 @@ 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 0000000000..daed064252 --- /dev/null +++ b/slither/detectors/functions/codex.py @@ -0,0 +1,80 @@ +import logging +import os +from typing import List + +from slither.detectors.abstract_detector import AbstractDetector, DetectorClassification +from slither.utils.output import Output + +logger = logging.getLogger("Slither") + + +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 _detect(self) -> List[Output]: + results: List[Output] = [] + + if not self.slither.codex_enabled: + return [] + + try: + # pylint: disable=import-outside-toplevel + import openai + except ImportError: + logging.info("OpenAI was not installed") + logging.info('run "pip install openai"') + return [] + + api_key = os.getenv("OPENAI_API_KEY") + if api_key is None: + logging.info( + "Please provide an Open API Key in OPENAI_API_KEY (https://beta.openai.com/account/api-keys)" + ) + return [] + openai.api_key = api_key + + for contract in self.compilation_unit.contracts: + prompt = "Is there a vulnerability in this solidity contracts?\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 = openai.Completion.create( # type: ignore + model="text-davinci-003", prompt=prompt, temperature=0, max_tokens=200 + ) + + if "choices" in answer: + if answer["choices"]: + if "text" in answer["choices"][0]: + if "Yes," in answer["choices"][0]["text"]: + info = [ + "Codex detected a potential bug in ", + contract, + "\n", + answer["choices"][0]["text"], + "\n", + ] + + res = self.generate_result(info) + results.append(res) + + return results diff --git a/slither/slither.py b/slither/slither.py index dcfc0ad7e5..a61e8255ff 100644 --- a/slither/slither.py +++ b/slither/slither.py @@ -83,6 +83,9 @@ def __init__(self, target: Union[str, CryticCompile], **kwargs): self.line_prefix = kwargs.get("change_line_prefix", "#") + # Indicate if codex-related features should be used + self.codex_enabled = kwargs.get("codex", False) + self._parsers: List[SlitherCompilationUnitSolc] = [] try: if isinstance(target, CryticCompile): From 5763c7431bcf4190b3b59f9b5c6de07777b4d31f Mon Sep 17 00:00:00 2001 From: Josselin Feist Date: Tue, 6 Dec 2022 13:58:16 +0100 Subject: [PATCH 2/4] black --- setup.py | 2 +- slither/__main__.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index b29244713c..03fe64c426 100644 --- a/setup.py +++ b/setup.py @@ -26,7 +26,7 @@ "deepdiff", "numpy", "solc-select>=v1.0.0b1", - "openai" + "openai", ] }, license="AGPL-3.0", diff --git a/slither/__main__.py b/slither/__main__.py index 00ea308d2a..8e32a1c433 100644 --- a/slither/__main__.py +++ b/slither/__main__.py @@ -302,7 +302,10 @@ def parse_args( parser.add_argument("filename", help=argparse.SUPPRESS) parser.add_argument( - "--codex", help="Enable codex (require an OpenAI API Key)", action="store_true", default=False + "--codex", + help="Enable codex (require an OpenAI API Key)", + action="store_true", + default=False, ) cryticparser.init(parser) From 00d33c6e781f05e6f279a4c07fe5754a457c03a0 Mon Sep 17 00:00:00 2001 From: Richie Date: Tue, 6 Dec 2022 13:56:38 -0800 Subject: [PATCH 3/4] feat: parameterize OpenAI inputs --- slither/__main__.py | 30 ++++++++++- slither/detectors/functions/codex.py | 75 ++++++++++++++++++++-------- slither/slither.py | 6 ++- slither/utils/command_line.py | 5 ++ 4 files changed, 94 insertions(+), 22 deletions(-) diff --git a/slither/__main__.py b/slither/__main__.py index 8e32a1c433..9426a6de9d 100644 --- a/slither/__main__.py +++ b/slither/__main__.py @@ -305,7 +305,35 @@ def parse_args( "--codex", help="Enable codex (require an OpenAI API Key)", action="store_true", - default=False, + default=defaults_flag_in_config["codex"], + ) + + parser.add_argument( + "--codex-contracts", + help="Comma separated list of contracts to submit to OpenAI Codex", + action="store", + default=defaults_flag_in_config["codex_contracts"], + ) + + parser.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"], + ) + + parser.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"], + ) + + parser.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"], ) cryticparser.init(parser) diff --git a/slither/detectors/functions/codex.py b/slither/detectors/functions/codex.py index daed064252..4dca447756 100644 --- a/slither/detectors/functions/codex.py +++ b/slither/detectors/functions/codex.py @@ -7,6 +7,7 @@ logger = logging.getLogger("Slither") +VULN_FOUND = "VULN_FOUND" class Codex(AbstractDetector): """ @@ -52,29 +53,63 @@ def _detect(self) -> List[Output]: openai.api_key = api_key for contract in self.compilation_unit.contracts: - prompt = "Is there a vulnerability in this solidity contracts?\n" + if self.slither.codex_contracts != "all" and contract.name not in self.slither.codex_contracts.split(","): + continue + prompt = "Analyze this Solidity contract and find the vulnerabilities. If you find any vulnerabilities, begin the response with {}".format(VULN_FOUND) 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 = openai.Completion.create( # type: ignore - model="text-davinci-003", prompt=prompt, temperature=0, max_tokens=200 - ) - - if "choices" in answer: - if answer["choices"]: - if "text" in answer["choices"][0]: - if "Yes," in answer["choices"][0]["text"]: - info = [ - "Codex detected a potential bug in ", - contract, - "\n", - answer["choices"][0]["text"], - "\n", - ] - - res = self.generate_result(info) - results.append(res) - + logging.info("Querying OpenAI") + print("Querying OpenAI") + answer = "" + res = {} + try: + res = openai.Completion.create( # type: ignore + prompt=prompt, + model=self.slither.codex_model, + temperature=self.slither.codex_temperature, + max_tokens=self.slither.codex_max_tokens, + ) + except Exception as e: + print("OpenAI request failed: " + str(e)) + logging.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 len(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(": ", "") + + if len(answer): + info = [ + "Codex detected a potential bug in ", + contract, + "\n", + answer, + "\n", + ] + + res = self.generate_result(info) + results.append(res) return results diff --git a/slither/slither.py b/slither/slither.py index a61e8255ff..0b1f57a37e 100644 --- a/slither/slither.py +++ b/slither/slither.py @@ -83,8 +83,12 @@ def __init__(self, target: Union[str, CryticCompile], **kwargs): self.line_prefix = kwargs.get("change_line_prefix", "#") - # Indicate if codex-related features should be used + # Indicate if Codex related features should be used self.codex_enabled = kwargs.get("codex", False) + self.codex_contracts = kwargs.get("codex_contracts") + self.codex_model = kwargs.get("codex_model") + self.codex_temperature = kwargs.get("codex_temperature") + self.codex_max_tokens = kwargs.get("codex_max_tokens") self._parsers: List[SlitherCompilationUnitSolc] = [] try: diff --git a/slither/utils/command_line.py b/slither/utils/command_line.py index c2fef5eca0..f774437a1a 100644 --- a/slither/utils/command_line.py +++ b/slither/utils/command_line.py @@ -29,6 +29,11 @@ # 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, "detectors_to_run": "all", "printers_to_run": None, "detectors_to_exclude": None, From f62433bd50fbc8d95a0a328b67a051cc305b6498 Mon Sep 17 00:00:00 2001 From: Josselin Feist Date: Wed, 7 Dec 2022 11:07:43 +0100 Subject: [PATCH 4/4] Create utils.codex Refactor functions/codex Minor improvements --- slither/__main__.py | 80 +++++++------- slither/detectors/functions/codex.py | 149 +++++++++++++++------------ slither/slither.py | 9 +- slither/utils/codex.py | 53 ++++++++++ slither/utils/command_line.py | 1 + 5 files changed, 187 insertions(+), 105 deletions(-) create mode 100644 slither/utils/codex.py diff --git a/slither/__main__.py b/slither/__main__.py index 9426a6de9d..75707af06b 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" @@ -301,41 +299,6 @@ def parse_args( parser.add_argument("filename", help=argparse.SUPPRESS) - parser.add_argument( - "--codex", - help="Enable codex (require an OpenAI API Key)", - action="store_true", - default=defaults_flag_in_config["codex"], - ) - - parser.add_argument( - "--codex-contracts", - help="Comma separated list of contracts to submit to OpenAI Codex", - action="store", - default=defaults_flag_in_config["codex_contracts"], - ) - - parser.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"], - ) - - parser.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"], - ) - - parser.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"], - ) - cryticparser.init(parser) parser.add_argument( @@ -351,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", @@ -591,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/functions/codex.py b/slither/detectors/functions/codex.py index 4dca447756..fb00f64c04 100644 --- a/slither/detectors/functions/codex.py +++ b/slither/detectors/functions/codex.py @@ -1,14 +1,16 @@ import logging -import os -from typing import List +import uuid +from typing import List, Union from slither.detectors.abstract_detector import AbstractDetector, DetectorClassification -from slither.utils.output import Output +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 @@ -30,79 +32,98 @@ class Codex(AbstractDetector): 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 [] - try: - # pylint: disable=import-outside-toplevel - import openai - except ImportError: - logging.info("OpenAI was not installed") - logging.info('run "pip install openai"') - return [] - - api_key = os.getenv("OPENAI_API_KEY") - if api_key is None: - logging.info( - "Please provide an Open API Key in OPENAI_API_KEY (https://beta.openai.com/account/api-keys)" - ) - return [] - openai.api_key = api_key + 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(","): + if ( + self.slither.codex_contracts != "all" + and contract.name not in self.slither.codex_contracts.split(",") + ): continue - prompt = "Analyze this Solidity contract and find the vulnerabilities. If you find any vulnerabilities, begin the response with {}".format(VULN_FOUND) + 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] - logging.info("Querying OpenAI") - print("Querying OpenAI") - answer = "" - res = {} - try: - res = openai.Completion.create( # type: ignore - prompt=prompt, - model=self.slither.codex_model, - temperature=self.slither.codex_temperature, - max_tokens=self.slither.codex_max_tokens, - ) - except Exception as e: - print("OpenAI request failed: " + str(e)) - logging.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 len(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(": ", "") - - if len(answer): - info = [ + + answer = self._run_codex(logging_file, prompt) + + if answer: + info: List[Union[str, SupportedOutput]] = [ "Codex detected a potential bug in ", contract, "\n", @@ -110,6 +131,6 @@ def _detect(self) -> List[Output]: "\n", ] - res = self.generate_result(info) - results.append(res) + new_result = self.generate_result(info) + results.append(new_result) return results diff --git a/slither/slither.py b/slither/slither.py index 0b1f57a37e..81e920d013 100644 --- a/slither/slither.py +++ b/slither/slither.py @@ -85,10 +85,11 @@ def __init__(self, target: Union[str, CryticCompile], **kwargs): # Indicate if Codex related features should be used self.codex_enabled = kwargs.get("codex", False) - self.codex_contracts = kwargs.get("codex_contracts") - self.codex_model = kwargs.get("codex_model") - self.codex_temperature = kwargs.get("codex_temperature") - self.codex_max_tokens = kwargs.get("codex_max_tokens") + 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: diff --git a/slither/utils/codex.py b/slither/utils/codex.py new file mode 100644 index 0000000000..0040fb03c5 --- /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 f774437a1a..71305c56e2 100644 --- a/slither/utils/command_line.py +++ b/slither/utils/command_line.py @@ -34,6 +34,7 @@ "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,