From 0de6aa36cd75235836918074379ce06e4673c915 Mon Sep 17 00:00:00 2001
From: Adam Souzis <adam@souzis.com>
Date: Wed, 13 Nov 2024 11:19:30 -0800
Subject: [PATCH] dsl: loader: preload allowed modules in safe mode

---
 smoketest.sh                       | 2 +-
 tests/test_constraints.py          | 2 --
 tests/test_docs.py                 | 3 ++-
 tests/test_dsl.py                  | 6 ++++--
 tosca-package/tosca/loader.py      | 7 +++++++
 tosca-package/tosca/python2yaml.py | 4 ++--
 unfurl/dsl.py                      | 2 +-
 7 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/smoketest.sh b/smoketest.sh
index 9833512c..69e7ff1a 100755
--- a/smoketest.sh
+++ b/smoketest.sh
@@ -1,3 +1,3 @@
 set -e
 .tox/${1:-py310}/bin/mypy unfurl --install-types --non-interactive
-UNFURL_TEST_SKIP=docker+slow+k8s+helm+$UNFURL_TEST_SKIP tox --skip-pkg-install -e ${1:-py310} -- -v --no-cov -n auto --dist loadfile $2 $3 $4 $5 $6 $7
+UNFURL_TEST_SKIP_BUILD_RUST=1 UNFURL_TEST_SKIP=docker+slow+k8s+helm+$UNFURL_TEST_SKIP tox --skip-pkg-install -e ${1:-py310} -- -v --no-cov -n auto --dist loadfile $2 $3 $4 $5 $6 $7
diff --git a/tests/test_constraints.py b/tests/test_constraints.py
index cd5d20c0..3739c439 100644
--- a/tests/test_constraints.py
+++ b/tests/test_constraints.py
@@ -328,8 +328,6 @@ def test_computed_properties():
             "skipped": 0,
             "changed": 1,
         }
-        # XXX we need to delete this module because mytypes gets re-evaluated, breaking class identity
-        # is this a scenario we need to worry about outside unit tests?
         result, job, summary = run_job_cmd(
             cli_runner, ["-vvv", "undeploy"], print_result=True
         )
diff --git a/tests/test_docs.py b/tests/test_docs.py
index e4a2fab4..eb3484cd 100644
--- a/tests/test_docs.py
+++ b/tests/test_docs.py
@@ -107,4 +107,5 @@ def test_quickstart():
         with open("service_template.py", "a") as f:
             f.write(deployment_blueprint)
         run_cmd(runner, "plan production")
-        run_cmd(runner, "deploy --dryrun --approve development")
+        if "slow" not in os.getenv("UNFURL_TEST_SKIP", ""):
+            run_cmd(runner, "deploy --dryrun --approve development")
diff --git a/tests/test_dsl.py b/tests/test_dsl.py
index 0c00b38d..64b25d56 100644
--- a/tests/test_dsl.py
+++ b/tests/test_dsl.py
@@ -1212,14 +1212,16 @@ def test_sandbox(capsys):
         """foo = dict(); foo[1] = 2; bar = list(); bar.append(1); baz = tuple()""",
         """import math; math.floor(1.0)""",
         """from unfurl.configurators.templates.dns import unfurl_relationships_DNSRecords""",
-         # """from unfurl.tosca_plugins import k8s; k8s.kube_artifacts""",
+        """from unfurl import artifacts""",
+        """import unfurl; unfurl.artifacts""",
+        """from unfurl.tosca_plugins import k8s; k8s.kube_artifacts""",
         """import tosca
 node = tosca.nodes.Root()
 node._name = "test"
         """,
     ]
     for src in allowed:
-        print("allowed", src)
+        # print("allowed?", src)
         assert _to_yaml(src, True)
 
 
diff --git a/tosca-package/tosca/loader.py b/tosca-package/tosca/loader.py
index 51379f67..3b018ed6 100644
--- a/tosca-package/tosca/loader.py
+++ b/tosca-package/tosca/loader.py
@@ -439,6 +439,13 @@ def get_descriptions(body):
     "urllib.parse",
 )
 
+def get_allowed_modules() -> Dict[str, ModuleType]:
+  allowed = {}
+  for name in ALLOWED_MODULES:
+      if name in sys.modules:
+         allowed[name] = ImmutableModule(name, sys.modules[name])
+  return allowed
+
 # XXX have the unfurl package set these:
 ALLOWED_PRIVATE_PACKAGES = [
     "unfurl.tosca_plugins",
diff --git a/tosca-package/tosca/python2yaml.py b/tosca-package/tosca/python2yaml.py
index 079fe8fa..cfd3181c 100644
--- a/tosca-package/tosca/python2yaml.py
+++ b/tosca-package/tosca/python2yaml.py
@@ -32,7 +32,7 @@
     EvalData,
     Namespace,
 )
-from .loader import restricted_exec, get_module_path
+from .loader import restricted_exec, get_module_path, get_allowed_modules
 
 
 class PythonToYaml:
@@ -441,7 +441,7 @@ def python_src_to_yaml_obj(
     import_resolver=None,
 ) -> dict:
     if modules is None:
-        modules = {}
+        modules = get_allowed_modules()
     global_state.modules = modules
     if namespace is None:
         namespace = {}
diff --git a/unfurl/dsl.py b/unfurl/dsl.py
index bfbb82fe..ffd6b2bf 100644
--- a/unfurl/dsl.py
+++ b/unfurl/dsl.py
@@ -123,7 +123,7 @@ def convert_to_yaml(
     if safe_mode_override:
         safe_mode = safe_mode_override != "never"
     if import_resolver.manifest.modules is None:
-        import_resolver.manifest.modules = {}
+        import_resolver.manifest.modules = tosca.loader.get_allowed_modules()
     yaml_src = python_src_to_yaml_obj(
         contents,
         namespace,