paulmataruso/tacticalrmm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cleanup agent update func

Browse Source
This commit is contained in:
wh1te909
2020-12-10 10:31:58 +00:00
parent 3474b1c471
commit e5dbb66d53
2 changed files with 122 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,6 +4,7 @@ from time import sleep
import random
import requests
from packaging import version as pyver
from typing import List
from django.conf import settings
@@ -14,153 +15,78 @@ from logs.models import PendingAction
logger.configure(**settings.LOG_CONFIG)
OLD_64_PY_AGENT = "https://github.com/wh1te909/winagent/releases/download/v0.11.2/winagent-v0.11.2.exe"
OLD_32_PY_AGENT = "https://github.com/wh1te909/winagent/releases/download/v0.11.2/winagent-v0.11.2-x86.exe"
def agent_update(pk: int) -> str:
agent = Agent.objects.get(pk=pk)
# skip if we can't determine the arch
if agent.arch is None:
logger.warning(f"Unable to determine arch on {agent.hostname}. Skipping.")
return "noarch"
if agent.has_nats:
if agent.pendingactions.filter(
action_type="agentupdate", status="pending"
).exists():
action = agent.pendingactions.filter(
action_type="agentupdate", status="pending"
).last()
if pyver.parse(action.details["version"]) < pyver.parse(
settings.LATEST_AGENT_VER
):
action.delete()
else:
return "pending"
PendingAction.objects.create(
agent=agent,
action_type="agentupdate",
details={
"url": agent.winagent_dl,
"version": settings.LATEST_AGENT_VER,
"inno": agent.win_inno_exe,
},
)
return "created"
# TODO
# Salt is deprecated, remove this once salt is gone
else:
agent.salt_api_async(
func="win_agent.do_agent_update_v2",
kwargs={
"inno": agent.win_inno_exe,
"url": agent.winagent_dl,
},
)
return "salt"
@app.task
def send_agent_update_task(pks, version):
assert isinstance(pks, list)
def send_agent_update_task(pks: List[int], version: str) -> None:
q = Agent.objects.filter(pk__in=pks)
agents = [i.pk for i in q if pyver.parse(i.version) < pyver.parse(version)]
agents: List[int] = [
i.pk for i in q if pyver.parse(i.version) < pyver.parse(version)
]
chunks = (agents[i : i + 30] for i in range(0, len(agents), 30))
for chunk in chunks:
for pk in chunk:
agent = Agent.objects.get(pk=pk)
# skip if we can't determine the arch
if agent.arch is None:
logger.warning(
f"Unable to determine arch on {agent.salt_id}. Skipping."
)
continue
# golang agent only backwards compatible with py agent 0.11.2
# force an upgrade to the latest python agent if version < 0.11.2
if pyver.parse(agent.version) < pyver.parse("0.11.2"):
url = OLD_64_PY_AGENT if agent.arch == "64" else OLD_32_PY_AGENT
inno = (
"winagent-v0.11.2.exe"
if agent.arch == "64"
else "winagent-v0.11.2-x86.exe"
)
else:
url = agent.winagent_dl
inno = agent.win_inno_exe
if agent.has_nats:
if agent.pendingactions.filter(
action_type="agentupdate", status="pending"
).exists():
action = agent.pendingactions.filter(
action_type="agentupdate", status="pending"
).last()
if pyver.parse(action.details["version"]) < pyver.parse(
settings.LATEST_AGENT_VER
):
action.delete()
else:
continue
PendingAction.objects.create(
agent=agent,
action_type="agentupdate",
details={
"url": agent.winagent_dl,
"version": settings.LATEST_AGENT_VER,
"inno": agent.win_inno_exe,
},
)
# TODO
# Salt is deprecated, remove this once salt is gone
else:
r = agent.salt_api_async(
func="win_agent.do_agent_update_v2",
kwargs={
"inno": inno,
"url": url,
},
)
sleep(5)
for pk in agents:
agent_update(pk)
@app.task
def auto_self_agent_update_task():
def auto_self_agent_update_task() -> None:
core = CoreSettings.objects.first()
if not core.agent_auto_update:
logger.info("Agent auto update is disabled. Skipping.")
return
q = Agent.objects.only("pk", "version")
agents = [
pks: List[int] = [
i.pk
for i in q
if pyver.parse(i.version) < pyver.parse(settings.LATEST_AGENT_VER)
]
chunks = (agents[i : i + 30] for i in range(0, len(agents), 30))
for chunk in chunks:
for pk in chunk:
agent = Agent.objects.get(pk=pk)
# skip if we can't determine the arch
if agent.arch is None:
logger.warning(
f"Unable to determine arch on {agent.salt_id}. Skipping."
)
continue
# golang agent only backwards compatible with py agent 0.11.2
# force an upgrade to the latest python agent if version < 0.11.2
if pyver.parse(agent.version) < pyver.parse("0.11.2"):
url = OLD_64_PY_AGENT if agent.arch == "64" else OLD_32_PY_AGENT
inno = (
"winagent-v0.11.2.exe"
if agent.arch == "64"
else "winagent-v0.11.2-x86.exe"
)
else:
url = agent.winagent_dl
inno = agent.win_inno_exe
if agent.has_nats:
if agent.pendingactions.filter(
action_type="agentupdate", status="pending"
).exists():
action = agent.pendingactions.filter(
action_type="agentupdate", status="pending"
).last()
if pyver.parse(action.details["version"]) < pyver.parse(
settings.LATEST_AGENT_VER
):
action.delete()
else:
continue
PendingAction.objects.create(
agent=agent,
action_type="agentupdate",
details={
"url": agent.winagent_dl,
"version": settings.LATEST_AGENT_VER,
"inno": agent.win_inno_exe,
},
)
# TODO
# Salt is deprecated, remove this once salt is gone
else:
r = agent.salt_api_async(
func="win_agent.do_agent_update_v2",
kwargs={
"inno": inno,
"url": url,
},
)
sleep(5)
for pk in pks:
agent_update(pk)
@app.task

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

@@ -5,19 +5,20 @@ from unittest.mock import patch
from model_bakery import baker
from itertools import cycle
from django.test import TestCase, override_settings
from django.conf import settings
from django.utils import timezone as djangotime
from logs.models import PendingAction
from tacticalrmm.test import TacticalTestCase
from .serializers import AgentSerializer
from winupdate.serializers import WinUpdatePolicySerializer
from .models import Agent
from .tasks import (
agent_recovery_sms_task,
auto_self_agent_update_task,
sync_salt_modules_task,
batch_sync_modules_task,
OLD_64_PY_AGENT,
OLD_32_PY_AGENT,
)
from winupdate.models import WinUpdatePolicy
@@ -786,6 +787,70 @@ class TestAgentTasks(TacticalTestCase):
self.assertEqual(ret.status, "SUCCESS")
@patch("agents.models.Agent.salt_api_async")
def test_agent_update(self, salt_api_async):
from agents.tasks import agent_update
agent_noarch = baker.make_recipe(
"agents.agent",
operating_system="Error getting OS",
version="1.1.0",
)
r = agent_update(agent_noarch.pk)
self.assertEqual(r, "noarch")
self.assertEqual(
PendingAction.objects.filter(
agent=agent_noarch, action_type="agentupdate"
).count(),
0,
)
agent64_nats = baker.make_recipe(
"agents.agent",
operating_system="Windows 10 Pro, 64 bit (build 19041.450)",
version="1.1.0",
)
r = agent_update(agent64_nats.pk)
self.assertEqual(r, "created")
action = PendingAction.objects.get(agent__pk=agent64_nats.pk)
self.assertEqual(action.action_type, "agentupdate")
self.assertEqual(action.status, "pending")
self.assertEqual(action.details["url"], settings.DL_64)
self.assertEqual(
action.details["inno"], f"winagent-v{settings.LATEST_AGENT_VER}.exe"
)
self.assertEqual(action.details["version"], settings.LATEST_AGENT_VER)
agent64_salt = baker.make_recipe(
"agents.agent",
operating_system="Windows 10 Pro, 64 bit (build 19041.450)",
version="1.0.0",
)
salt_api_async.return_value = True
r = agent_update(agent64_salt.pk)
self.assertEqual(r, "salt")
salt_api_async.assert_called_with(
func="win_agent.do_agent_update_v2",
kwargs={
"inno": f"winagent-v{settings.LATEST_AGENT_VER}.exe",
"url": settings.DL_64,
},
)
salt_api_async.reset_mock()
agent32_nats = baker.make_recipe(
"agents.agent",
operating_system="Windows 7 Professional, 32 bit (build 7601.23964)",
version="1.1.0",
)
agent32_salt = baker.make_recipe(
"agents.agent",
operating_system="Windows 7 Professional, 32 bit (build 7601.23964)",
version="1.0.0",
)
""" @patch("agents.models.Agent.salt_api_async")
@patch("agents.tasks.sleep", return_value=None)
def test_auto_self_agent_update_task(self, mock_sleep, salt_api_async):
# test 64bit golang agent
@@ -888,4 +953,4 @@ class TestAgentTasks(TacticalTestCase):
"url": OLD_32_PY_AGENT,
},
)
self.assertEqual(ret.status, "SUCCESS")
self.assertEqual(ret.status, "SUCCESS") """