From 85166b6e8b3d966b97da026d8e2199f90679f8b4 Mon Sep 17 00:00:00 2001
From: wh1te909 <7434746+wh1te909@users.noreply.github.com>
Date: Thu, 17 Oct 2024 20:27:15 +0000
Subject: [PATCH] add run on server option to run script endpoint #1923

---
 api/tacticalrmm/agents/tests/test_agents.py | 63 ++++++++++++++++++++-
 api/tacticalrmm/agents/views.py             | 27 +++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/api/tacticalrmm/agents/tests/test_agents.py b/api/tacticalrmm/agents/tests/test_agents.py
index addcaec0..13e41bd5 100644
--- a/api/tacticalrmm/agents/tests/test_agents.py
+++ b/api/tacticalrmm/agents/tests/test_agents.py
@@ -2,7 +2,7 @@ import json
 import os
 from itertools import cycle
 from typing import TYPE_CHECKING
-from unittest.mock import patch
+from unittest.mock import PropertyMock, patch
 from zoneinfo import ZoneInfo
 
 from django.conf import settings
@@ -768,6 +768,67 @@ class TestAgentViews(TacticalTestCase):
 
         self.assertEqual(Note.objects.get(agent=self.agent).note, "ok")
 
+        # test run on server
+        with patch("core.utils.run_server_script") as mock_run_server_script:
+            mock_run_server_script.return_value = ("output", "error", 1.23456789, 0)
+            data = {
+                "script": script.pk,
+                "output": "wait",
+                "args": ["arg1", "arg2"],
+                "timeout": 15,
+                "run_as_user": False,
+                "env_vars": ["key1=val1", "key2=val2"],
+                "run_on_server": True,
+            }
+
+            r = self.client.post(url, data, format="json")
+            self.assertEqual(r.status_code, 200)
+            hist = AgentHistory.objects.filter(agent=self.agent, script=script).last()
+            if not hist:
+                raise AgentHistory.DoesNotExist
+
+            mock_run_server_script.assert_called_with(
+                body=script.script_body,
+                args=script.parse_script_args(self.agent, script.shell, data["args"]),
+                env_vars=script.parse_script_env_vars(
+                    self.agent, script.shell, data["env_vars"]
+                ),
+                shell=script.shell,
+                timeout=18,
+            )
+
+            expected_ret = {
+                "stdout": "output",
+                "stderr": "error",
+                "execution_time": "1.2346",
+                "retcode": 0,
+            }
+
+            self.assertEqual(r.data, expected_ret)
+
+            hist.refresh_from_db()
+            expected_script_results = {**expected_ret, "id": hist.pk}
+            self.assertEqual(hist.script_results, expected_script_results)
+
+            # test run on server with server scripts disabled
+            with patch(
+                "core.models.CoreSettings.server_scripts_enabled",
+                new_callable=PropertyMock,
+            ) as server_scripts_enabled:
+                server_scripts_enabled.return_value = False
+
+                data = {
+                    "script": script.pk,
+                    "output": "wait",
+                    "args": ["arg1", "arg2"],
+                    "timeout": 15,
+                    "run_as_user": False,
+                    "env_vars": ["key1=val1", "key2=val2"],
+                    "run_on_server": True,
+                }
+                r = self.client.post(url, data, format="json")
+                self.assertEqual(r.status_code, 400)
+
     def test_get_notes(self):
         url = f"{base_url}/notes/"
 
diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py
index b3ac5bb5..78ca3181 100644
--- a/api/tacticalrmm/agents/views.py
+++ b/api/tacticalrmm/agents/views.py
@@ -768,6 +768,10 @@ def run_script(request, agent_id):
     run_as_user: bool = request.data["run_as_user"]
     env_vars: list[str] = request.data["env_vars"]
     req_timeout = int(request.data["timeout"]) + 3
+    run_on_server: bool | None = request.data.get("run_on_server")
+
+    if run_on_server and not get_core_settings().server_scripts_enabled:
+        return notify_error("This feature is disabled.")
 
     AuditLog.audit_script_run(
         username=request.user.username,
@@ -784,6 +788,29 @@ def run_script(request, agent_id):
     )
     history_pk = hist.pk
 
+    if run_on_server:
+        from core.utils import run_server_script
+
+        r = run_server_script(
+            body=script.script_body,
+            args=script.parse_script_args(agent, script.shell, args),
+            env_vars=script.parse_script_env_vars(agent, script.shell, env_vars),
+            shell=script.shell,
+            timeout=req_timeout,
+        )
+
+        ret = {
+            "stdout": r[0],
+            "stderr": r[1],
+            "execution_time": "{:.4f}".format(r[2]),
+            "retcode": r[3],
+        }
+
+        hist.script_results = {**ret, "id": history_pk}
+        hist.save(update_fields=["script_results"])
+
+        return Response(ret)
+
     if output == "wait":
         r = agent.run_script(
             scriptpk=script.pk,