Update readme, add exception

README.md

@@ -1,23 +1,19 @@
-<p align="center">
-<img src="https://github.com/trailofbits/fickling/blob/sh/readme/fickling_image.png" width="600" height="312">
-</p>
-
 # Fickling
 
 Fickling is a decompiler, static analyzer, and bytecode rewriter for Python
 [pickle](https://docs.python.org/3/library/pickle.html) object serializations.
-You can use fickling to detect, analyze, reverse engineer, or even create 
-malicious pickle or pickle-based files, including PyTorch files. 
+You can use fickling to detect, analyze, reverse engineer, or even create
+malicious pickle or pickle-based files, including PyTorch files.
 
 Fickling can be used both as a **python library** and a **CLI**.
 
 - [Installation](#installation)
 - [Malicious file detection](#malicious-file-detection)
 - [Advanced usage](#advanced-usage)
-    - [Trace pickle execution](#trace-pickle-execution)
-    - [Pickle code injection](#pickle-code-injection)
-    - [Pickle decompilation](#pickle-decompilation)
-    - [PyTorch polyglots](#pytorch-polyglots)
+  - [Trace pickle execution](#trace-pickle-execution)
+  - [Pickle code injection](#pickle-code-injection)
+  - [Pickle decompilation](#pickle-decompilation)
+  - [PyTorch polyglots](#pytorch-polyglots)
 - [About pickle](#about-pickle)
 - [Contact](#contact)
 
@@ -34,9 +30,10 @@ python -m pip install fickling
 
 Fickling can seamlessly be integrated into your codebase to detect and halt the loading of malicious files at runtime.
 
-Below we show the different ways you can use fickling to enforce safety checks on pickle files. Under the hood, it hooks the `pickle` library to add safety checks so that loading a pickle file raises an `UnsafePickleFile` exception if malicious content is detected in the file.
+Below we show the different ways you can use fickling to enforce safety checks on pickle files. Under the hood, it hooks the `pickle` library to add safety checks so that loading a pickle file raises an `UnsafeFileError` exception if malicious content is detected in the file.
 
 #### Option 1 (recommended): check safety of all pickle files loaded
+
 ```python
 # This enforces safety checks every time pickle.load() is used
 fickling.always_check_safety()
@@ -45,48 +42,51 @@ fickling.always_check_safety()
 with open("file.pkl", "rb") as f:
     try:
         pickle.load(f)
-    except fickling.UnsafePickleFile:
+    except fickling.UnsafeFileError:
         print("Unsafe file!")
 ```
 
 #### Option 2: use a context manager
+
 ```python
 with fickling.check_safety():
     # All pickle files loaded within the context manager are checked for safety
     try:
         with open("file.pkl", "rb") as f:
             pickle.load("file.pkl")
-    except fickling.UnsafePickleFile:
+    except fickling.UnsafeFileError:
         print("Unsafe file!")
 
 # Files loaded outside of context manager are NOT checked
 pickle.load("file.pkl")
 ```
 
-
 #### Option 3: check and load a single file
+
 ```python
 # Use fickling.load() in place of pickle.load() to check safety and load a single pickle file 
 try:
     fickling.load("file.pkl")
-except fickling.UnsafePickleFile as e:
+except fickling.UnsafeFileError as e:
     print("Unsafe file!")
 ```
 
 #### Option 4: only check pickle file safety without loading
+
 ```python3
 # Perform a safety check on a pickle file without loading it
 if not fickling.is_likely_safe("file.pkl"):
     print("Unsafe file!")
 ```
 
 #### Accessing the safety analysis results
+
 You can access the details of fickling's safety analysis from within the raised exception:
 
 ```python
 >>> try:
 ...     fickling.load("unsafe.pkl")
-... except fickling.UnsafePickleFile as e:
+... except fickling.UnsafeFileError as e:
 ...     print(e.info)
 
 {
@@ -107,28 +107,31 @@ You can access the details of fickling's safety analysis from within the raised
 If you are using another language than Python, you can still use fickling's `CLI` to safety-check pickle files:
 
 ```console
-$ fickling --check-safety -p pickled.data
+fickling --check-safety -p pickled.data
 ```
 
 ## Advanced usage
 
 ### Trace pickle execution
+
 Fickling's `CLI` allows to safely trace the execution of the Pickle virtual machine without
 exercising any malicious code:
 
 ```console
-$ fickling --trace file.pkl
+fickling --trace file.pkl
 ```
 
 ### Pickle code injection
+
 Fickling allows to inject arbitrary code in a pickle file that will run every time the file is loaded
 
 ```console
-$ fickling --inject "print('Malicious')" file.pkl
+fickling --inject "print('Malicious')" file.pkl
 ```
 
 ### Pickle decompilation
-Fickling can be used to decompile a pickle file for further analysis 
+
+Fickling can be used to decompile a pickle file for further analysis
 
 ```python
 >>> import ast, pickle
@@ -150,8 +153,10 @@ Module(
     type_ignores=[])
 ```
 
-### PyTorch polyglots 
+### PyTorch polyglots
+
 We currently support inspecting, identifying, and creating file polyglots between the following PyTorch file formats:
+
 - **PyTorch v0.1.1**: Tar file with sys_info, pickle, storages, and tensors
 - **PyTorch v0.1.10**: Stacked pickle files
 - **TorchScript v1.0**: ZIP file with model.json and constants.pkl (a JSON file and a pickle file)
@@ -177,13 +182,14 @@ Your file is most likely of this format:  PyTorch v1.3
 
 Check out [our examples](https://github.com/trailofbits/fickling/tree/master/example) to learn more about using fickling!
 
-## About pickle 
+## About pickle
+
 Pickled Python objects are in fact bytecode that is interpreted by a stack-based
 virtual machine built into Python called the "Pickle Machine". Fickling can take
 pickled data streams and decompile them into human-readable Python code that,
-when executed, will deserialize to the original serialized object. This is made 
-possible by Fickling’s custom implementation of the PM. Fickling is safe to run 
-on potentially malicious files because its PM symbolically executes code rather 
+when executed, will deserialize to the original serialized object. This is made
+possible by Fickling’s custom implementation of the PM. Fickling is safe to run
+on potentially malicious files because its PM symbolically executes code rather
 than overtly executing it.
 
 The authors do not prescribe any meaning to the “F” in Fickling; it could stand
@@ -194,7 +200,8 @@ Learn more about fickling in our [blog post](https://blog.trailofbits.com/2021/0
 and [DEF CON AI Village 2021 talk](https://www.youtube.com/watch?v=bZ0m_H_dEJI).
 
 ## Contact
-If you'd like to file a bug report or feature request, please use our [issues](https://github.com/trailofbits/fickling/issues) page. 
+
+If you'd like to file a bug report or feature request, please use our [issues](https://github.com/trailofbits/fickling/issues) page.
 Feel free to contact us or reach out in [Empire Hacking](https://slack.empirehacking.nyc/) for help using or extending fickling.
 
 ## License
@@ -205,7 +212,3 @@ It is licensed under the [GNU Lesser General Public License v3.0](LICENSE).
 exception to the terms.
 
 © 2021, Trail of Bits.
-
-<p align="center">
-<strong><i>We relish the thought of a day when pickling will no longer be used to deserialize untrusted files.</i></strong>
-</p>

example/hook_functions.py

@@ -16,6 +16,7 @@
 with open("safe.pkl", "wb") as file:
     pickle.dump(test_list, file)
 
+
 class Payload:
     def __init__(self):
         self.a = 1
@@ -41,5 +42,5 @@ def __reduce__(self):
 
 # This hook works when pickle.load is called under the hood in Python as well
 # Note that this does not always work for torch.load()
-# This should raise "SafetyError"
+# This should raise "UnsafeFileError"
 numpy.load("unsafe.pkl", allow_pickle=True)

fickling/analysis.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from enum import Enum
-from typing import BinaryIO, ByteString, Dict, Iterable, Iterator, Optional, Set, Tuple, Type, Union
+from typing import Dict, Iterable, Iterator, Optional, Set, Tuple, Type
 
 if sys.version_info < (3, 9):
     from astunparse import unparse

fickling/exception.py

@@ -0,0 +1,8 @@
+class UnsafeFileError(Exception):
+    def __init__(self, filepath, info):
+        super().__init__()
+        self.filepath = filepath
+        self.info = info
+
+    def __str__(self):
+        return f"Safety results for {self.filepath} : {str(self.info)}"

fickling/loader.py

@@ -1,15 +1,10 @@
 import pickle
 
 from fickling.analysis import Severity, check_safety
+from fickling.exception import UnsafeFileError
 from fickling.fickle import Pickled
 
 
-class SafetyError(Exception):
-    """Exception raised when a file is deemed unsafe by fickling."""
-
-    pass
-
-
 def load(
     file,
     max_acceptable_severity=Severity.LIKELY_SAFE,
@@ -27,4 +22,4 @@ def load(
         # loaded after the analysis.
         return pickle.loads(pickled_data.dumps(), *args, **kwargs)
     else:
-        raise SafetyError(f"File is unsafe: {result.severity.name}")
+        raise UnsafeFileError(file, result.to_dict())

test/test_hook.py

@@ -6,7 +6,7 @@
 import numpy
 
 import fickling.hook as hook
-from fickling.loader import SafetyError
+from fickling.exception import UnsafeFileError
 
 
 class TestHook(unittest.TestCase):
@@ -47,7 +47,7 @@ def __reduce__(self):
         try:
             numpy.load("unsafe.pickle", allow_pickle=True)
         except UnpicklingError as e:
-            if isinstance(e.__cause__, SafetyError):
+            if isinstance(e.__cause__, UnsafeFileError):
                 pass
             else:
                 self.fail(e)