Rooba
/
tacticalrmm
mirror of https://github.com/amidaware/tacticalrmm.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

move wmi to nats

This commit is contained in:
wh1te909 2020-11-24 05:14:45 +00:00
parent e90e527603
commit 9b52b4efd9
4 changed files with 68 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
api/tacticalrmm/agents/models.py
View File

@ -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:

33
api/tacticalrmm/agents/tasks.py
View File

@ -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

55
api/tacticalrmm/agents/tests.py
View File

@ -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()

23
api/tacticalrmm/scripts/tasks.py
View File

@ -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))