From a13f6f1e680224c8589a5a453d456e6ec92fdfdf Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Fri, 15 Jan 2021 10:19:01 +0000
Subject: [PATCH] move recovery to natsapi

---
 api/tacticalrmm/agents/tasks.py  | 20 ++++++++++++++++++--
 api/tacticalrmm/natsapi/views.py | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/api/tacticalrmm/agents/tasks.py b/api/tacticalrmm/agents/tasks.py
index 83a1c7fa..8127e186 100644
--- a/api/tacticalrmm/agents/tasks.py
+++ b/api/tacticalrmm/agents/tasks.py
@@ -37,7 +37,7 @@ def _check_in_full(pk: int) -> None:
 def check_in_task() -> None:
     q = Agent.objects.only("pk", "version")
     agents: List[int] = [
-        i.pk for i in q if pyver.parse(i.version) >= pyver.parse("1.1.12")
+        i.pk for i in q if pyver.parse(i.version) == pyver.parse("1.1.12")
     ]
     chunks = (agents[i : i + 50] for i in range(0, len(agents), 50))
     for chunk in chunks:
@@ -139,7 +139,9 @@ def sync_sysinfo_task():
     online = [
         i
         for i in agents
-        if pyver.parse(i.version) >= pyver.parse("1.1.3") and i.status == "online"
+        if pyver.parse(i.version) >= pyver.parse("1.1.3")
+        and pyver.parse(i.version) <= pyver.parse("1.1.12")
+        and i.status == "online"
     ]
 
     chunks = (online[i : i + 50] for i in range(0, len(online), 50))
@@ -309,6 +311,20 @@ def install_salt_task(pk: int) -> None:
     asyncio.run(agent.nats_cmd({"func": "installsalt"}, wait=False))
 
 
+@app.task
+def handle_agent_recovery_task(pk: int) -> None:
+    sleep(10)
+    from agents.models import RecoveryAction
+
+    action = RecoveryAction.objects.get(pk=pk)
+    if action.mode == "command":
+        data = {"func": "recoverycmd", "recoverycommand": action.command}
+    else:
+        data = {"func": "recover", "payload": {"mode": action.mode}}
+
+    asyncio.run(action.agent.nats_cmd(data, wait=False))
+
+
 @app.task
 def run_script_email_results_task(
     agentpk: int, scriptpk: int, nats_timeout: int, nats_data: dict, emails: List[str]
diff --git a/api/tacticalrmm/natsapi/views.py b/api/tacticalrmm/natsapi/views.py
index 8ed438ad..45a82a15 100644
--- a/api/tacticalrmm/natsapi/views.py
+++ b/api/tacticalrmm/natsapi/views.py
@@ -15,6 +15,11 @@ from agents.models import Agent
 from software.models import InstalledSoftware
 from checks.utils import bytes2human
 from agents.serializers import WinAgentSerializer
+from agents.tasks import (
+    agent_recovery_email_task,
+    agent_recovery_sms_task,
+    handle_agent_recovery_task,
+)
 
 from tacticalrmm.utils import notify_error, filter_software, SoftwareList
 
@@ -36,6 +41,28 @@ class NatsCheckIn(APIView):
         agent.version = request.data["version"]
         agent.last_seen = djangotime.now()
         agent.save(update_fields=["version", "last_seen"])
+
+        if agent.agentoutages.exists() and agent.agentoutages.last().is_active:
+            last_outage = agent.agentoutages.last()
+            last_outage.recovery_time = djangotime.now()
+            last_outage.save(update_fields=["recovery_time"])
+
+            if agent.overdue_email_alert:
+                agent_recovery_email_task.delay(pk=last_outage.pk)
+            if agent.overdue_text_alert:
+                agent_recovery_sms_task.delay(pk=last_outage.pk)
+
+        recovery = agent.recoveryactions.filter(last_run=None).last()
+        if recovery is not None:
+            recovery.last_run = djangotime.now()
+            recovery.save(update_fields=["last_run"])
+            handle_agent_recovery_task.delay(pk=recovery.pk)
+            return Response("ok")
+
+        # get any pending actions
+        if agent.pendingactions.filter(status="pending").exists():
+            agent.handle_pending_actions()
+
         return Response("ok")
 
     def put(self, request):