fix related agents method in policy model

api/tacticalrmm/automation/models.py

@@ -30,24 +30,20 @@ class Policy(BaseAuditModel):
         explicit_clients = self.server_clients.all()
         explicit_sites = self.server_sites.all()
 
-        filtered_agents_pks = list()
+        filtered_agents_pks = Policy.objects.none()
 
         for site in explicit_sites:
             if site.client not in explicit_clients:
-                filtered_agents_pks.append(
-                    Agent.objects.filter(
+                filtered_agents_pks |= Agent.objects.filter(
                     client=site.client.client,
                     site=site.site,
                     monitoring_type="server",
                 ).values_list("pk", flat=True)
-                )
 
-        for client in explicit_clients:
-            filtered_agents_pks.append(
-                Agent.objects.filter(
-                    client=client.client, monitoring_type="server"
+        filtered_agents_pks |= Agent.objects.filter(
+            client__in=[client.client for client in explicit_clients],
+            monitoring_type="server",
         ).values_list("pk", flat=True)
-            )
 
         return Agent.objects.filter(
             models.Q(pk__in=filtered_agents_pks)
@@ -59,24 +55,20 @@ class Policy(BaseAuditModel):
         explicit_clients = self.workstation_clients.all()
         explicit_sites = self.workstation_sites.all()
 
-        filtered_agents_pks = list()
+        filtered_agents_pks = Policy.objects.none()
 
         for site in explicit_sites:
             if site.client not in explicit_clients:
-                filtered_agents_pks.append(
-                    Agent.objects.filter(
+                filtered_agents_pks |= Agent.objects.filter(
                     client=site.client.client,
                     site=site.site,
                     monitoring_type="workstation",
                 ).values_list("pk", flat=True)
-                )
 
-        for client in explicit_clients:
-            filtered_agents_pks.append(
-                Agent.objects.filter(
-                    client=client.client, monitoring_type="workstation"
+        filtered_agents_pks |= Agent.objects.filter(
+            client__in=[client.client for client in explicit_clients],
+            monitoring_type="workstation",
         ).values_list("pk", flat=True)
-            )
 
         return Agent.objects.filter(
             models.Q(pk__in=filtered_agents_pks)

api/tacticalrmm/automation/tasks.py

@@ -13,14 +13,14 @@ def generate_agent_checks_from_policies_task(
     # clear: clears all policy checks first
     # create_tasks: also create tasks after checks are generated
     ###
-    policypk, clear=False, create_tasks=False
+    policypk,
+    clear=False,
+    create_tasks=False,
 ):
 
     policy = Policy.objects.get(pk=policypk)
     for agent in policy.related_agents():
-        agent.generate_checks_from_policies(
-            clear=clear
-        )
+        agent.generate_checks_from_policies(clear=clear)
         if create_tasks:
             agent.generate_tasks_from_policies(
                 clear=clear,
@@ -63,10 +63,9 @@ def update_policy_check_fields_task(checkpk):
         threshold=check.threshold,
         name=check.name,
         fails_b4_alert=check.fails_b4_alert,
-        disk=check.disk,
         ip=check.ip,
-        script=check.script,
         script_args=check.script_args,
+        timeout=check.timeout,
         pass_if_start_pending=check.pass_if_start_pending,
         pass_if_svc_not_exist=check.pass_if_svc_not_exist,
         restart_if_stopped=check.restart_if_stopped,
@@ -84,9 +83,7 @@ def update_policy_check_fields_task(checkpk):
 
 
 @app.task
-def generate_agent_tasks_from_policies_task(
-    policypk, clear=False
-):
+def generate_agent_tasks_from_policies_task(policypk, clear=False):
 
     policy = Policy.objects.get(pk=policypk)
     for agent in policy.related_agents():
@@ -94,9 +91,7 @@ def generate_agent_tasks_from_policies_task(
 
 
 @app.task
-def generate_agent_tasks_by_location_task(
-    location, mon_type, clear=False
-):
+def generate_agent_tasks_by_location_task(location, mon_type, clear=False):
 
     for agent in Agent.objects.filter(**location).filter(monitoring_type=mon_type):
         agent.generate_tasks_from_policies(clear=clear)
@@ -126,9 +121,7 @@ def update_policy_task_fields_task(taskpk, enabled):
 
     tasks = AutomatedTask.objects.filter(parent_task=taskpk)
 
-    tasks.update(
-        enabled=enabled
-    )
+    tasks.update(enabled=enabled)
 
     for autotask in tasks:
         enable_or_disable_win_task(autotask.pk, enabled)

api/tacticalrmm/automation/tests.py

@@ -521,7 +521,9 @@ class TestPolicyViews(TacticalTestCase):
         task = baker.make("autotasks.AutomatedTask", policy=policy)
 
         # create policy managed tasks
-        policy_tasks = baker.make("autotasks.AutomatedTask", parent_task=task.id, _quantity=5)
+        policy_tasks = baker.make(
+            "autotasks.AutomatedTask", parent_task=task.id, _quantity=5
+        )
 
         url = f"/automation/policyautomatedtaskstatus/{task.id}/task/"
 
@@ -823,7 +825,12 @@ class TestPolicyTasks(TacticalTestCase):
 
         # make sure each agent check says overriden_by_policy
         self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 14)
-        self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.filter(overriden_by_policy=True).count(), 7)
+        self.assertEqual(
+            Agent.objects.get(pk=agent.id)
+            .agentchecks.filter(overriden_by_policy=True)
+            .count(),
+            7,
+        )
 
     def test_generating_agent_policy_checks_by_location(self):
         from .tasks import generate_agent_checks_by_location_task
@@ -831,23 +838,46 @@ class TestPolicyTasks(TacticalTestCase):
         # setup data
         policy = baker.make("automation.Policy", active=True)
         self.create_checks(policy=policy)
-        clients = baker.make("clients.Client", client=seq("Default"), _quantity=2, server_policy=policy, workstation_policy=policy)
-        baker.make("clients.Site", client=cycle(clients), site=seq("Default"), _quantity=4)
-        server_agent = baker.make_recipe("agents.server_agent", client="Default1", site="Default1")
-        workstation_agent = baker.make_recipe("agents.workstation_agent", client="Default1", site="Default3")
+        clients = baker.make(
+            "clients.Client",
+            client=seq("Default"),
+            _quantity=2,
+            server_policy=policy,
+            workstation_policy=policy,
+        )
+        baker.make(
+            "clients.Site", client=cycle(clients), site=seq("Default"), _quantity=4
+        )
+        server_agent = baker.make_recipe(
+            "agents.server_agent", client="Default1", site="Default1"
+        )
+        workstation_agent = baker.make_recipe(
+            "agents.workstation_agent", client="Default1", site="Default3"
+        )
         agent1 = baker.make_recipe("agents.agent", client="Default2", site="Default2")
         agent2 = baker.make_recipe("agents.agent", client="Default2", site="Default4")
 
-        generate_agent_checks_by_location_task({"client": "Default1", "site": "Default1"}, "server", clear=True, create_tasks=True)
+        generate_agent_checks_by_location_task(
+            {"client": "Default1", "site": "Default1"},
+            "server",
+            clear=True,
+            create_tasks=True,
+        )
 
         # server_agent should have policy checks and the other agents should not
         self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
-        self.assertEqual(Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
         self.assertEqual(Agent.objects.get(pk=agent1.id).agentchecks.count(), 0)
 
-        generate_agent_checks_by_location_task({"client": "Default1"}, "workstation", clear=True, create_tasks=True)
+        generate_agent_checks_by_location_task(
+            {"client": "Default1"}, "workstation", clear=True, create_tasks=True
+        )
         # workstation_agent should now have policy checks and the other agents should not
-        self.assertEqual(Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 7)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 7
+        )
         self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
         self.assertEqual(Agent.objects.get(pk=agent1.id).agentchecks.count(), 0)
         self.assertEqual(Agent.objects.get(pk=agent2.id).agentchecks.count(), 0)
@@ -860,9 +890,15 @@ class TestPolicyTasks(TacticalTestCase):
         policy = baker.make("automation.Policy", active=True)
         self.create_checks(policy=policy)
         clients = baker.make("clients.Client", client=seq("Default"), _quantity=2)
-        baker.make("clients.Site", client=cycle(clients), site=seq("Default"), _quantity=4)
-        server_agent = baker.make_recipe("agents.server_agent", client="Default1", site="Default1")
-        workstation_agent = baker.make_recipe("agents.workstation_agent", client="Default1", site="Default3")
+        baker.make(
+            "clients.Site", client=cycle(clients), site=seq("Default"), _quantity=4
+        )
+        server_agent = baker.make_recipe(
+            "agents.server_agent", client="Default1", site="Default1"
+        )
+        workstation_agent = baker.make_recipe(
+            "agents.workstation_agent", client="Default1", site="Default3"
+        )
         agent1 = baker.make_recipe("agents.agent", client="Default2", site="Default2")
         agent2 = baker.make_recipe("agents.agent", client="Default2", site="Default4")
         core = CoreSettings.objects.first()
@@ -873,7 +909,9 @@ class TestPolicyTasks(TacticalTestCase):
         generate_all_agent_checks_task("server", clear=True, create_tasks=True)
 
         # all servers should have 7 checks
-        self.assertEqual(Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
         self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
         self.assertEqual(Agent.objects.get(pk=agent1.id).agentchecks.count(), 7)
         self.assertEqual(Agent.objects.get(pk=agent2.id).agentchecks.count(), 0)
@@ -881,7 +919,9 @@ class TestPolicyTasks(TacticalTestCase):
         generate_all_agent_checks_task("workstation", clear=True, create_tasks=True)
 
         # all agents should have 7 checks now
-        self.assertEqual(Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 7)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 7
+        )
         self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
         self.assertEqual(Agent.objects.get(pk=agent1.id).agentchecks.count(), 7)
         self.assertEqual(Agent.objects.get(pk=agent2.id).agentchecks.count(), 7)
@@ -894,7 +934,9 @@ class TestPolicyTasks(TacticalTestCase):
         self.create_checks(policy=policy)
         client = baker.make("clients.Client", client="Default", server_policy=policy)
         baker.make("clients.Site", client=client, site="Default")
-        agent = baker.make_recipe("agents.server_agent", client="Default", site="Default")
+        agent = baker.make_recipe(
+            "agents.server_agent", client="Default", site="Default"
+        )
         agent.generate_checks_from_policies()
 
         # make sure agent has 7 checks
@@ -907,7 +949,11 @@ class TestPolicyTasks(TacticalTestCase):
 
         # make sure policy check doesn't exist on agent
         self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 6)
-        self.assertFalse(Agent.objects.get(pk=agent.id).agentchecks.filter(parent_check=policy_check_id).exists())
+        self.assertFalse(
+            Agent.objects.get(pk=agent.id)
+            .agentchecks.filter(parent_check=policy_check_id)
+            .exists()
+        )
 
     def update_policy_check_fields(self):
         from .tasks import update_policy_check_fields_task
@@ -917,31 +963,46 @@ class TestPolicyTasks(TacticalTestCase):
         self.create_checks(policy=policy)
         client = baker.make("clients.Client", client="Default", server_policy=policy)
         baker.make("clients.Site", client=client, site="Default")
-        agent = baker.make_recipe("agents.server_agent", client="Default", site="Default")
+        agent = baker.make_recipe(
+            "agents.server_agent", client="Default", site="Default"
+        )
         agent.generate_checks_from_policies()
 
         # make sure agent has 7 checks
         self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 7)
 
         # pick a policy check and update it with new values
-        ping_check = Policy.objects.get(pk=policy.id).policychecks.filter(check_type="ping").first()
+        ping_check = (
+            Policy.objects.get(pk=policy.id)
+            .policychecks.filter(check_type="ping")
+            .first()
+        )
         ping_check.ip = "12.12.12.12"
         ping_check.save()
 
         update_policy_check_fields_task(ping_check.id)
 
         # make sure policy check was updated on the agent
-        self.assertEquals(Agent.objects.get(pk=agent.id).agentchecks.filter(parent_check=ping_check.id).ip, "12.12.12.12")
+        self.assertEquals(
+            Agent.objects.get(pk=agent.id)
+            .agentchecks.filter(parent_check=ping_check.id)
+            .ip,
+            "12.12.12.12",
+        )
 
     def test_generate_agent_tasks(self):
         from .tasks import generate_agent_tasks_from_policies_task
 
         # create test data
         policy = baker.make("automation.Policy", active=True)
-        tasks = baker.make("autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3)
+        tasks = baker.make(
+            "autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3
+        )
         client = baker.make("clients.Client", client="Default")
         baker.make("clients.Site", client=client, site="Default")
-        agent = baker.make_recipe("agents.server_agent", client="Default", site="Default", policy=policy)
+        agent = baker.make_recipe(
+            "agents.server_agent", client="Default", site="Default", policy=policy
+        )
 
         generate_agent_tasks_from_policies_task(policy.id, clear=True)
 
@@ -967,26 +1028,48 @@ class TestPolicyTasks(TacticalTestCase):
 
         # setup data
         policy = baker.make("automation.Policy", active=True)
-        tasks = baker.make("autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3)
-        clients = baker.make("clients.Client", client=seq("Default"), _quantity=2, server_policy=policy, workstation_policy=policy)
-        baker.make("clients.Site", client=cycle(clients), site=seq("Default"), _quantity=4)
-        server_agent = baker.make_recipe("agents.server_agent", client="Default1", site="Default1")
-        workstation_agent = baker.make_recipe("agents.workstation_agent", client="Default1", site="Default3")
+        tasks = baker.make(
+            "autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3
+        )
+        clients = baker.make(
+            "clients.Client",
+            client=seq("Default"),
+            _quantity=2,
+            server_policy=policy,
+            workstation_policy=policy,
+        )
+        baker.make(
+            "clients.Site", client=cycle(clients), site=seq("Default"), _quantity=4
+        )
+        server_agent = baker.make_recipe(
+            "agents.server_agent", client="Default1", site="Default1"
+        )
+        workstation_agent = baker.make_recipe(
+            "agents.workstation_agent", client="Default1", site="Default3"
+        )
         agent1 = baker.make_recipe("agents.agent", client="Default2", site="Default2")
         agent2 = baker.make_recipe("agents.agent", client="Default2", site="Default4")
 
-        generate_agent_tasks_by_location_task({"client": "Default1", "site": "Default1"}, "server", clear=True)
+        generate_agent_tasks_by_location_task(
+            {"client": "Default1", "site": "Default1"}, "server", clear=True
+        )
 
         # all servers in Default1 and site Default1 should have 3 tasks
-        self.assertEqual(Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 0)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 0
+        )
         self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3)
         self.assertEqual(Agent.objects.get(pk=agent1.id).autotasks.count(), 0)
         self.assertEqual(Agent.objects.get(pk=agent2.id).autotasks.count(), 0)
 
-        generate_agent_tasks_by_location_task({"client": "Default1"}, "workstation", clear=True)
+        generate_agent_tasks_by_location_task(
+            {"client": "Default1"}, "workstation", clear=True
+        )
 
         # all workstations in Default1 should have 3 tasks
-        self.assertEqual(Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3
+        )
         self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3)
         self.assertEqual(Agent.objects.get(pk=agent1.id).autotasks.count(), 0)
         self.assertEqual(Agent.objects.get(pk=agent2.id).autotasks.count(), 0)
@@ -999,14 +1082,15 @@ class TestPolicyTasks(TacticalTestCase):
         tasks = baker.make("autotasks.AutomatedTask", policy=policy, _quantity=3)
         client = baker.make("clients.Client", client="Default", server_policy=policy)
         baker.make("clients.Site", client=client, site="Default")
-        agent = baker.make_recipe("agents.server_agent", client="Default", site="Default")
+        agent = baker.make_recipe(
+            "agents.server_agent", client="Default", site="Default"
+        )
         agent.generate_tasks_from_policies()
 
         delete_policy_autotask_task(tasks[0].id)
 
         delete_win_task_schedule.assert_called_with(agent.autotasks.first().id)
 
-
     @patch("autotasks.tasks.run_win_task.delay")
     def test_run_policy_task(self, run_win_task):
         from .tasks import run_win_policy_autotask_task
@@ -1026,10 +1110,14 @@ class TestPolicyTasks(TacticalTestCase):
 
         # setup data
         policy = baker.make("automation.Policy", active=True)
-        tasks = baker.make("autotasks.AutomatedTask", enabled=True, policy=policy, _quantity=3)
+        tasks = baker.make(
+            "autotasks.AutomatedTask", enabled=True, policy=policy, _quantity=3
+        )
         client = baker.make("clients.Client", client="Default", server_policy=policy)
         baker.make("clients.Site", client=client, site="Default")
-        agent = baker.make_recipe("agents.server_agent", client="Default", site="Default")
+        agent = baker.make_recipe(
+            "agents.server_agent", client="Default", site="Default"
+        )
         agent.generate_tasks_from_policies()
 
         tasks[0].enabled = False

web/src/components/automation/modals/PolicyAdd.vue

@@ -73,7 +73,7 @@
         </q-select>
       </q-card-section>
       <q-card-section class="row items-center">
-        <q-btn label="Add Policies" color="primary" type="submit" />
+        <q-btn label="Edit Policies" color="primary" type="submit" />
       </q-card-section>
     </q-form>
   </q-card>