diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/README.md b/README.md index 451a6d1..0d1e7fa 100644 --- a/README.md +++ b/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/). diff --git a/example/unvetted_dependencies.py b/example/unvetted_dependencies.py new file mode 100644 index 0000000..4882517 --- /dev/null +++ b/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("❌") + diff --git a/fickling/pickle.py b/fickling/pickle.py index 62e5d4b..fb5ea72 100644 --- a/fickling/pickle.py +++ b/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] = ()):