Fix for arbitrary code execution while unpickling in ipycache.load_va…

…rs() method.

Resolution for issue rossant#47 in the original repo.

Any malicious command trying to process through the unpickle command would have to go through the restricted_loads() method which only allows io.StringsIO to parse.
Anything else, and it would raise a UnpicklingError.

ipycache.py

@@ -19,7 +19,6 @@
 from IPython.display import clear_output
 import hashlib
 
-
 #------------------------------------------------------------------------------
 # Six utility functions for Python 2/3 compatibility
 #------------------------------------------------------------------------------
@@ -115,6 +114,7 @@ def load_vars(path, vars):
     with open(path, 'rb') as f:
         # Load the variables from the cache.
         try:
+            restricted_loads(f.read())
             cache = pickle.load(f)
         except EOFError as e:
             cache={}
@@ -151,8 +151,26 @@ def save_vars(path, vars_d):
     """
     with open(path, 'wb') as f:
         dump(vars_d, f)
-
-
+
+
+# ------------------------------------------------------------------------------
+# RestrictedUnpickler - For mitigating arbitrary code execution while unpickling
+# This function provides restriction of using only the io module
+# ------------------------------------------------------------------------------
+class RestrictedUnpickler(pickle.Unpickler):
+
+    def find_class(self, module, name):
+        if module == '_io' and name == 'StringIO':
+            return getattr(sys.modules[module], name)
+        # Forbid everything else.
+        raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
+                                     (module, name))
+
+
+def restricted_loads(s):
+    """Helper function analogous to pickle.loads()."""
+    return RestrictedUnpickler(io.BytesIO(s)).load()
+
 #------------------------------------------------------------------------------
 # CapturedIO
 #------------------------------------------------------------------------------

test_ipycache.py

@@ -206,3 +206,18 @@ def ip_push(vars):
           ip_user_ns=user_ns, ip_run_cell=ip_run_cell, ip_push=ip_push)
 
     os.remove(path)
+
+
+def test_load_exploitPickle():
+    class vulnLoad():
+        def __init__(self):
+            self.a = 1
+
+        def __reduce__(self):
+            return (os.system, ('uname -a',))
+
+    payload = vulnLoad()
+    path = "malicious.pkl"
+    with open("malicious.pkl", "wb") as f:
+        pickle.dump(payload, f)
+    assert_raises(pickle.UnpicklingError, load_vars, path, ['a'])