From 6e3cad454c69311e384b67e624c8224e3dfd84d2 Mon Sep 17 00:00:00 2001
From: wh1te909 <7434746+wh1te909@users.noreply.github.com>
Date: Wed, 5 Jun 2024 00:53:50 +0000
Subject: [PATCH] add error handling for server script

---
 api/tacticalrmm/core/utils.py | 38 ++++++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/api/tacticalrmm/core/utils.py b/api/tacticalrmm/core/utils.py
index 8345977f..d399cd43 100644
--- a/api/tacticalrmm/core/utils.py
+++ b/api/tacticalrmm/core/utils.py
@@ -313,23 +313,37 @@ def run_server_script(
         var_split = var.split("=")
         custom_env[var_split[0]] = var_split[1]
 
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp_script:
+    with tempfile.NamedTemporaryFile(
+        mode="w", delete=False, prefix="trmm-"
+    ) as tmp_script:
         tmp_script.write(body.replace("\r\n", "\n"))
         tmp_script_path = tmp_script.name
 
     os.chmod(tmp_script_path, 0o550)
 
+    stdout, stderr = "", ""
+    retcode = 0
+
     start_time = time.time()
-    result = subprocess.run(
-        [tmp_script_path] + parsed_args,
-        capture_output=True,
-        text=True,
-        env=custom_env,
-        timeout=timeout,
-    )
-    execution_time = time.time() - start_time
+    try:
+        ret = subprocess.run(
+            [tmp_script_path] + parsed_args,
+            capture_output=True,
+            text=True,
+            env=custom_env,
+            timeout=timeout,
+        )
+        stdout, stderr, retcode = ret.stdout, ret.stderr, ret.returncode
+    except subprocess.TimeoutExpired:
+        stderr = f"Error: Timed out after {timeout} seconds."
+        retcode = 98
+    except Exception as e:
+        stderr = f"Error: {e}"
+        retcode = 99
+    finally:
+        execution_time = time.time() - start_time
 
-    with suppress(Exception):
-        os.remove(tmp_script_path)
+        with suppress(Exception):
+            os.remove(tmp_script_path)
 
-    return result.stdout, result.stderr, execution_time, result.returncode
+    return stdout, stderr, execution_time, retcode