Use JSON instead of Pickle (#3390)

cumulusci/core/config/base_config.py

@@ -1,7 +1,9 @@
+import json
 import logging
 import os
 import typing as T
 import warnings
+from datetime import date, datetime
 from functools import lru_cache
 
 # turn on strictness.
@@ -10,6 +12,13 @@
 CHECK_CONFIG_TYPES = os.environ.get("CHECK_CONFIG_TYPES")
 
 
+def load_dates(x):
+    if isinstance(x, datetime):
+        return {"$type": "datetime", "$value": x.isoformat()}
+    elif isinstance(x, date):
+        return {"$type": "date", "$value": x.isoformat()}
+
+
 class BaseConfig(object):
     """BaseConfig provides a common interface for nested access for all Config objects in CCI."""
 
@@ -43,6 +52,9 @@ def _load_config(self):
         """Subclasses may override this method to initialize :py:attr:`~config`"""
         pass
 
+    def _serialize(self):
+        return json.dumps(self.config, default=load_dates).encode("utf-8")
+
     @classmethod
     def _allowed_names(cls) -> T.Dict[str, type]:
         return getattr(cls, "__annotations__", {})

cumulusci/core/keychain/encrypted_file_project_keychain.py

@@ -1,9 +1,9 @@
 import base64
 import json
 import os
-import pickle
 import sys
 import typing as T
+from datetime import date, datetime
 from pathlib import Path
 from shutil import rmtree
 
@@ -26,14 +26,15 @@
 from cumulusci.core.keychain import BaseProjectKeychain
 from cumulusci.core.keychain.base_project_keychain import DEFAULT_CONNECTED_APP_NAME
 from cumulusci.core.utils import import_class, import_global
+from cumulusci.utils.pickle import restricted_loads
 from cumulusci.utils.yaml.cumulusci_yml import ScratchOrg
 
 DEFAULT_SERVICES_FILENAME = "DEFAULT_SERVICES.json"
 
 # The file permissions that we want set on all
 # .org and .service files. Equivalent to -rw-------
 SERVICE_ORG_FILE_MODE = 0o600
-OS_FILE_FLAGS = os.O_WRONLY | os.O_CREAT
+OS_FILE_FLAGS = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
 if sys.platform.startswith("win"):
     # O_BINARY only available on Windows
     OS_FILE_FLAGS |= os.O_BINARY
@@ -43,17 +44,29 @@
 backend = default_backend()
 
 
-def pad(s):
-    return s + (BS - len(s) % BS) * chr(BS - len(s) % BS).encode("ascii")
-
-
 scratch_org_class = os.environ.get("CUMULUSCI_SCRATCH_ORG_CLASS")
 if scratch_org_class:
     scratch_org_factory = import_global(scratch_org_class)  # pragma: no cover
 else:
     scratch_org_factory = ScratchOrgConfig
 
 
+def extract_dates(x):
+    if x.get("$type") == "date":
+        return date.fromisoformat(x["$value"])
+    if x.get("$type") == "datetime":
+        return datetime.fromisoformat(x["$value"])
+    assert "$type" not in x, f"Unknown $type: {x['$type']}"
+    return x
+
+
+def safe_load_json_or_pickle(data):
+    try:
+        return json.loads(data, object_hook=extract_dates)
+    except ValueError:
+        return restricted_loads(data)
+
+
 """
 EncryptedFileProjectKeychain
 ----------------------------
@@ -112,10 +125,10 @@ def _get_cipher(self, iv=None):
         return cipher, iv
 
     def _encrypt_config(self, config):
-        pickled = pickle.dumps(config.config, protocol=2)
-        pickled = pad(pickled)
+        json_obj = config._serialize()
+        encryptor_value = json_obj + (BS - len(json_obj) % BS) * b" "
         cipher, iv = self._get_cipher()
-        return base64.b64encode(iv + cipher.encryptor().update(pickled))
+        return base64.b64encode(iv + cipher.encryptor().update(encryptor_value))
 
     def _decrypt_config(self, config_class, encrypted_config, extra=None, context=None):
         if self.key:
@@ -127,17 +140,17 @@ def _decrypt_config(self, config_class, encrypted_config, extra=None, context=No
             encrypted_config = base64.b64decode(encrypted_config)
             iv = encrypted_config[:16]
             cipher, iv = self._get_cipher(iv)
-            pickled = cipher.decryptor().update(encrypted_config[16:])
+            decrypted = cipher.decryptor().update(encrypted_config[16:])
             try:
-                unpickled = pickle.loads(pickled, encoding="bytes")
-            except Exception:
+                config_obj = safe_load_json_or_pickle(decrypted)
+            except Exception as exc:
                 raise KeychainKeyNotFound(
                     f"Unable to decrypt{' ' + context if context else ''}. "
                     "It was probably stored using a different CUMULUSCI_KEY."
-                )
+                ) from exc
             # Convert bytes created in Python 2
             config_dict = {}
-            for k, v in unpickled.items():
+            for k, v in config_obj.items():
                 if isinstance(k, bytes):
                     k = k.decode("utf-8")
                 if isinstance(v, bytes):
@@ -165,7 +178,7 @@ def _get_config_bytes(self, config) -> bytes:
         if self.key:
             org_bytes = self._encrypt_config(config)
         else:
-            org_bytes = pickle.dumps(config.config)
+            org_bytes = config._serialize()
 
         assert org_bytes is not None, "org_bytes should have a value"
         return org_bytes
@@ -367,7 +380,7 @@ def _config_from_bytes(self, config, name):
                 context=f"org config ({name})",
             )
         else:
-            config = pickle.loads(config)
+            config = safe_load_json_or_pickle(config)
             org = self._construct_config(OrgConfig, [config, name, self])
 
         return org
@@ -592,7 +605,7 @@ def _set_service(
         elif self.key and not config_encrypted:
             config = self._encrypt_config(service_config)
         else:
-            config = pickle.dumps(service_config.config)
+            config = service_config._serialize()
 
         self.services[service_type][alias] = config
 
@@ -643,7 +656,7 @@ def _get_service(self, service_type, alias):
                 context=f"service config ({service_type}:{alias})",
             )
         else:
-            config = pickle.loads(config)
+            config = safe_load_json_or_pickle(config)
             org = self._construct_config(ConfigClass, [config, alias, self])
 
         return org

cumulusci/core/keychain/tests/test_encrypted_file_project_keychain.py

@@ -1,7 +1,6 @@
 import datetime
 import json
 import os
-import pickle
 import re
 import sys
 import tempfile
@@ -74,11 +73,20 @@ def test_raise_service_not_configured(self, keychain):
     #               Orgs                  #
     #######################################
 
-    def test_set_and_get_org__global(self, keychain, org_config):
+    def test_set_and_get_org_with_dates__global(self, keychain, org_config):
         org_config.global_org = True
+
+        custom_datetime = datetime.datetime.now()
+        custom_date = datetime.datetime.now().date()
+        org_config.config["custom_datetime"] = custom_datetime
+        org_config.config["custom_date"] = custom_date
+
         keychain.set_org(org_config, True)
         assert list(keychain.orgs.keys()) == ["test"]
-        assert keychain.get_org("test").config == org_config.config
+        config = keychain.get_org("test").config
+        assert config == org_config.config
+        assert config["custom_datetime"] == custom_datetime
+        assert config["custom_date"] == custom_date
 
     def test_get_org__with_config_properly_overridden(
         self, keychain, scratch_org_config
@@ -112,7 +120,7 @@ def test_set_org__no_key_should_save_to_unencrypted_file(
 
         filepath = Path(keychain.project_local_dir, "test.org")
         with open(filepath, "rb") as f:
-            assert pickle.load(f) == org_config.config
+            assert json.load(f) == org_config.config
 
     def test_set_org__should_not_save_when_environment_project_keychain_set(
         self, keychain, org_config

cumulusci/utils/pickle.py

@@ -0,0 +1,44 @@
+"""Tools for safely saving and loading pickles using an AllowedList"""
+
+import io
+import pickle
+import typing as T
+import warnings
+
+
+class Type_Cannot_Be_Used_With_Random_Reference(T.NamedTuple):
+    """This type cannot be unpickled."""
+
+    module: str
+    name: str
+
+
+_SAFE_CLASSES = {
+    ("decimal", "Decimal"),
+    ("datetime", "date"),
+    ("datetime", "datetime"),
+    ("datetime", "timedelta"),
+    ("datetime", "timezone"),
+}
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """Safe unpickler with an allowed-list"""
+
+    count = 0
+
+    def find_class(self, module, name):
+        # Only allow safe classes from builtins.
+        if (module, name) in _SAFE_CLASSES:
+            return super().find_class(module, name)
+        else:
+            # Return a "safe" object that does nothing.
+            if RestrictedUnpickler.count < 10:
+                warnings.warn(f"Cannot save and refer to {module}, {name}")
+                RestrictedUnpickler.count += 1
+            return lambda *args: Type_Cannot_Be_Used_With_Random_Reference(module, name)
+
+
+def restricted_loads(data):
+    """Helper function analogous to pickle.loads()."""
+    return RestrictedUnpickler(io.BytesIO(data)).load()