From 9b52b4efd90c395a9f3cf65e4fffe19f0251861d Mon Sep 17 00:00:00 2001 From: wh1te909 Date: Tue, 24 Nov 2020 05:14:45 +0000 Subject: [PATCH] move wmi to nats --- api/tacticalrmm/agents/models.py | 5 +-- api/tacticalrmm/agents/tasks.py | 33 ++++++++++++------- api/tacticalrmm/agents/tests.py | 55 ++++++++++++++++++++------------ api/tacticalrmm/scripts/tasks.py | 23 +++++++------ 4 files changed, 68 insertions(+), 48 deletions(-) diff --git a/api/tacticalrmm/agents/models.py b/api/tacticalrmm/agents/models.py index d328ce66..3bb50d62 100644 --- a/api/tacticalrmm/agents/models.py +++ b/api/tacticalrmm/agents/models.py @@ -630,10 +630,7 @@ class Agent(BaseAuditModel): return "failed" def not_supported(self, version_added): - if pyver.parse(self.version) < pyver.parse(version_added): - return True - - return False + return pyver.parse(self.version) < pyver.parse(version_added) def delete_superseded_updates(self): try: diff --git a/api/tacticalrmm/agents/tasks.py b/api/tacticalrmm/agents/tasks.py index cfc680b1..524723a6 100644 --- a/api/tacticalrmm/agents/tasks.py +++ b/api/tacticalrmm/agents/tasks.py @@ -1,3 +1,4 @@ +import asyncio from loguru import logger from time import sleep import random @@ -55,7 +56,7 @@ def send_agent_update_task(pks, version): f"Updating {agent.salt_id} current version {agent.version} using {inno}" ) - if pyver.parse(agent.version) >= pyver.parse("1.1.0"): + if agent.has_nats: if agent.pendingactions.filter( action_type="agentupdate", status="pending" ).exists(): @@ -127,7 +128,7 @@ def auto_self_agent_update_task(): f"Updating {agent.salt_id} current version {agent.version} using {inno}" ) - if pyver.parse(agent.version) >= pyver.parse("1.1.0"): + if agent.has_nats: if agent.pendingactions.filter( action_type="agentupdate", status="pending" ).exists(): @@ -177,7 +178,11 @@ def update_salt_minion_task(): @app.task def get_wmi_detail_task(pk): agent = Agent.objects.get(pk=pk) - r = agent.salt_api_async(timeout=30, func="win_agent.local_sys_info") + if agent.has_nats: + asyncio.run(agent.nats_cmd({"func": "sysinfo"}, wait=False)) + else: + agent.salt_api_async(timeout=30, func="win_agent.local_sys_info") + return "ok" @@ -197,7 +202,7 @@ def sync_salt_modules_task(pk): def batch_sync_modules_task(): # sync modules, split into chunks of 50 agents to not overload salt agents = Agent.objects.all() - online = [i.salt_id for i in agents if i.status == "online"] + online = [i.salt_id for i in agents] chunks = (online[i : i + 50] for i in range(0, len(online), 50)) for chunk in chunks: Agent.salt_batch_async(minions=chunk, func="saltutil.sync_modules") @@ -208,15 +213,19 @@ def batch_sync_modules_task(): def batch_sysinfo_task(): # update system info using WMI agents = Agent.objects.all() - online = [ - i.salt_id - for i in agents - if not i.not_supported("0.11.0") and i.status == "online" + + agents_nats = [agent for agent in agents if agent.has_nats] + minions = [ + agent.salt_id + for agent in agents + if not agent.has_nats and pyver.parse(agent.version) >= pyver.parse("0.11.0") ] - chunks = (online[i : i + 30] for i in range(0, len(online), 30)) - for chunk in chunks: - Agent.salt_batch_async(minions=chunk, func="win_agent.local_sys_info") - sleep(10) + + if minions: + Agent.salt_batch_async(minions=minions, func="win_agent.local_sys_info") + + for agent in agents_nats: + asyncio.run(agent.nats_cmd({"func": "sysinfo"}, wait=False)) @app.task diff --git a/api/tacticalrmm/agents/tests.py b/api/tacticalrmm/agents/tests.py index 2c53f947..1e518870 100644 --- a/api/tacticalrmm/agents/tests.py +++ b/api/tacticalrmm/agents/tests.py @@ -736,13 +736,19 @@ class TestAgentTasks(TacticalTestCase): self.authenticate() self.setup_coresettings() + @patch("agents.models.Agent.nats_cmd") @patch("agents.models.Agent.salt_api_async", return_value=None) - def test_get_wmi_detail_task(self, salt_api_async): - self.agent = baker.make_recipe("agents.agent") - ret = get_wmi_detail_task.s(self.agent.pk).apply() + def test_get_wmi_detail_task(self, salt_api_async, nats_cmd): + self.agent_salt = baker.make_recipe("agents.agent", version="1.0.2") + ret = get_wmi_detail_task.s(self.agent_salt.pk).apply() salt_api_async.assert_called_with(timeout=30, func="win_agent.local_sys_info") self.assertEqual(ret.status, "SUCCESS") + self.agent_nats = baker.make_recipe("agents.agent", version="1.1.0") + ret = get_wmi_detail_task.s(self.agent_nats.pk).apply() + nats_cmd.assert_called_with({"func": "sysinfo"}, wait=False) + self.assertEqual(ret.status, "SUCCESS") + @patch("agents.models.Agent.salt_api_cmd") def test_sync_salt_modules_task(self, salt_api_cmd): self.agent = baker.make_recipe("agents.agent") @@ -765,7 +771,7 @@ class TestAgentTasks(TacticalTestCase): @patch("agents.models.Agent.salt_batch_async", return_value=None) @patch("agents.tasks.sleep", return_value=None) def test_batch_sync_modules_task(self, mock_sleep, salt_batch_async): - # chunks of 50, 60 online should run only 2 times + # chunks of 50, should run 4 times baker.make_recipe( "agents.online_agent", last_seen=djangotime.now(), _quantity=60 ) @@ -775,32 +781,41 @@ class TestAgentTasks(TacticalTestCase): _quantity=115, ) ret = batch_sync_modules_task.s().apply() - self.assertEqual(salt_batch_async.call_count, 2) + self.assertEqual(salt_batch_async.call_count, 4) self.assertEqual(ret.status, "SUCCESS") + @patch("agents.models.Agent.nats_cmd") @patch("agents.models.Agent.salt_batch_async", return_value=None) @patch("agents.tasks.sleep", return_value=None) - def test_batch_sysinfo_task(self, mock_sleep, salt_batch_async): - # chunks of 30, 70 online should run only 3 times - self.online = baker.make_recipe( - "agents.online_agent", version=settings.LATEST_AGENT_VER, _quantity=70 - ) - self.overdue = baker.make_recipe( - "agents.overdue_agent", version=settings.LATEST_AGENT_VER, _quantity=115 + def test_batch_sysinfo_task(self, mock_sleep, salt_batch_async, nats_cmd): + + self.agents_nats = baker.make_recipe( + "agents.agent", version="1.1.0", _quantity=20 ) + # test nats ret = batch_sysinfo_task.s().apply() - self.assertEqual(salt_batch_async.call_count, 3) + self.assertEqual(nats_cmd.call_count, 20) + nats_cmd.assert_called_with({"func": "sysinfo"}, wait=False) + self.assertEqual(ret.status, "SUCCESS") + + self.agents_salt = baker.make_recipe( + "agents.agent", version="1.0.2", _quantity=70 + ) + + minions = [i.salt_id for i in self.agents_salt] + + ret = batch_sysinfo_task.s().apply() + self.assertEqual(salt_batch_async.call_count, 1) + salt_batch_async.assert_called_with( + minions=minions, func="win_agent.local_sys_info" + ) self.assertEqual(ret.status, "SUCCESS") salt_batch_async.reset_mock() - [i.delete() for i in self.online] - [i.delete() for i in self.overdue] + [i.delete() for i in self.agents_salt] # test old agents, should not run - self.online_old = baker.make_recipe( - "agents.online_agent", version="0.10.2", _quantity=70 - ) - self.overdue_old = baker.make_recipe( - "agents.overdue_agent", version="0.10.2", _quantity=115 + self.agents_old = baker.make_recipe( + "agents.agent", version="0.10.2", _quantity=70 ) ret = batch_sysinfo_task.s().apply() salt_batch_async.assert_not_called() diff --git a/api/tacticalrmm/scripts/tasks.py b/api/tacticalrmm/scripts/tasks.py index 31d0320b..3f8ea182 100644 --- a/api/tacticalrmm/scripts/tasks.py +++ b/api/tacticalrmm/scripts/tasks.py @@ -60,15 +60,14 @@ def handle_bulk_script_task(scriptpk, agentpks, args, timeout): }, ) - if agents_nats: - nats_data = { - "func": "runscript", - "timeout": timeout, - "script_args": args, - "payload": { - "code": script.code, - "shell": script.shell, - }, - } - for agent in agents_nats: - asyncio.run(agent.nats_cmd(nats_data, wait=False)) + nats_data = { + "func": "runscript", + "timeout": timeout, + "script_args": args, + "payload": { + "code": script.code, + "shell": script.shell, + }, + } + for agent in agents_nats: + asyncio.run(agent.nats_cmd(nats_data, wait=False))