Supporting secure unpickling in PyRosetta (#523)

Currently, `PackedPose` objects are serialized/deserialized using the
`pickle` module (introduced in ~2019), and the `Pose.cache` dictionary
(introduced in #430) supports caching arbitrary datatypes in the `Pose`
object using the `pickle` module. Additionally, #462 enables saving
compressed `PackedPose` objects to disk (i.e., as `*.b64_pose` and
`*.pkl_pose` files) for sharing PyRosetta `Pose` objects with the
scientific community. However, use of the `pickle` module is not secure
(see warning [here](https://docs.python.org/3/library/pickle.html) as
outlined in #519).

Herein this PR, a secure `pickle.loads` method is developed and slotted
into the `PackedPose` and `Pose.cache` infrastructure to permanently
disallow certain risky packages, modules, and namespaces from being
unpickled/loaded (e.g., `exec`, `eval`, `os.system`, `subprocess.run`,
etc., and will be updated over time as needed), thus significantly
improving the security of handling `PackedPose` and `Pose` objects in
memory if received from a second party (i.e., over a socket, queue,
interprocess communication, etc.) or when reading a file received from a
second party (i.e., using `pyrosetta.distributed.io.pose_from_file` with
a `*.b64_pose` and `*.pkl_pose` file). By default, only `pyrosetta` and
`numpy` packages, and certain `builtins` modules (like `dict`,
`complex`, `tuple`, etc.), are considered secure and permitted to be
unpickled/loaded. Other packages that the user may want to
serialize/deserialize may be assigned as secure per-process by the user
in-code (see methods below). It is worth noting that PyTorch developers
have implemented a similar strategy with the
[torch.serialization.add_safe_globals()](https://docs.pytorch.org/docs/stable/notes/serialization.html#torch.serialization.add_safe_globals)
method.

Another aim of this PR is to implement an optional Hash-based Message
Authentication Code (HMAC) key in the `Pose.cache` dictionary for data
integrity verification. While not a security feature, this new API
allows the user to set a HMAC key to be prepended to every score value
in the `Pose.cache` dictionary that effectively says "this was saved by
PyRosetta", so that it intentionally raises an error when the HMAC key
is missing or differs upon retrieval, indicating that the data appears
to have been tampered with or modified. By default, the HMAC key is
disabled (being set to `None`) in order to reduce memory overhead of the
`Pose.cache` dictionary; e.g., if 32 bytes are prepended to each score
value, with 1,000 score values that's 32,000 bytes or 32 KB of overhead,
and with a million score values that's 32 MB of overhead.

The following are newly added functions:
- `pyrosetta.secure_unpickle.add_secure_package`: Add a package to the
unpickle allowed list
- `pyrosetta.secure_unpickle.remove_secure_package`: Remove a package
from the unpickle allowed list
- `pyrosetta.secure_unpickle.clear_secure_packages`: Remove all packages
from the unpickle allowed list
- `pyrosetta.secure_unpickle.get_disallowed_packages`: Return all
permanently disallowed packages/modules/prefixes
- `pyrosetta.secure_unpickle.get_secure_packages`: Return all packages
in the unpickle allowed list
- `pyrosetta.secure_unpickle.set_secure_packages`: Set all packages in
the unpickle allowed list
- `pyrosetta.secure_unpickle.set_unpickle_hmac_key`: Set the HMAC key
for the `Pose.cache` dictionary
- `pyrosetta.secure_unpickle.get_unpickle_hmac_key`: Return the HMAC key
for the `Pose.cache` dictionary

---------

Co-authored-by: Rachel Clune <rachel.clune@omsf.io>

Author: klimaj

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):