Release 0.15.10

api/tacticalrmm/agents/models.py

@@ -2,7 +2,6 @@
 import re
 from collections import Counter
 from contextlib import suppress
-from distutils.version import LooseVersion
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union, cast
 
 import msgpack
@@ -16,6 +15,7 @@
 from django.utils import timezone as djangotime
 from nats.errors import TimeoutError
 from packaging import version as pyver
+from packaging.version import Version as LooseVersion
 
 from agents.utils import get_agent_url
 from checks.models import CheckResult
@@ -876,8 +876,10 @@ def delete_superseded_updates(self) -> None:
                 # extract the version from the title and sort from oldest to newest
                 # skip if no version info is available therefore nothing to parse
                 try:
+                    matches = r"(Version|Versão)"
+                    pattern = r"\(" + matches + r"(.*?)\)"
                     vers = [
-                        re.search(r"\(Version(.*?)\)", i).group(1).strip()
+                        re.search(pattern, i, flags=re.IGNORECASE).group(2).strip()
                         for i in titles
                     ]
                     sorted_vers = sorted(vers, key=LooseVersion)

api/tacticalrmm/agents/tests/test_agents.py

@@ -3,8 +3,8 @@
 from itertools import cycle
 from typing import TYPE_CHECKING
 from unittest.mock import patch
+from zoneinfo import ZoneInfo
 
-import pytz
 from django.conf import settings
 from django.utils import timezone as djangotime
 from model_bakery import baker
@@ -866,7 +866,7 @@ def test_get_agent_history(self):
 
         # test pulling data
         r = self.client.get(url, format="json")
-        ctx = {"default_tz": pytz.timezone("America/Los_Angeles")}
+        ctx = {"default_tz": ZoneInfo("America/Los_Angeles")}
         data = AgentHistorySerializer(history, many=True, context=ctx).data
         self.assertEqual(r.status_code, 200)
         self.assertEqual(r.data, data)  # type:ignore

api/tacticalrmm/agents/views.py

@@ -6,19 +6,12 @@
 from io import StringIO
 from pathlib import Path
 
-from core.utils import (
-    get_core_settings,
-    get_mesh_ws_url,
-    remove_mesh_agent,
-    token_is_valid,
-)
 from django.conf import settings
 from django.db.models import Exists, OuterRef, Prefetch, Q
 from django.http import HttpResponse
 from django.shortcuts import get_object_or_404
 from django.utils import timezone as djangotime
 from django.utils.dateparse import parse_datetime
-from logs.models import AuditLog, DebugLog, PendingAction
 from meshctrl.utils import get_login_token
 from packaging import version as pyver
 from rest_framework import serializers
@@ -27,8 +20,16 @@
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
 from rest_framework.views import APIView
+
+from core.utils import (
+    get_core_settings,
+    get_mesh_ws_url,
+    remove_mesh_agent,
+    token_is_valid,
+)
+from logs.models import AuditLog, DebugLog, PendingAction
 from scripts.models import Script
