diff --git a/api/tacticalrmm/agents/models.py b/api/tacticalrmm/agents/models.py index 2918241c..cf52a454 100644 --- a/api/tacticalrmm/agents/models.py +++ b/api/tacticalrmm/agents/models.py @@ -322,7 +322,7 @@ class Agent(BaseAuditModel): online = [ agent for agent in Agent.objects.only( - "pk", "last_seen", "overdue_time", "offline_time" + "pk", "agent_id", "last_seen", "overdue_time", "offline_time" ) if agent.status == "online" ] diff --git a/api/tacticalrmm/agents/tasks.py b/api/tacticalrmm/agents/tasks.py index 852c93ba..3d367e3e 100644 --- a/api/tacticalrmm/agents/tasks.py +++ b/api/tacticalrmm/agents/tasks.py @@ -1,9 +1,6 @@ import asyncio import datetime as dt -import json import random -import subprocess -import tempfile from time import sleep from typing import Union @@ -17,6 +14,7 @@ from core.models import CoreSettings from logs.models import PendingAction from scripts.models import Script from tacticalrmm.celery import app +from tacticalrmm.utils import run_nats_api_cmd logger.configure(**settings.LOG_CONFIG) @@ -257,30 +255,13 @@ def run_script_email_results_task( logger.error(e) -def _get_nats_config() -> dict: - return { - "key": settings.SECRET_KEY, - "natsurl": f"tls://{settings.ALLOWED_HOSTS[0]}:4222", - } - - @app.task def monitor_agents_task() -> None: agents = Agent.objects.only( "pk", "agent_id", "last_seen", "overdue_time", "offline_time" ) - ret = [i.agent_id for i in agents if i.status != "online"] - config = _get_nats_config() - config["agents"] = ret - with tempfile.NamedTemporaryFile() as fp: - with open(fp.name, "w") as f: - json.dump(config, f) - - cmd = ["/usr/local/bin/nats-api", "-c", fp.name, "-m", "monitor"] - try: - subprocess.run(cmd, capture_output=True, timeout=30) - except Exception as e: - logger.error(e) + ids = [i.agent_id for i in agents if i.status != "online"] + run_nats_api_cmd("monitor", ids) @app.task @@ -288,15 +269,5 @@ def get_wmi_task() -> None: agents = Agent.objects.only( "pk", "agent_id", "last_seen", "overdue_time", "offline_time" ) - ret = [i.agent_id for i in agents if i.status == "online"] - config = _get_nats_config() - config["agents"] = ret - with tempfile.NamedTemporaryFile() as fp: - with open(fp.name, "w") as f: - json.dump(config, f) - - cmd = ["/usr/local/bin/nats-api", "-c", fp.name, "-m", "wmi"] - try: - subprocess.run(cmd, capture_output=True, timeout=30) - except Exception as e: - logger.error(e) + ids = [i.agent_id for i in agents if i.status == "online"] + run_nats_api_cmd("wmi", ids) diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py index eea65a0e..a92b1a62 100644 --- a/api/tacticalrmm/agents/views.py +++ b/api/tacticalrmm/agents/views.py @@ -40,7 +40,7 @@ logger.configure(**settings.LOG_CONFIG) @api_view() def get_agent_versions(request): - agents = Agent.objects.only("pk") + agents = Agent.objects.prefetch_related("site").only("pk", "hostname") return Response( { "versions": [settings.LATEST_AGENT_VER], diff --git a/api/tacticalrmm/tacticalrmm/utils.py b/api/tacticalrmm/tacticalrmm/utils.py index e3bf4db6..01f6489e 100644 --- a/api/tacticalrmm/tacticalrmm/utils.py +++ b/api/tacticalrmm/tacticalrmm/utils.py @@ -228,3 +228,20 @@ class KnoxAuthMiddlewareInstance: KnoxAuthMiddlewareStack = lambda inner: KnoxAuthMiddlewareInstance( AuthMiddlewareStack(inner) ) + + +def run_nats_api_cmd(mode: str, ids: list[str], timeout: int = 30) -> None: + config = { + "key": settings.SECRET_KEY, + "natsurl": f"tls://{settings.ALLOWED_HOSTS[0]}:4222", + } + config["agents"] = ids + with tempfile.NamedTemporaryFile() as fp: + with open(fp.name, "w") as f: + json.dump(config, f) + + cmd = ["/usr/local/bin/nats-api", "-c", fp.name, "-m", mode] + try: + subprocess.run(cmd, capture_output=True, timeout=timeout) + except Exception as e: + logger.error(e)