From 5835632dab53c335699ca489642aafde5ee8a19d Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Wed, 5 May 2021 06:50:25 +0000
Subject: [PATCH] add button to force code signing

---
 api/tacticalrmm/agents/tasks.py               | 46 +++++++++++++------
 api/tacticalrmm/core/views.py                 | 17 +++++++
 .../modals/coresettings/CodeSign.vue          | 24 +++++++++-
 3 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/api/tacticalrmm/agents/tasks.py b/api/tacticalrmm/agents/tasks.py
index b2ea94a1..4324b977 100644
--- a/api/tacticalrmm/agents/tasks.py
+++ b/api/tacticalrmm/agents/tasks.py
@@ -20,7 +20,7 @@ from tacticalrmm.utils import run_nats_api_cmd
 logger.configure(**settings.LOG_CONFIG)
 
 
-def agent_update(pk: int, codesigntoken: str = None) -> str:
+def agent_update(pk: int, codesigntoken: str = None, force: bool = False) -> str:
     from agents.utils import get_exegen_url
 
     agent = Agent.objects.get(pk=pk)
@@ -45,22 +45,23 @@ def agent_update(pk: int, codesigntoken: str = None) -> str:
     else:
         url = agent.winagent_dl
 
-    if agent.pendingactions.filter(
-        action_type="agentupdate", status="pending"
-    ).exists():
-        agent.pendingactions.filter(
+    if not force:
+        if agent.pendingactions.filter(
             action_type="agentupdate", status="pending"
-        ).delete()
+        ).exists():
+            agent.pendingactions.filter(
+                action_type="agentupdate", status="pending"
+            ).delete()
 
-    PendingAction.objects.create(
-        agent=agent,
-        action_type="agentupdate",
-        details={
-            "url": url,
-            "version": version,
-            "inno": inno,
-        },
-    )
+        PendingAction.objects.create(
+            agent=agent,
+            action_type="agentupdate",
+            details={
+                "url": url,
+                "version": version,
+                "inno": inno,
+            },
+        )
 
     nats_data = {
         "func": "agentupdate",
@@ -74,6 +75,21 @@ def agent_update(pk: int, codesigntoken: str = None) -> str:
     return "created"
 
 
+@app.task
+def force_code_sign(pks: list[int]) -> None:
+    try:
+        token = CodeSignToken.objects.first().token
+    except:
+        return
+
+    chunks = (pks[i : i + 50] for i in range(0, len(pks), 50))
+    for chunk in chunks:
+        for pk in chunk:
+            agent_update(pk=pk, codesigntoken=token, force=True)
+            sleep(0.05)
+        sleep(4)
+
+
 @app.task
 def send_agent_update_task(pks: list[int]) -> None:
     try:
diff --git a/api/tacticalrmm/core/views.py b/api/tacticalrmm/core/views.py
index 025ec5d7..e2b531c6 100644
--- a/api/tacticalrmm/core/views.py
+++ b/api/tacticalrmm/core/views.py
@@ -234,6 +234,23 @@ class CodeSign(APIView):
             ret = "Something went wrong"
         return notify_error(ret)
 
+    def post(self, request):
+        from agents.models import Agent
+        from agents.tasks import force_code_sign
+
+        err = "A valid token must be saved first"
+        try:
+            t = CodeSignToken.objects.first().token
+        except:
+            return notify_error(err)
+
+        if t is None or t == "":
+            return notify_error(err)
+
+        pks: list[int] = list(Agent.objects.only("pk").values_list("pk", flat=True))
+        force_code_sign.delay(pks=pks)
+        return Response("Agents will be code signed shortly")
+
 
 class GetAddKeyStore(APIView):
     def get(self, request):
diff --git a/web/src/components/modals/coresettings/CodeSign.vue b/web/src/components/modals/coresettings/CodeSign.vue
index d08476b5..7986ff65 100644
--- a/web/src/components/modals/coresettings/CodeSign.vue
+++ b/web/src/components/modals/coresettings/CodeSign.vue
@@ -5,6 +5,17 @@
       <q-space />
       <q-btn icon="close" flat round dense v-close-popup />
     </q-card-section>
+    <q-card-section class="row">
+      <q-btn
+        :disable="!settings.token"
+        label="Code sign all agents"
+        color="positive"
+        class="full-width"
+        @click="doCodeSign"
+      >
+        <q-tooltip>Force all existing agents to be updated to the code-signed version</q-tooltip>
+      </q-btn>
+    </q-card-section>
     <q-form @submit.prevent="editToken">
       <q-card-section class="row">
         <div class="col-2">Token:</div>
@@ -52,12 +63,23 @@ export default {
         .then(r => {
           this.$q.loading.hide();
           this.notifySuccess(r.data);
-          this.$emit("close");
         })
         .catch(e => {
           this.$q.loading.hide();
         });
     },
+    doCodeSign() {
+      this.$q.loading.show();
+      this.$axios
+        .post("/core/codesign/")
+        .then(r => {
+          this.$q.loading.hide();
+          this.notifySuccess(r.data);
+        })
+        .catch(() => {
+          this.$q.loading.hide();
+        });
+    },
   },
   created() {
     this.getToken();