Repository: SFDO-Tooling/CumulusCI
Pull request #3390 (2022-10-11 22:37:27)

Use JSON instead of Pickle
- All new keychain details are serialized using json, not pickle.

prescod: @bethbrains 

We're going to have to support Dates and DateTimes.

https://github.com/SFDO-Tooling/CumulusCI/actions/runs/3230706675/jobs/5289461956#step:8:41

We need to use the `default=` parameter:

https://docs.python.org/3/library/json.html#json.dump

We should also do some kind of survey to figure out what other types we might need to support.

Patch:
diff --git a/cumulusci/core/config/base_config.py b/cumulusci/core/config/base_config.py
index 915455b23..d2bc79deb 100644
--- a/cumulusci/core/config/base_config.py
+++ b/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 @@ STRICT_GETATTR = os.environ.get("STRICT_GETATTR")
 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 @@ class BaseConfig(object):
         """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__", {})
diff --git a/cumulusci/core/keychain/encrypted_file_project_keychain.py b/cumulusci/core/keychain/encrypted_file_project_keychain.py
index 854df2ac5..66c365f6e 100644
--- a/cumulusci/core/keychain/encrypted_file_project_keychain.py
+++ b/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,6 +26,7 @@ from cumulusci.core.exceptions import (
 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"
@@ -33,7 +34,7 @@ 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,10 +44,6 @@ BS = 16
 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
@@ -54,6 +51,22 @@ 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 @@ class EncryptedFileProjectKeychain(BaseProjectKeychain):
         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 @@ class EncryptedFileProjectKeychain(BaseProjectKeychain):
             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 @@ class EncryptedFileProjectKeychain(BaseProjectKeychain):
         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 @@ class EncryptedFileProjectKeychain(BaseProjectKeychain):
                 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 @@ class EncryptedFileProjectKeychain(BaseProjectKeychain):
         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 @@ class EncryptedFileProjectKeychain(BaseProjectKeychain):
                 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
diff --git a/cumulusci/utils/pickle.py b/cumulusci/utils/pickle.py
new file mode 100644
index 000000000..7bb0ebd9e
--- /dev/null
+++ b/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()


Test patch:
diff --git a/cumulusci/core/keychain/tests/test_encrypted_file_project_keychain.py b/cumulusci/core/keychain/tests/test_encrypted_file_project_keychain.py
index 73a7c2809..0b4f33e1f 100644
--- a/cumulusci/core/keychain/tests/test_encrypted_file_project_keychain.py
+++ b/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 @@ class TestEncryptedFileProjectKeychain:
     #               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 @@ class TestEncryptedFileProjectKeychain:
 
         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