diff --git a/api/tacticalrmm/agents/tests.py b/api/tacticalrmm/agents/tests.py index ff8a9be4..1a8b5786 100644 --- a/api/tacticalrmm/agents/tests.py +++ b/api/tacticalrmm/agents/tests.py @@ -83,29 +83,29 @@ class TestAgentViews(TacticalTestCase): self.check_not_authenticated("post", url) - @patch("agents.models.Agent.salt_api_cmd") - def test_ping(self, mock_ret): + @patch("agents.models.Agent.nats_cmd") + def test_ping(self, nats_cmd): url = f"/agents/{self.agent.pk}/ping/" - mock_ret.return_value = "timeout" + nats_cmd.return_value = "timeout" r = self.client.get(url) self.assertEqual(r.status_code, 200) ret = {"name": self.agent.hostname, "status": "offline"} self.assertEqual(r.json(), ret) - mock_ret.return_value = "error" + nats_cmd.return_value = "natsdown" r = self.client.get(url) self.assertEqual(r.status_code, 200) ret = {"name": self.agent.hostname, "status": "offline"} self.assertEqual(r.json(), ret) - mock_ret.return_value = True + nats_cmd.return_value = "pong" r = self.client.get(url) self.assertEqual(r.status_code, 200) ret = {"name": self.agent.hostname, "status": "online"} self.assertEqual(r.json(), ret) - mock_ret.return_value = False + nats_cmd.return_value = "asdasjdaksdasd" r = self.client.get(url) self.assertEqual(r.status_code, 200) ret = {"name": self.agent.hostname, "status": "offline"} @@ -113,34 +113,18 @@ class TestAgentViews(TacticalTestCase): self.check_not_authenticated("get", url) + @patch("agents.models.Agent.nats_cmd") @patch("agents.tasks.uninstall_agent_task.delay") - def test_uninstall(self, mock_task): + @patch("agents.views.reload_nats") + def test_uninstall(self, reload_nats, mock_task, nats_cmd): url = "/agents/uninstall/" data = {"pk": self.agent.pk} r = self.client.delete(url, data, format="json") self.assertEqual(r.status_code, 200) - mock_task.assert_called_with(self.agent.salt_id) - - self.check_not_authenticated("delete", url) - - @patch("agents.tasks.uninstall_agent_task.delay") - def test_uninstall_catch_no_user(self, mock_task): - # setup data - agent_user = User.objects.create_user( - username=self.agent.agent_id, password=User.objects.make_random_password(60) - ) - agent_token = Token.objects.create(user=agent_user) - - url = "/agents/uninstall/" - data = {"pk": self.agent.pk} - - agent_user.delete() - - r = self.client.delete(url, data, format="json") - self.assertEqual(r.status_code, 200) - + nats_cmd.assert_called_with({"func": "uninstall"}, wait=False) + reload_nats.assert_called_once() mock_task.assert_called_with(self.agent.salt_id) self.check_not_authenticated("delete", url) @@ -167,23 +151,19 @@ class TestAgentViews(TacticalTestCase): self.check_not_authenticated("get", url) - @patch("agents.models.Agent.salt_api_cmd") - def test_kill_proc(self, mock_ret): + @patch("agents.models.Agent.nats_cmd") + def test_kill_proc(self, nats_cmd): url = f"/agents/{self.agent.pk}/8234/killproc/" - mock_ret.return_value = True + nats_cmd.return_value = "ok" r = self.client.get(url) self.assertEqual(r.status_code, 200) - mock_ret.return_value = False + nats_cmd.return_value = "timeout" r = self.client.get(url) self.assertEqual(r.status_code, 400) - mock_ret.return_value = "timeout" - r = self.client.get(url) - self.assertEqual(r.status_code, 400) - - mock_ret.return_value = "error" + nats_cmd.return_value = "process doesn't exist" r = self.client.get(url) self.assertEqual(r.status_code, 400) diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py index bb35263d..df6bc472 100644 --- a/api/tacticalrmm/agents/views.py +++ b/api/tacticalrmm/agents/views.py @@ -34,7 +34,7 @@ from .tasks import uninstall_agent_task, send_agent_update_task from winupdate.tasks import bulk_check_for_updates_task from scripts.tasks import run_bulk_script_task -from tacticalrmm.utils import notify_error +from tacticalrmm.utils import notify_error, reload_nats logger.configure(**settings.LOG_CONFIG) @@ -84,6 +84,7 @@ def uninstall(request): salt_id = agent.salt_id name = agent.hostname agent.delete() + reload_nats() uninstall_agent_task.delay(salt_id) return Response(f"{name} will now be uninstalled.") @@ -165,15 +166,17 @@ def get_processes(request, pk): @api_view() def kill_proc(request, pk, pid): agent = get_object_or_404(Agent, pk=pk) - r = agent.salt_api_cmd(timeout=25, func="ps.kill_pid", arg=int(pid)) + if not agent.has_nats: + return notify_error("Requires agent version 1.1.0 or greater") + + r = asyncio.run( + agent.nats_cmd({"func": "killproc", "procpid": int(pid)}, timeout=15) + ) if r == "timeout": return notify_error("Unable to contact the agent") - elif r == "error": - return notify_error("Something went wrong") - - if isinstance(r, bool) and not r: - return notify_error("Unable to kill the process") + elif r != "ok": + return notify_error(r) return Response("ok") diff --git a/api/tacticalrmm/apiv3/views.py b/api/tacticalrmm/apiv3/views.py index 41c34814..4b41f061 100644 --- a/api/tacticalrmm/apiv3/views.py +++ b/api/tacticalrmm/apiv3/views.py @@ -33,7 +33,7 @@ from agents.tasks import ( from winupdate.tasks import check_for_updates_task from software.tasks import get_installed_software, install_chocolatey from checks.utils import bytes2human -from tacticalrmm.utils import notify_error +from tacticalrmm.utils import notify_error, reload_nats logger.configure(**settings.LOG_CONFIG) @@ -457,6 +457,8 @@ class NewAgent(APIView): else: WinUpdatePolicy(agent=agent).save() + reload_nats() + # Generate policies for new agent agent.generate_checks_from_policies() agent.generate_tasks_from_policies()