From cb01891adc9d6e0f241620314ccf0ab1e1ecef49 Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Sun, 17 Nov 2019 09:28:23 +0000
Subject: [PATCH] add version info to agent

---
 .../agents/migrations/0008_agent_version.py    | 18 ++++++++++++++++++
 api/djangormm/agents/models.py                 |  1 +
 api/djangormm/agents/tasks.py                  | 15 ++++++++++++++-
 api/djangormm/api/views.py                     |  3 +++
 4 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 api/djangormm/agents/migrations/0008_agent_version.py

diff --git a/api/djangormm/agents/migrations/0008_agent_version.py b/api/djangormm/agents/migrations/0008_agent_version.py
new file mode 100644
index 00000000..b652d133
--- /dev/null
+++ b/api/djangormm/agents/migrations/0008_agent_version.py
@@ -0,0 +1,18 @@
+# Generated by Django 2.2.7 on 2019-11-17 06:44
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('agents', '0007_auto_20191027_0411'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='agent',
+            name='version',
+            field=models.CharField(default='0.1.0', max_length=255),
+        ),
+    ]
diff --git a/api/djangormm/agents/models.py b/api/djangormm/agents/models.py
index 56f0a414..1fa2cccb 100644
--- a/api/djangormm/agents/models.py
+++ b/api/djangormm/agents/models.py
@@ -13,6 +13,7 @@ from django.conf import settings
 from django.contrib.postgres.fields import JSONField
 
 class Agent(models.Model):
+    version = models.CharField(default="0.1.0", max_length=255)
     operating_system = models.CharField(null=True, max_length=255)
     plat = models.CharField(max_length=255, null=True)
     plat_release = models.CharField(max_length=255, null=True)
diff --git a/api/djangormm/agents/tasks.py b/api/djangormm/agents/tasks.py
index 231e364b..2071f198 100644
--- a/api/djangormm/agents/tasks.py
+++ b/api/djangormm/agents/tasks.py
@@ -81,7 +81,7 @@ def uninstall_agent_task(pk, wait=True):
 
 
 @app.task
-def update_agent_task(pk, version="0.2.0"):
+def update_agent_task(pk, version="0.1.4"):
     app_dir = "C:\\Program Files\\TacticalAgent"
     temp_dir = "C:\\Windows\\Temp"
     agent = Agent.objects.get(pk=pk)
@@ -124,6 +124,19 @@ def update_agent_task(pk, version="0.2.0"):
         logger.error(f"{agent.hostname} unable to stop checkrunner service")
         return f"{agent.hostname} unable to stop checkrunner service"
     
+    update_version = agent.salt_api_cmd(
+        hostname=agent.hostname,
+        timeout=45,
+        func="sqlite3.modify",
+        arg=[
+            "C:\\Program Files\\TacticalAgent\\winagent\\agentdb.db",
+            f'UPDATE agentstorage SET version = "{version}"'
+        ]
+    )
+    versiondata = update_version.json()
+    if not versiondata["return"][0][agent.hostname]:
+        logger.error(f"{agent.hostname} unable to update sql version")
+        return f"{agent.hostname} unable to update sql version"
 
     resp3 = agent.salt_api_cmd(
         hostname=agent.hostname,
diff --git a/api/djangormm/api/views.py b/api/djangormm/api/views.py
index 82a07e1a..244b73c1 100644
--- a/api/djangormm/api/views.py
+++ b/api/djangormm/api/views.py
@@ -275,6 +275,7 @@ def update(request):
     cpu_info = data["cpu_info"]
     plat = data["platform"]
     plat_release = data["platform_release"]
+    version = data["version"]
 
     agent = get_object_or_404(Agent, agent_id=agent_id)
 
@@ -284,6 +285,7 @@ def update(request):
     agent.cpu_info = cpu_info
     agent.plat = plat
     agent.plat_release = plat_release
+    agent.version = version
 
     agent.save(update_fields=[
         "last_seen",
@@ -293,6 +295,7 @@ def update(request):
         "cpu_info",
         "plat",
         "plat_release",
+        "version",
     ])
 
     sync_salt_modules_task.delay(agent.pk)