Add detectors_to_include to override exclude args

crytic · Apr 26, 2024 · e2daa8a · e2daa8a
1 parent ded705d
commit e2daa8a
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 29 deletions.
diff --git a/slither/__main__.py b/slither/__main__.py
@@ -11,7 +11,7 @@
 import sys
 import traceback
 from importlib import metadata
-from typing import Tuple, Optional, List, Dict, Type, Union, Any, Sequence
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Type, Union
 
 
 from crytic_compile import cryticparser, CryticCompile, InvalidCompilation
@@ -211,47 +211,54 @@ def choose_detectors(
 
     if args.detectors_to_run == "all":
         detectors_to_run = all_detector_classes
-        if args.detectors_to_exclude:
-            detectors_excluded = args.detectors_to_exclude.split(",")
-            for detector in detectors:
-                if detector in detectors_excluded:
-                    detectors_to_run.remove(detectors[detector])
     else:
-        for detector in args.detectors_to_run.split(","):
-            if detector in detectors:
-                detectors_to_run.append(detectors[detector])
-            else:
-                raise ValueError(f"Error: {detector} is not a detector")
-        detectors_to_run = sorted(detectors_to_run, key=lambda x: x.IMPACT)
+        detectors_to_run = __include_detectors(
+            set(detectors_to_run), args.detectors_to_run, detectors
+        )
         return detectors_to_run
 
-    if args.exclude_optimization:
-        detectors_to_run = [
-            d for d in detectors_to_run if d.IMPACT != DetectorClassification.OPTIMIZATION
-        ]
+    classification_map = {
+        DetectorClassification.HIGH: args.exclude_high,
+        DetectorClassification.MEDIUM: args.exclude_medium,
+        DetectorClassification.LOW: args.exclude_low,
+        DetectorClassification.INFORMATIONAL: args.exclude_informational,
+        DetectorClassification.OPTIMIZATION: args.exclude_optimization,
+    }
+    excluded_classification = [
+        classification for classification, included in classification_map.items() if included
+    ]
+    detectors_to_run = [d for d in detectors_to_run if d.IMPACT not in excluded_classification]
 
-    if args.exclude_informational:
-        detectors_to_run = [
-            d for d in detectors_to_run if d.IMPACT != DetectorClassification.INFORMATIONAL
-        ]
-    if args.exclude_low:
-        detectors_to_run = [d for d in detectors_to_run if d.IMPACT != DetectorClassification.LOW]
-    if args.exclude_medium:
-        detectors_to_run = [
-            d for d in detectors_to_run if d.IMPACT != DetectorClassification.MEDIUM
-        ]
-    if args.exclude_high:
-        detectors_to_run = [d for d in detectors_to_run if d.IMPACT != DetectorClassification.HIGH]
     if args.detectors_to_exclude:
         detectors_to_run = [
             d for d in detectors_to_run if d.ARGUMENT not in args.detectors_to_exclude
         ]
 
-    detectors_to_run = sorted(detectors_to_run, key=lambda x: x.IMPACT)
+    if args.detectors_to_include:
+        detectors_to_run = __include_detectors(
+            set(detectors_to_run), args.detectors_to_include, detectors
+        )
 
     return detectors_to_run
 
 
+def __include_detectors(
+    detectors_to_run: Set[Type[AbstractDetector]],
+    detectors_to_include: str,
+    detectors: Dict[str, Type[AbstractDetector]],
+) -> List[Type[AbstractDetector]]:
+    include_detectors = detectors_to_include.split(",")
+
+    for detector in include_detectors:
+        if detector in detectors:
+            detectors_to_run.add(detectors[detector])
+        else:
+            raise ValueError(f"Error: {detector} is not a detector")
+
+    detectors_to_run = sorted(detectors_to_run, key=lambda x: x.IMPACT)
+    return detectors_to_run
+
+
 def choose_printers(
     args: argparse.Namespace, all_printer_classes: List[Type[AbstractPrinter]]
 ) -> List[Type[AbstractPrinter]]:
@@ -407,6 +414,14 @@ def parse_args(
         default=defaults_flag_in_config["exclude_high"],
     )
 
+    group_detector.add_argument(
+        "--include-detectors",
+        help="Comma-separated list of detectors that should be included",
+        action="store",
+        dest="detectors_to_include",
+        default=defaults_flag_in_config["detectors_to_include"],
+    )
+
     fail_on_group = group_detector.add_mutually_exclusive_group()
     fail_on_group.add_argument(
         "--fail-pedantic",

diff --git a/slither/utils/command_line.py b/slither/utils/command_line.py
@@ -48,6 +48,7 @@ class FailOnLevel(enum.Enum):
     "detectors_to_run": "all",
     "printers_to_run": None,
     "detectors_to_exclude": None,
+    "detectors_to_include": None,
     "exclude_dependencies": False,
     "exclude_informational": False,
     "exclude_optimization": False,