From db8917a76996a99131cd6a74d8c9b09004f07737 Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Mon, 23 Nov 2020 05:09:06 +0000
Subject: [PATCH] move reboot to nats

---
 api/tacticalrmm/agents/tests.py | 13 +++++--------
 api/tacticalrmm/agents/views.py | 20 +++++++-------------
 2 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/api/tacticalrmm/agents/tests.py b/api/tacticalrmm/agents/tests.py
index 1a8b5786..9213df26 100644
--- a/api/tacticalrmm/agents/tests.py
+++ b/api/tacticalrmm/agents/tests.py
@@ -187,20 +187,17 @@ class TestAgentViews(TacticalTestCase):
 
         self.check_not_authenticated("get", url)
 
-    @patch("agents.models.Agent.salt_api_cmd")
-    def test_power_action(self, mock_ret):
+    @patch("agents.models.Agent.nats_cmd")
+    def test_power_action(self, nats_cmd):
         url = f"/agents/poweraction/"
 
         data = {"pk": self.agent.pk, "action": "rebootnow"}
-        mock_ret.return_value = True
+        nats_cmd.return_value = "ok"
         r = self.client.post(url, data, format="json")
         self.assertEqual(r.status_code, 200)
+        nats_cmd.assert_called_with({"func": "rebootnow"}, timeout=10)
 
-        mock_ret.return_value = "error"
-        r = self.client.post(url, data, format="json")
-        self.assertEqual(r.status_code, 400)
-
-        mock_ret.return_value = False
+        nats_cmd.return_value = "timeout"
         r = self.client.post(url, data, format="json")
         self.assertEqual(r.status_code, 400)
 
diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py
index df6bc472..9c8bd3ac 100644
--- a/api/tacticalrmm/agents/views.py
+++ b/api/tacticalrmm/agents/views.py
@@ -203,19 +203,13 @@ def get_event_log(request, pk, logtype, days):
 
 @api_view(["POST"])
 def power_action(request):
-    pk = request.data["pk"]
-    action = request.data["action"]
-    agent = get_object_or_404(Agent, pk=pk)
-    if action == "rebootnow":
-        logger.info(f"{agent.hostname} was scheduled for immediate reboot")
-        r = agent.salt_api_cmd(
-            timeout=30,
-            func="system.reboot",
-            arg=3,
-            kwargs={"in_seconds": True},
-        )
-    if r == "timeout" or r == "error" or (isinstance(r, bool) and not r):
-        return notify_error("Unable to contact the agent")
+    agent = get_object_or_404(Agent, pk=request.data["pk"])
+    if not agent.has_nats:
+        return notify_error("Requires agent version 1.1.0 or greater")
+    if request.data["action"] == "rebootnow":
+        r = asyncio.run(agent.nats_cmd({"func": "rebootnow"}, timeout=10))
+        if r != "ok":
+            return notify_error("Unable to contact the agent")
 
     return Response("ok")