diff --git a/slither/detectors/functions/codex.py b/slither/detectors/functions/codex.py index fb00f64c04..6b694155e0 100644 --- a/slither/detectors/functions/codex.py +++ b/slither/detectors/functions/codex.py @@ -52,6 +52,10 @@ def _run_codex(self, logging_file: str, prompt: str) -> str: answer = "" res = {} + + if self.slither.codex_organization: + openai_module.organization = self.slither.codex_organization + try: res = openai_module.Completion.create( prompt=prompt, diff --git a/slither/slither.py b/slither/slither.py index 45d99906fd..227a37365e 100644 --- a/slither/slither.py +++ b/slither/slither.py @@ -1,5 +1,5 @@ import logging -from typing import Union, List, ValuesView, Type, Dict +from typing import Union, List, ValuesView, Type, Dict, Optional from crytic_compile import CryticCompile, InvalidCompilation @@ -90,6 +90,7 @@ def __init__(self, target: Union[str, CryticCompile], **kwargs): 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.codex_organization: Optional[str] = kwargs.get("codex_organization", None) self.no_fail = kwargs.get("no_fail", False) diff --git a/slither/tools/documentation/__main__.py b/slither/tools/documentation/__main__.py index 8e545fb095..39b1eacdac 100644 --- a/slither/tools/documentation/__main__.py +++ b/slither/tools/documentation/__main__.py @@ -36,6 +36,13 @@ def parse_args() -> argparse.Namespace: default=False, ) + parser.add_argument( + "--include-tests", + help="Include the tests", + action="store_true", + default=False, + ) + parser.add_argument( "--retry", help="Retry failed query (default 1). Each retry increases the temperature by 0.1", @@ -202,6 +209,7 @@ def _handle_compilation_unit( overwrite: bool, force: bool, retry: int, + include_test: bool, ) -> None: logging_file: Optional[str] if slither.codex_log: @@ -210,9 +218,8 @@ def _handle_compilation_unit( logging_file = None for scope in compilation_unit.scopes.values(): - # Dont send tests file - if ( + if not include_test and ( ".t.sol" in scope.filename.absolute or "mock" in scope.filename.absolute.lower() or "test" in scope.filename.absolute.lower() @@ -271,6 +278,7 @@ def main() -> None: args.overwrite, args.force_answer_parsing, int(args.retry), + args.include_tests, ) except ImportError: pass diff --git a/slither/utils/codex.py b/slither/utils/codex.py index 3b06efe6f5..8b15656708 100644 --- a/slither/utils/codex.py +++ b/slither/utils/codex.py @@ -64,6 +64,13 @@ def init_parser(parser: ArgumentParser, always_enable_codex: bool = False) -> No default=defaults_flag_in_config["codex_max_tokens"], ) + group_codex.add_argument( + "--codex-organization", + help="Codex organization", + action="store", + default=None, + ) + # TODO: investigate how to set the correct return type # So that the other modules can work with openai