From ac892be0e0a3d40f066079e5d6371836d7bfeb93 Mon Sep 17 00:00:00 2001
From: Ajmal Aboobacker <43377443+B3EF@users.noreply.github.com>
Date: Thu, 7 Jan 2021 18:20:02 +0530
Subject: [PATCH 1/2] Update heartbeat.py

---
 cup/services/heartbeat.py | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/cup/services/heartbeat.py b/cup/services/heartbeat.py
index 37437e9..8ef08a4 100644
--- a/cup/services/heartbeat.py
+++ b/cup/services/heartbeat.py
@@ -11,10 +11,35 @@
 import pickle
 import platform
 import threading
+import io
+import builtins
 
 from cup import log
 from cup import net
 from cup.util import conf
+safe_builtins = {
+    'range',
+    'complex',
+    'set',
+    'frozenset',
+    'slice',
+}
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+
+    def find_class(self, module, name):
+        """Only allow safe classes from builtins"""
+        if module == "builtins" and name in safe_builtins:
+            return getattr(builtins, 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()
+
 if platform.system() == 'Linux':
     from cup.res import linux
 
@@ -63,7 +88,7 @@ def deserilize(self, binary):
         deserilize it from binary
         """
         try:
-            self._dict_info = pickle.loads(binary)
+            self._dict_info = pickle.loads(pickle.loads(restricted_loads(binary))
             return True
         # pylint: disable=W0703
         except Exception as error:
@@ -374,7 +399,7 @@ def _test():
     localhost = LinuxHost(name='localhost', init_this_host=True)
     binary = localhost.serilize()
     print('binary:{0}'.format(binary))
-    print(pickle.loads(binary))
+    print(pickle.loads(restricted_loads(binary)))
 
 
 if __name__ == '__main__':

From 9c84492b35de3bbd067f335c19244d87485b463e Mon Sep 17 00:00:00 2001
From: Ajmal Aboobacker <43377443+B3EF@users.noreply.github.com>
Date: Thu, 7 Jan 2021 20:03:55 +0530
Subject: [PATCH 2/2] Update heartbeat.py

---
 cup/services/heartbeat.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cup/services/heartbeat.py b/cup/services/heartbeat.py
index 8ef08a4..0d4bbe2 100644
--- a/cup/services/heartbeat.py
+++ b/cup/services/heartbeat.py
@@ -399,7 +399,8 @@ def _test():
     localhost = LinuxHost(name='localhost', init_this_host=True)
     binary = localhost.serilize()
     print('binary:{0}'.format(binary))
-    print(pickle.loads(restricted_loads(binary)))
+    restricted_loads(binary)
+    print(pickle.loads(binary))
 
 
 if __name__ == '__main__':