-from scripts.tasks import handle_bulk_command_task, handle_bulk_script_task
+from scripts.tasks import bulk_command_task, bulk_script_task
 from tacticalrmm.constants import (
     AGENT_DEFER,
     AGENT_STATUS_OFFLINE,
@@ -561,10 +562,11 @@ def patch(self, request, agent_id):
 @api_view(["POST"])
 @permission_classes([IsAuthenticated, InstallAgentPerms])
 def install_agent(request):
+    from knox.models import AuthToken
+
     from accounts.models import User
     from agents.utils import get_agent_url
     from core.utils import token_is_valid
-    from knox.models import AuthToken
 
     # TODO rework this ghetto validation hack
     # https://github.com/amidaware/tacticalrmm/issues/1461
@@ -947,7 +949,7 @@ def bulk(request):
     agents: list[int] = [agent.pk for agent in q]
 
     if not agents:
-        return notify_error("No agents where found meeting the selected criteria")
+        return notify_error("No agents were found meeting the selected criteria")
 
     AuditLog.audit_bulk_action(
         request.user,
@@ -962,27 +964,29 @@ def bulk(request):
         else:
             shell = request.data["shell"]
 
-        handle_bulk_command_task.delay(
-            agents,
-            request.data["cmd"],
-            shell,
-            request.data["timeout"],
-            request.user.username[:50],
-            request.data["run_as_user"],
+        bulk_command_task.delay(
+            agent_pks=agents,
+            cmd=request.data["cmd"],
+            shell=shell,
+            timeout=request.data["timeout"],
+            username=request.user.username[:50],
+            run_as_user=request.data["run_as_user"],
         )
         return Response(f"Command will now be run on {len(agents)} agents")
 
     elif request.data["mode"] == "script":
         script = get_object_or_404(Script, pk=request.data["script"])
-        handle_bulk_script_task.delay(
-            script.pk,
-            agents,
-            request.data["args"],
-            request.data["timeout"],
-            request.user.username[:50],
-            request.data["run_as_user"],
-            request.data["env_vars"],
+
+        bulk_script_task.delay(
+            script_pk=script.pk,
+            agent_pks=agents,
+            args=request.data["args"],
+            timeout=request.data["timeout"],
+            username=request.user.username[:50],
+            run_as_user=request.data["run_as_user"],
+            env_vars=request.data["env_vars"],
         )
+
         return Response(f"{script.name} will now be run on {len(agents)} agents")
 
     elif request.data["mode"] == "patch":

api/tacticalrmm/autotasks/models.py

@@ -3,8 +3,8 @@
 import string
 from contextlib import suppress
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from zoneinfo import ZoneInfo
 
-import pytz
 from django.core.cache import cache
 from django.core.validators import MaxValueValidator, MinValueValidator
 from django.db import models
@@ -270,12 +270,12 @@ def generate_nats_task_payload(
                 and self.task_type == TaskType.RUN_ONCE
                 and self.run_asap_after_missed
                 and agent
-                and self.run_time_date
-                < djangotime.now().astimezone(pytz.timezone(agent.timezone))
+                and self.run_time_date.replace(tzinfo=ZoneInfo(agent.timezone))
+                < djangotime.now().astimezone(ZoneInfo(agent.timezone))
             ):
                 self.run_time_date = (
                     djangotime.now() + djangotime.timedelta(minutes=5)
-                ).astimezone(pytz.timezone(agent.timezone))
+                ).astimezone(ZoneInfo(agent.timezone))
 
             task["start_year"] = int(self.run_time_date.strftime("%Y"))
             task["start_month"] = int(self.run_time_date.strftime("%-m"))

api/tacticalrmm/autotasks/tasks.py

@@ -149,6 +149,7 @@ async def _run() -> None:
                 for item in items
             ]
             await asyncio.gather(*tasks)
+            await nc.flush()
             await nc.close()
 
         asyncio.run(_run())

api/tacticalrmm/checks/views.py

@@ -1,9 +1,6 @@
 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
@@ -17,16 +14,15 @@
 from alerts.models import Alert
 from automation.models import Policy
 from tacticalrmm.constants import CheckStatus, CheckType
-from tacticalrmm.helpers import notify_error, setup_nats_options
+from tacticalrmm.exceptions import NatsDown
+from tacticalrmm.helpers import notify_error
+from tacticalrmm.nats_utils import abulk_nats_command
 from tacticalrmm.permissions import _has_perm_on_agent
 
 from .models import Check, CheckHistory, CheckResult
 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]
@@ -189,29 +185,22 @@ def bulk_run_checks(request, target, pk):
         case "site":
             q = Q(site__id=pk)
 
-    agents = list(
+    agent_ids = list(
         Agent.objects.only("agent_id", "site")
         .filter(q)
         .values_list("agent_id", flat=True)
     )
 
-    if not agents:
+    if not agent_ids:
         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))
+    payload = {"func": "runchecks"}
+    items = [(agent_id, payload) for agent_id in agent_ids]
 
-        tasks = [_run_check(nc=nc, sub=agent) for agent in agents]
-        await asyncio.gather(*tasks)
-        await nc.close()
+    try:
+        asyncio.run(abulk_nats_command(items=items))
+    except NatsDown as e:
+        return notify_error(str(e))
 
-    asyncio.run(_run())
-    ret = f"Checks will now be run on {len(agents)} agents"
+    ret = f"Checks will now be run on {len(agent_ids)} agents"
     return Response(ret)

