From e4b1f39fdcc06f01a7b4eff4d0a73901e95a1b23 Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Sun, 22 Nov 2020 10:32:21 +0000
Subject: [PATCH] move run script to nats

---
 api/tacticalrmm/agents/models.py | 25 +++++++++-----
 api/tacticalrmm/agents/views.py  | 58 ++++++++------------------------
 api/tacticalrmm/scripts/tasks.py | 17 ----------
 3 files changed, 30 insertions(+), 70 deletions(-)

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