diff --git a/api/tacticalrmm/checks/permissions.py b/api/tacticalrmm/checks/permissions.py index 392675e4..8a49e199 100644 --- a/api/tacticalrmm/checks/permissions.py +++ b/api/tacticalrmm/checks/permissions.py @@ -1,6 +1,11 @@ from rest_framework import permissions -from tacticalrmm.permissions import _has_perm, _has_perm_on_agent +from tacticalrmm.permissions import ( + _has_perm, + _has_perm_on_agent, + _has_perm_on_client, + _has_perm_on_site, +) class ChecksPerms(permissions.BasePermission): @@ -21,3 +26,17 @@ class RunChecksPerms(permissions.BasePermission): return _has_perm(r, "can_run_checks") and _has_perm_on_agent( r.user, view.kwargs["agent_id"] ) + + +class BulkRunChecksPerms(permissions.BasePermission): + def has_permission(self, r, view) -> bool: + if not _has_perm(r, "can_run_checks"): + return False + + if view.kwargs["target"] == "client": + return _has_perm_on_client(user=r.user, client_id=view.kwargs["pk"]) + + elif view.kwargs["target"] == "site": + return _has_perm_on_site(user=r.user, site_id=view.kwargs["pk"]) + + return False diff --git a/api/tacticalrmm/checks/urls.py b/api/tacticalrmm/checks/urls.py index 77652906..084e697b 100644 --- a/api/tacticalrmm/checks/urls.py +++ b/api/tacticalrmm/checks/urls.py @@ -8,4 +8,5 @@ urlpatterns = [ path("/reset/", views.ResetCheck.as_view()), path("/run/", views.run_checks), path("/history/", views.GetCheckHistory.as_view()), + path("//csbulkrun/", views.bulk_run_checks), ] diff --git a/api/tacticalrmm/checks/views.py b/api/tacticalrmm/checks/views.py index b534312d..3543b9c4 100644 --- a/api/tacticalrmm/checks/views.py +++ b/api/tacticalrmm/checks/views.py @@ -1,6 +1,9 @@ import asyncio from datetime import datetime as dt +from typing import TYPE_CHECKING +import msgpack +import nats from django.db.models import Q from django.shortcuts import get_object_or_404 from django.utils import timezone as djangotime @@ -14,13 +17,16 @@ from agents.models import Agent from alerts.models import Alert from automation.models import Policy from tacticalrmm.constants import CheckStatus, CheckType -from tacticalrmm.helpers import notify_error +from tacticalrmm.helpers import notify_error, setup_nats_options from tacticalrmm.permissions import _has_perm_on_agent from .models import Check, CheckHistory, CheckResult -from .permissions import ChecksPerms, RunChecksPerms +from .permissions import BulkRunChecksPerms, ChecksPerms, RunChecksPerms from .serializers import CheckHistorySerializer, CheckSerializer +if TYPE_CHECKING: + from nats.aio.client import Client as NATSClient + class GetAddChecks(APIView): permission_classes = [IsAuthenticated, ChecksPerms] @@ -171,3 +177,41 @@ def run_checks(request, agent_id): return Response(f"Checks will now be run on {agent.hostname}") return notify_error("Unable to contact the agent") + + +@api_view(["POST"]) +@permission_classes([IsAuthenticated, BulkRunChecksPerms]) +def bulk_run_checks(request, target, pk): + q = Q() + match target: + case "client": + q = Q(site__client__id=pk) + case "site": + q = Q(site__id=pk) + + agents = list( + Agent.objects.only("agent_id", "site") + .filter(q) + .values_list("agent_id", flat=True) + ) + + if not agents: + return notify_error("No agents matched query") + + async def _run_check(nc: "NATSClient", sub) -> None: + await nc.publish(subject=sub, payload=msgpack.dumps({"func": "runchecks"})) + + async def _run() -> None: + opts = setup_nats_options() + try: + nc = await nats.connect(**opts) + except Exception as e: + return notify_error(str(e)) + + tasks = [_run_check(nc=nc, sub=agent) for agent in agents] + await asyncio.gather(*tasks) + await nc.close() + + asyncio.run(_run()) + ret = f"Checks will now be run on {len(agents)} agents" + return Response(ret)