From acf03ece9629fbdef6eb65b3f4bd9c324e9193fe Mon Sep 17 00:00:00 2001
From: Nicole Rappe <nicole.rappe@bunny-lab.io>
Date: Fri, 3 Oct 2025 18:17:39 -0600
Subject: [PATCH] Ensure script variables are available under original names

---
 .../Roles/role_ScriptExec_CURRENTUSER.py      | 33 ++++++++++++++++--
 Data/Agent/Roles/role_ScriptExec_SYSTEM.py    | 34 +++++++++++++++++--
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/Data/Agent/Roles/role_ScriptExec_CURRENTUSER.py b/Data/Agent/Roles/role_ScriptExec_CURRENTUSER.py
index b4c3d25..f9049e6 100644
--- a/Data/Agent/Roles/role_ScriptExec_CURRENTUSER.py
+++ b/Data/Agent/Roles/role_ScriptExec_CURRENTUSER.py
@@ -15,6 +15,11 @@ ROLE_CONTEXTS = ['interactive']
 IS_WINDOWS = os.name == 'nt'
 
 
+def _canonical_env_key(name: str) -> str:
+    cleaned = re.sub(r"[^A-Za-z0-9_]", "_", (name or "").strip())
+    return cleaned.upper()
+
+
 def _sanitize_env_map(raw) -> Dict[str, str]:
     env: Dict[str, str] = {}
     if isinstance(raw, dict):
@@ -24,7 +29,7 @@ def _sanitize_env_map(raw) -> Dict[str, str]:
             name = str(key).strip()
             if not name:
                 continue
-            env_key = re.sub(r"[^A-Za-z0-9_]", "_", name).upper()
+            env_key = _canonical_env_key(name)
             if not env_key:
                 continue
             if isinstance(value, bool):
@@ -37,6 +42,29 @@ def _sanitize_env_map(raw) -> Dict[str, str]:
     return env
 
 
+def _apply_variable_aliases(env_map: Dict[str, str], variables: List[Dict[str, str]]) -> Dict[str, str]:
+    if not isinstance(env_map, dict) or not isinstance(variables, list):
+        return env_map
+    for var in variables:
+        if not isinstance(var, dict):
+            continue
+        name = str(var.get('name') or '').strip()
+        if not name:
+            continue
+        canonical = _canonical_env_key(name)
+        if not canonical or canonical not in env_map:
+            continue
+        value = env_map[canonical]
+        alias = re.sub(r"[^A-Za-z0-9_]", "_", name)
+        if alias and alias not in env_map:
+            env_map[alias] = value
+        if alias == name:
+            continue
+        if re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", name) and name not in env_map:
+            env_map[name] = value
+    return env_map
+
+
 def _ps_literal(value: str) -> str:
     return "'" + value.replace("'", "''") + "'"
 
@@ -201,7 +229,7 @@ class Role:
                     name = str(var.get('name') or '').strip()
                     if not name:
                         continue
-                    key = re.sub(r"[^A-Za-z0-9_]", "_", name).upper()
+                    key = _canonical_env_key(name)
                     if key in env_map:
                         continue
                     default_val = var.get('default')
@@ -211,6 +239,7 @@ class Role:
                         env_map[key] = ""
                     else:
                         env_map[key] = str(default_val)
+                env_map = _apply_variable_aliases(env_map, variables)
                 try:
                     timeout_seconds = max(0, int(payload.get('timeout_seconds') or 0))
                 except Exception:
diff --git a/Data/Agent/Roles/role_ScriptExec_SYSTEM.py b/Data/Agent/Roles/role_ScriptExec_SYSTEM.py
index c3c4794..0ba5852 100644
--- a/Data/Agent/Roles/role_ScriptExec_SYSTEM.py
+++ b/Data/Agent/Roles/role_ScriptExec_SYSTEM.py
@@ -16,6 +16,11 @@ def _project_root():
     return os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
 
 
+def _canonical_env_key(name: str) -> str:
+    cleaned = re.sub(r"[^A-Za-z0-9_]", "_", (name or "").strip())
+    return cleaned.upper()
+
+
 def _sanitize_env_map(raw) -> Dict[str, str]:
     env: Dict[str, str] = {}
     if isinstance(raw, dict):
@@ -25,7 +30,7 @@ def _sanitize_env_map(raw) -> Dict[str, str]:
             name = str(key).strip()
             if not name:
                 continue
-            env_key = re.sub(r"[^A-Za-z0-9_]", "_", name).upper()
+            env_key = _canonical_env_key(name)
             if not env_key:
                 continue
             if isinstance(value, bool):
@@ -38,6 +43,30 @@ def _sanitize_env_map(raw) -> Dict[str, str]:
     return env
 
 
+def _apply_variable_aliases(env_map: Dict[str, str], variables: List[Dict[str, str]]) -> Dict[str, str]:
+    if not isinstance(env_map, dict) or not isinstance(variables, list):
+        return env_map
+    for var in variables:
+        if not isinstance(var, dict):
+            continue
+        name = str(var.get('name') or '').strip()
+        if not name:
+            continue
+        canonical = _canonical_env_key(name)
+        if not canonical or canonical not in env_map:
+            continue
+        value = env_map[canonical]
+        alias = re.sub(r"[^A-Za-z0-9_]", "_", name)
+        if alias and alias not in env_map:
+            env_map[alias] = value
+        if alias == name:
+            continue
+        # Only add the original name when it results in a valid identifier.
+        if re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", name) and name not in env_map:
+            env_map[name] = value
+    return env_map
+
+
 def _ps_literal(value: str) -> str:
     return "'" + value.replace("'", "''") + "'"
 
@@ -194,7 +223,7 @@ class Role:
                     name = str(var.get('name') or '').strip()
                     if not name:
                         continue
-                    key = re.sub(r"[^A-Za-z0-9_]", "_", name).upper()
+                    key = _canonical_env_key(name)
                     if key in env_map:
                         continue
                     default_val = var.get('default')
@@ -204,6 +233,7 @@ class Role:
                         env_map[key] = ""
                     else:
                         env_map[key] = str(default_val)
+                env_map = _apply_variable_aliases(env_map, variables)
                 try:
                     timeout_seconds = max(0, int(payload.get('timeout_seconds') or 0))
                 except Exception: