diff --git a/api/tacticalrmm/agents/models.py b/api/tacticalrmm/agents/models.py index e8de9f2c..58c4e04c 100644 --- a/api/tacticalrmm/agents/models.py +++ b/api/tacticalrmm/agents/models.py @@ -436,7 +436,7 @@ class Agent(BaseAuditModel): except Exception: return "err" - async def nats_cmd(self, data, timeout=30): + async def nats_cmd(self, data, timeout=30, wait=True): nc = NATS() options = { "servers": f"tls://{settings.ALLOWED_HOSTS[0]}:4222", @@ -450,15 +450,22 @@ class Agent(BaseAuditModel): except: return "natsdown" - try: - msg = await nc.request(self.agent_id, msgpack.dumps(data), timeout=timeout) - except ErrTimeout: - ret = "timeout" - else: - ret = msgpack.loads(msg.data) + if wait: + try: + msg = await nc.request( + self.agent_id, msgpack.dumps(data), timeout=timeout + ) + except ErrTimeout: + ret = "timeout" + else: + ret = msgpack.loads(msg.data) - await nc.close() - return ret + await nc.close() + return ret + else: + await nc.publish(self.agent_id, msgpack.dumps(data)) + await nc.flush() + await nc.close() def salt_api_cmd(self, **kwargs): diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py index 3b01bb81..24d636ff 100644 --- a/api/tacticalrmm/agents/views.py +++ b/api/tacticalrmm/agents/views.py @@ -38,7 +38,7 @@ from winupdate.serializers import WinUpdatePolicySerializer from .tasks import uninstall_agent_task, send_agent_update_task from winupdate.tasks import bulk_check_for_updates_task -from scripts.tasks import run_script_bg_task, run_bulk_script_task +from scripts.tasks import run_bulk_script_task from tacticalrmm.utils import notify_error @@ -660,10 +660,7 @@ def recover(request): def run_script(request): agent = get_object_or_404(Agent, pk=request.data["pk"]) script = get_object_or_404(Script, pk=request.data["scriptPK"]) - output = request.data["output"] - args = request.data["args"] - req_timeout = int(request.data["timeout"]) + 3 AuditLog.audit_script_run( @@ -672,48 +669,21 @@ def run_script(request): script=script.name, ) + data = { + "func": "runscript", + "timeout": request.data["timeout"], + "script_args": request.data["args"], + "payload": { + "code": script.code, + "shell": script.shell, + }, + } + if output == "wait": - r = agent.salt_api_cmd( - timeout=req_timeout, - func="win_agent.run_script", - kwargs={ - "filepath": script.filepath, - "filename": script.filename, - "shell": script.shell, - "timeout": request.data["timeout"], - "args": args, - }, - ) - - if isinstance(r, dict): - if r["stdout"]: - return Response(r["stdout"]) - elif r["stderr"]: - return Response(r["stderr"]) - else: - try: - r["retcode"] - except KeyError: - return notify_error("Something went wrong") - - return Response(f"Return code: {r['retcode']}") - - else: - if r == "timeout": - return notify_error("Unable to contact the agent") - elif r == "error": - return notify_error("Something went wrong") - else: - return notify_error(str(r)) - + r = asyncio.run(agent.nats_cmd(data, timeout=req_timeout)) + return Response(r) else: - data = { - "agentpk": agent.pk, - "scriptpk": script.pk, - "timeout": request.data["timeout"], - "args": args, - } - run_script_bg_task.delay(data) + asyncio.run(agent.nats_cmd(data, wait=False)) return Response(f"{script.name} will now be run on {agent.hostname}") diff --git a/api/tacticalrmm/scripts/tasks.py b/api/tacticalrmm/scripts/tasks.py index d3fa63ec..cff93b71 100644 --- a/api/tacticalrmm/scripts/tasks.py +++ b/api/tacticalrmm/scripts/tasks.py @@ -3,23 +3,6 @@ from agents.models import Agent from .models import Script -@app.task -def run_script_bg_task(data): - agent = Agent.objects.get(pk=data["agentpk"]) - script = Script.objects.get(pk=data["scriptpk"]) - - agent.salt_api_async( - func="win_agent.run_script", - kwargs={ - "filepath": script.filepath, - "filename": script.filename, - "shell": script.shell, - "timeout": data["timeout"], - "args": data["args"], - }, - ) - - @app.task def run_bulk_script_task(data): # for powershell and batch scripts only, workaround for salt bg script bug