Feat: add vetted properties

.gitignore

@@ -0,0 +1 @@
+__pycache__/

README.md

@@ -65,6 +65,12 @@ You can also safely trace the execution of the Pickle virtual machine without ex
 Finally, you can inject arbitrary Python code that will be run on unpickling into an existing pickle file with the
 `--inject` option.
 
+### Unvetted Dependencies
+
+You can check for unvetted dependencies in `fickling` by giving the `Pickled` class a list of "vetted" function calls from given modules.
+
+See [example/unvetted_dependencies.py](example/unvetted_dependencies.py)
+
 ## License
 
 This utility was developed by [Trail of Bits](https://www.trailofbits.com/).

example/unvetted_dependencies.py

@@ -0,0 +1,33 @@
+from fickling.pickle import Pickled
+import numpy as np
+import pickle
+from sklearn import svm, datasets
+
+def check_vetted(p: Pickled):
+    if p.has_unvetted_dependency:
+        print(f"unvetted deps : {p.unvetted_dependencies}")
+
+# sklearn
+clf = svm.SVC()
+X, y = datasets.load_iris(return_X_y=True)
+clf.fit(X, y)
+s = pickle.dumps(clf)
+
+p = Pickled.load(s)
+p.vetted_dependencies = ["numpy.ndarray", "sklearn.svm._classes.SVC", "numpy.core.multiarray._reconstruct", "numpy.dtype", "numpy.core.multiarray.scalar"]
+check_vetted(p)
+if p.is_likely_safe:
+    print("✅")
+else:
+    print("❌")
+
+# numpy
+arr = np.ndarray([1, 2, 3])
+p = Pickled.load(pickle.dumps(arr))
+p.vetted_dependencies = ["numpy.ndarray"]
+check_vetted(p)
+if p.is_likely_safe:
+    print("✅")
+else:
+    print("❌")
+

fickling/pickle.py

@@ -246,6 +246,8 @@ def __init__(self, opcodes: Iterable[Opcode]):
         self._opcodes: List[Opcode] = list(opcodes)
         self._ast: Optional[ast.Module] = None
         self._properties: Optional[ASTProperties] = None
+        self._vetted_dependencies: Optional[List[str]] = None
+        self._unvetted_dependencies: Optional[List[str]] = None
 
     def __len__(self) -> int:
         return len(self._opcodes)
@@ -404,6 +406,24 @@ def has_import(self) -> bool:
         """Checks whether unpickling would cause an import to be run"""
         return bool(self.properties.imports)
 
+    @property
+    def has_unvetted_dependency(self) -> bool:
+        if self._vetted_dependencies is None:
+            raise ValueError("Cannot call has_unvetted_import when vetted_dependencies is not set")
+        if self._unvetted_dependencies is None:
+            self._unvetted_dependencies = []
+            for st in self.ast.body:
+                if isinstance(st, ast.ImportFrom):
+                    for imp in st.names:
+                        import_path = f"{st.module}.{imp.name}"
+                        if import_path not in self._vetted_dependencies:
+                            self._unvetted_dependencies.append(import_path)
+            return len(self._unvetted_dependencies) > 0
+        elif len(self._unvetted_dependencies) > 0:
+            return True
+        else:
+            return False
+
     @property
     def has_call(self) -> bool:
         """Checks whether unpickling would cause a function call"""
@@ -437,6 +457,18 @@ def ast(self) -> ast.Module:
             self._ast = Interpreter.interpret(self)
         return self._ast
 
+    @property
+    def vetted_dependencies(self) -> Optional[List[str]]:
+        return self._vetted_dependencies
+
+    @vetted_dependencies.setter
+    def vetted_dependencies(self, dependencies: List[str]):
+        self._vetted_dependencies = dependencies
+
+    @property
+    def unvetted_dependencies(self) -> Optional[List[str]]:
+        return self._unvetted_dependencies
+
 
 class Stack(GenericSequence, Generic[T]):
     def __init__(self, initial_value: Iterable[T] = ()):