Merge branch 'main' into dependency-error-messages

modelscan/issues.py

@@ -142,3 +142,6 @@ def output_json(self) -> Dict[str, str]:
             "source": f"{str(self.source)}",
             "scanner": f"{self.scanner}",
         }
+
+    def __repr__(self) -> str:
+        return f"<OperatorIssueDetails(module={self.module}, operator={self.operator}, source={str(self.source)})>"

modelscan/settings.py

@@ -61,6 +61,7 @@
                     "exec",
                     "open",
                     "breakpoint",
+                    "__import__",
                 ],  # Pickle versions 0, 1, 2 have those function under '__builtin__'
                 "builtins": [
                     "eval",
@@ -70,6 +71,7 @@
                     "exec",
                     "open",
                     "breakpoint",
+                    "__import__",
                 ],  # Pickle versions 3, 4 have those function under 'builtins'
                 "runpy": "*",
                 "os": "*",
@@ -78,7 +80,11 @@
                 "socket": "*",
                 "subprocess": "*",
                 "sys": "*",
-                "operator": "attrgetter",  # Ex of code execution: operator.attrgetter("system")(__import__("os"))("echo pwned")
+                "operator": [
+                    "attrgetter",  # Ex of code execution: operator.attrgetter("system")(__import__("os"))("echo pwned")
+                ],
+                "pty": "*",
+                "pickle": "*",
             },
             "HIGH": {
                 "webbrowser": "*",  # Includes webbrowser.open()

modelscan/tools/picklescanner.py

@@ -134,45 +134,41 @@ def scan_pickle_bytes(
         )
 
     logger.debug("Global imports in %s: %s", source, raw_globals)
+    severities = {
+        "CRITICAL": IssueSeverity.CRITICAL,
+        "HIGH": IssueSeverity.HIGH,
+        "MEDIUM": IssueSeverity.MEDIUM,
+        "LOW": IssueSeverity.LOW,
+    }
 
     for rg in raw_globals:
         global_module, global_name, severity = rg[0], rg[1], None
-        unsafe_critical_filter = settings["unsafe_globals"]["CRITICAL"].get(
-            global_module
-        )
-        unsafe_high_filter = settings["unsafe_globals"]["HIGH"].get(global_module)
-        unsafe_medium_filter = settings["unsafe_globals"]["MEDIUM"].get(global_module)
-        unsafe_low_filter = settings["unsafe_globals"]["LOW"].get(global_module)
-        if unsafe_critical_filter is not None and (
-            unsafe_critical_filter == "*" or global_name in unsafe_critical_filter
-        ):
-            severity = IssueSeverity.CRITICAL
-
-        elif unsafe_high_filter is not None and (
-            unsafe_high_filter == "*" or global_name in unsafe_high_filter
-        ):
-            severity = IssueSeverity.HIGH
-        elif unsafe_medium_filter is not None and (
-            unsafe_medium_filter == "*" or global_name in unsafe_medium_filter
-        ):
-            severity = IssueSeverity.MEDIUM
-        elif unsafe_low_filter is not None and (
-            unsafe_low_filter == "*" or global_name in unsafe_low_filter
-        ):
-            severity = IssueSeverity.LOW
-        elif "unknown" in global_module or "unknown" in global_name:
-            severity = IssueSeverity.MEDIUM
-        else:
-            continue
-        issues.append(
-            Issue(
-                code=IssueCode.UNSAFE_OPERATOR,
-                severity=severity,
-                details=OperatorIssueDetails(
-                    module=global_module, operator=global_name, source=source
-                ),
+        for severity_name in severities:
+            if global_module not in settings["unsafe_globals"][severity_name]:
+                continue
+            filter = settings["unsafe_globals"][severity_name][global_module]
+            if filter == "*":
+                severity = severities[severity_name]
+                break
+            for filter_value in filter:
+                if filter_value in global_name:
+                    severity = severities[severity_name]
+                    break
+            else:
+                continue
+            break
+        if "unknown" in global_module or "unknown" in global_name:
+            severity = IssueSeverity.CRITICAL  # we must assume it is RCE
+        if severity is not None:
+            issues.append(
+                Issue(
+                    code=IssueCode.UNSAFE_OPERATOR,
+                    severity=severity,
+                    details=OperatorIssueDetails(
+                        module=global_module, operator=global_name, source=source
+                    ),
+                )
             )
-        )
     return ScanResults(issues, [])
 
 

tests/test_modelscan.py

@@ -80,6 +80,52 @@ def __reduce__(self) -> Any:
         return sys.exit, (0,)
 
 
+def malicious12_gen() -> bytes:
+    p = pickle.PROTO + b"\x05"
+
+    # stack = [pickle.loads]
+    p += pickle.GLOBAL + b"pickle\nloads\n"
+
+    # stack = [pickle.loads, p2]
+    p2 = (
+        pickle.PROTO
+        + b"\x05"
+        + pickle.GLOBAL
+        + b"os\nsystem\n"
+        + pickle.UNICODE
+        + b"echo pwned!!!\n"
+        + pickle.TUPLE1
+        + pickle.REDUCE
+        + pickle.STOP
+    )
+    p += pickle.BINBYTES + len(p2).to_bytes(4, "little") + p2
+
+    # stack = [pickle.loads, (p2,)]
+    p += pickle.TUPLE1
+
+    # stack = [pickle.loads(p2)]
+    p += pickle.REDUCE
+
+    # return None
+    p += pickle.POP
+    p += pickle.NONE
+
+    p += pickle.STOP
+    return p
+
+
+def malicious13_gen() -> bytes:
+    p = pickle.PROTO + b"\x05"
+
+    p += pickle.GLOBAL + b"builtins\neval.__call__\n"
+    p += pickle.UNICODE + b'__import__("os").system("echo pwned!!!")\n'
+    p += pickle.TUPLE1
+    p += pickle.REDUCE
+
+    p += pickle.STOP
+    return p
+
+
 def initialize_pickle_file(path: str, obj: Any, version: int) -> None:
     if not os.path.exists(path):
         with open(path, "wb") as file:
@@ -209,6 +255,10 @@ def file_path(tmp_path_factory: Any) -> Any:
     )
     initialize_data_file(f"{tmp}/data/malicious10.pkl", malicious10_pickle_bytes)
 
+    initialize_data_file(f"{tmp}/data/malicious12.pkl", malicious12_gen())
+
+    initialize_data_file(f"{tmp}/data/malicious13.pkl", malicious13_gen())
+
     return tmp
 
 
@@ -643,6 +693,30 @@ def test_scan_pickle_operators(file_path: Any) -> None:
     malicious11 = ModelScan()
     malicious11.scan(Path(f"{file_path}/data/malicious11.pkl"))
     assert malicious11.issues.all_issues == expected_malicious11
+    expected_malicious12 = [
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "pickle", "loads", f"{file_path}/data/malicious12.pkl"
+            ),
+        )
+    ]
+    malicious12 = ModelScan()
+    malicious12.scan(Path(f"{file_path}/data/malicious12.pkl"))
+    assert malicious12.issues.all_issues == expected_malicious12
+    expected_malicious13 = [
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "builtins", "eval.__call__", f"{file_path}/data/malicious13.pkl"
+            ),
+        )
+    ]
+    malicious13 = ModelScan()
+    malicious13.scan(Path(f"{file_path}/data/malicious13.pkl"))
+    assert malicious13.issues.all_issues == expected_malicious13
 
 
 def test_scan_directory_path(file_path: str) -> None:
@@ -807,6 +881,20 @@ def test_scan_directory_path(file_path: str) -> None:
             IssueSeverity.CRITICAL,
             OperatorIssueDetails("os", "system", f"{file_path}/data/malicious11.pkl"),
         ),
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "pickle", "loads", f"{file_path}/data/malicious12.pkl"
+            ),
+        ),
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "builtins", "eval.__call__", f"{file_path}/data/malicious13.pkl"
+            ),
+        ),
     }
     ms = ModelScan()
     p = Path(f"{file_path}/data/")