Supporting secure unpickling in PyRosetta

source/src/python/PyRosetta/src/pyrosetta/__init__.py

@@ -67,6 +67,16 @@
     get_score_function,
     Pose,
 )
+from pyrosetta.secure_unpickle import (
+    add_secure_package,
+    remove_secure_package,
+    clear_secure_packages,
+    get_disallowed_packages,
+    get_secure_packages,
+    set_secure_packages,
+    UnpickleSecurityError,
+    UnpickleIntegrityError,
+)
 
 ###############################################################################
 # Exception handling.

source/src/python/PyRosetta/src/pyrosetta/bindings/scores/serialization.py

@@ -9,48 +9,36 @@
 __author__ = "Jason C. Klima"
 
 
-import base64
 import collections
-import pickle
+
+from pyrosetta.secure_unpickle import SecureSerializerBase
 
 
 class PoseScoreSerializerBase(object):
     """Base class for `PoseScoreSerializer` methods."""
-
     @staticmethod
     def to_pickle(value):
-        try:
-            return pickle.dumps(value)
-        except (TypeError, OverflowError, MemoryError, pickle.PicklingError) as ex:
-            raise TypeError(
-                "Only pickle-serializable object types are allowed to be set "
-                + "as score values. Received: %r. %s" % (type(value), ex)
-            )
+        return SecureSerializerBase.to_pickle(value)
 
     @staticmethod
     def from_pickle(value):
-        try:
-            return pickle.loads(value)
-        except (TypeError, OverflowError, MemoryError, EOFError, pickle.UnpicklingError) as ex:
-            raise TypeError(
-                "Could not deserialize score value of type %r. %s" % (type(value), ex)
-            )
+        return SecureSerializerBase.secure_loads(value)
 
     @staticmethod
     def to_base64(value):
-        return base64.b64encode(value).decode()
+        return SecureSerializerBase.to_base64(value)
 
     @staticmethod
     def from_base64(value):
-        return base64.b64decode(value, validate=True)
+        return SecureSerializerBase.from_base64(value)
 
     @staticmethod
-    def to_base64_pickle(value):
-        return PoseScoreSerializerBase.to_base64(PoseScoreSerializerBase.to_pickle(value))
+    def to_base64_pickle(obj):
+        return SecureSerializerBase.secure_to_base64_pickle(obj)
 
     @staticmethod
     def from_base64_pickle(value):
-        return PoseScoreSerializerBase.from_pickle(PoseScoreSerializerBase.from_base64(value))
+        return SecureSerializerBase.secure_from_base64_pickle(value)
 
     @staticmethod
     def bool_from_str(value):

source/src/python/PyRosetta/src/pyrosetta/distributed/cluster/__init__.py

@@ -49,7 +49,7 @@
     "run",
     "update_scores",
 ]
-__version__: str = "2.1.1"
+__version__: str = "2.1.2"
 
 _print_conda_warnings()
 

source/src/python/PyRosetta/src/pyrosetta/distributed/cluster/serialization.py

@@ -38,6 +38,7 @@
 
 from functools import singledispatch, wraps
 from pyrosetta.distributed.packed_pose.core import PackedPose
+from pyrosetta.secure_unpickle import SecureSerializerBase
 from typing import (
     Any,
     Dict,
@@ -189,7 +190,7 @@ def compress_packed_pose(self, packed_pose: Any) -> Union[NoReturn, None, bytes]
             compressed_packed_pose = None
         elif isinstance(packed_pose, PackedPose):
             packed_pose = update_scores(packed_pose)
-            compressed_packed_pose = self.encoder(packed_pose.pickled_pose)
+            compressed_packed_pose = self.encoder(io.to_pickle(packed_pose))
         else:
             raise TypeError(
                 "The 'packed_pose' argument parameter must be of type `NoneType` or `PackedPose`."
@@ -200,8 +201,8 @@ def compress_packed_pose(self, packed_pose: Any) -> Union[NoReturn, None, bytes]
     @requires_compression
     def decompress_packed_pose(self, compressed_packed_pose: Any) -> Union[NoReturn, None, PackedPose]:
         """
-        Decompress a `bytes` object with the custom serialization and `cloudpickle` modules. If the 'compressed_packed_pose'
-        argument parameter is `None`, then just return `None`.
+        Decompress a `bytes` object with the custom serialization module and secure implementation of the `pickle` module.
+        If the 'compressed_packed_pose' argument parameter is `None`, then just return `None`.
 
         Args:
             compressed_packed_pose: the input `bytes` object to decompress. If `None`, then just return `None`.
@@ -215,7 +216,7 @@ def decompress_packed_pose(self, compressed_packed_pose: Any) -> Union[NoReturn,
         if compressed_packed_pose is None:
             packed_pose = None
         elif isinstance(compressed_packed_pose, bytes):
-            pose = cloudpickle.loads(self.decoder(compressed_packed_pose))
+            pose = SecureSerializerBase.secure_loads(self.decoder(compressed_packed_pose))
             packed_pose = io.to_packed(pose)
         else:
             raise TypeError(

source/src/python/PyRosetta/src/pyrosetta/distributed/packed_pose/core.py

@@ -13,6 +13,9 @@
 import pyrosetta.rosetta.core.pose as pose
 import pyrosetta.distributed
 
+from pyrosetta.secure_unpickle import SecureSerializerBase
+
+
 __all__ = ["pack_result", "pose_result", "to_packed", "to_pose", "to_dict", "to_base64", "to_pickle", "PackedPose"]
 
 
@@ -38,7 +41,7 @@ class PackedPose:
     def __init__(self, pose_or_pack):
         """Create a packed pose from pose, pack, or pickled bytes."""
         if isinstance(pose_or_pack, pose.Pose):
-            self.pickled_pose = pickle.dumps(pose_or_pack)
+            self.pickled_pose = SecureSerializerBase.to_pickle(pose_or_pack)
             self.scores = dict(pose_or_pack.cache)
 
         elif isinstance(pose_or_pack, PackedPose):
@@ -55,7 +58,7 @@ def __init__(self, pose_or_pack):
     @property
     @pyrosetta.distributed.requires_init
     def pose(self):
-        return pickle.loads(self.pickled_pose)
+        return SecureSerializerBase.secure_loads(self.pickled_pose)
 
     def update_scores(self, *score_dicts, **score_kwargs):
         new_scores = {}
@@ -71,7 +74,9 @@ def update_scores(self, *score_dicts, **score_kwargs):
 
     def clone(self):
         result = PackedPose(self.pose)
-        result.scores = pickle.loads(pickle.dumps(self.scores))
+        result.scores = SecureSerializerBase.secure_loads(
+            SecureSerializerBase.to_pickle(self.scores)
+        )
         return result
 
     def empty(self):