add run on server option to run script endpoint #1923

This commit is contained in:
wh1te909 2024-10-17 20:27:15 +00:00
parent 5278599675
commit 85166b6e8b
2 changed files with 89 additions and 1 deletions

View File

@ -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/"

View File

@ -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,