api/tacticalrmm/core/migrations/0037_coresettings_open_ai_model_and_more.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2 on 2023-04-09 15:17
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('core', '0036_alter_coresettings_default_time_zone'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='coresettings',
+            name='open_ai_model',
+            field=models.CharField(blank=True, default='gpt-3.5-turbo', max_length=255),
+        ),
+        migrations.AddField(
+            model_name='coresettings',
+            name='open_ai_token',
+            field=models.CharField(blank=True, max_length=255, null=True),
+        ),
+    ]

api/tacticalrmm/core/models.py

@@ -98,6 +98,10 @@ class CoreSettings(BaseAuditModel):
     date_format = models.CharField(
         max_length=30, blank=True, default="MMM-DD-YYYY - HH:mm"
     )
+    open_ai_token = models.CharField(max_length=255, null=True, blank=True)
+    open_ai_model = models.CharField(
+        max_length=255, blank=True, default="gpt-3.5-turbo"
+    )
 
     def save(self, *args, **kwargs) -> None:
         from alerts.tasks import cache_agents_alert_template

api/tacticalrmm/core/urls.py

@@ -19,4 +19,5 @@
     path("smstest/", views.TwilioSMSTest.as_view()),
     path("clearcache/", views.clear_cache),
     path("status/", views.status),
+    path("openai/generate/", views.OpenAICodeCompletion.as_view()),
 ]

api/tacticalrmm/core/views.py

@@ -1,8 +1,10 @@
+import json
 import re
 from pathlib import Path
+from zoneinfo import ZoneInfo
 
 import psutil
-import pytz
+import requests
 from cryptography import x509
 from django.conf import settings
 from django.http import JsonResponse
@@ -12,6 +14,7 @@
 from rest_framework.decorators import api_view, permission_classes
 from rest_framework.exceptions import PermissionDenied
 from rest_framework.permissions import IsAuthenticated
+from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.views import APIView
 
@@ -77,6 +80,7 @@ def dashboard_info(request):
     from core.utils import token_is_expired
     from tacticalrmm.utils import get_latest_trmm_ver
 
+    core_settings = get_core_settings()
     return Response(
         {
             "trmm_version": settings.TRMM_VERSION,
@@ -94,8 +98,9 @@ def dashboard_info(request):
             "clear_search_when_switching": request.user.clear_search_when_switching,
             "hosted": getattr(settings, "HOSTED", False),
             "date_format": request.user.date_format,
-            "default_date_format": get_core_settings().date_format,
+            "default_date_format": core_settings.date_format,
             "token_is_expired": token_is_expired(),
+            "open_ai_integration_enabled": bool(core_settings.open_ai_token),
         }
     )
 
@@ -416,7 +421,7 @@ def status(request):
     cert_bytes = Path(cert_file).read_bytes()
 
     cert = x509.load_pem_x509_certificate(cert_bytes)
-    expires = pytz.utc.localize(cert.not_valid_after)
+    expires = cert.not_valid_after.replace(tzinfo=ZoneInfo("UTC"))
     now = djangotime.now()
     delta = expires - now
 
@@ -449,3 +454,55 @@ def status(request):
             "nginx": sysd_svc_is_running("nginx.service"),
         }
     return JsonResponse(ret, json_dumps_params={"indent": 2})
+
+
+class OpenAICodeCompletion(APIView):
+    permission_classes = [IsAuthenticated]
+
+    def post(self, request: Request) -> Response:
+        settings = get_core_settings()
+
+        if not settings.open_ai_token:
+            return notify_error(
+                "Open AI API Key not found. Open Global Settings > Open AI."
+            )
+
+        if not request.data["prompt"]:
+            return notify_error("Not prompt field found")
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {settings.open_ai_token}",
+        }
+
+        data = {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": request.data["prompt"],
+                },
+            ],
+            "model": settings.open_ai_model,
+            "temperature": 0.5,
+            "max_tokens": 1000,
+            "n": 1,
+            "stop": None,
+        }
+
+        try:
+            response = requests.post(
+                "https://api.openai.com/v1/chat/completions",
+                headers=headers,
+                data=json.dumps(data),
+            )
+        except Exception as e:
+            return notify_error(str(e))
+
+        response_data = json.loads(response.text)
+
+        if "error" in response_data:
+            return notify_error(
+                f"The Open AI API returned an error: {response_data['error']['message']}"
+            )
+
+        return Response(response_data["choices"][0]["message"]["content"])