From c3637afe696be928cbe4178f2176e078ca01c44c Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Sun, 10 Mar 2024 00:14:04 +0000
Subject: [PATCH] max websocket max size customizable

---
 api/tacticalrmm/core/management/commands/get_mesh_exe_url.py  | 4 ++--
 .../core/management/commands/initial_mesh_setup.py            | 4 ++--
 api/tacticalrmm/core/mesh_utils.py                            | 4 ++--
 api/tacticalrmm/core/utils.py                                 | 4 ++--
 api/tacticalrmm/tacticalrmm/constants.py                      | 3 ++-
 5 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py b/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py
index b1e8944b..fc0d2b88 100644
--- a/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py
+++ b/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py
@@ -5,14 +5,14 @@ import websockets
 from django.core.management.base import BaseCommand
 
 from core.utils import get_mesh_ws_url
-from tacticalrmm.constants import WS_MAX_SIZE
+from tacticalrmm.constants import TRMM_WS_MAX_SIZE
 
 
 class Command(BaseCommand):
     help = "Sets up initial mesh central configuration"
 
     async def websocket_call(self, uri):
-        async with websockets.connect(uri, max_size=WS_MAX_SIZE) as websocket:
+        async with websockets.connect(uri, max_size=TRMM_WS_MAX_SIZE) as websocket:
             # Get Invitation Link
             await websocket.send(
                 json.dumps(
diff --git a/api/tacticalrmm/core/management/commands/initial_mesh_setup.py b/api/tacticalrmm/core/management/commands/initial_mesh_setup.py
index 3a836c31..693c6004 100644
--- a/api/tacticalrmm/core/management/commands/initial_mesh_setup.py
+++ b/api/tacticalrmm/core/management/commands/initial_mesh_setup.py
@@ -6,14 +6,14 @@ from django.conf import settings
 from django.core.management.base import BaseCommand
 
 from core.utils import get_core_settings, get_mesh_ws_url
-from tacticalrmm.constants import WS_MAX_SIZE
+from tacticalrmm.constants import TRMM_WS_MAX_SIZE
 
 
 class Command(BaseCommand):
     help = "Sets up initial mesh central configuration"
 
     async def websocket_call(self, uri):
-        async with websockets.connect(uri, max_size=WS_MAX_SIZE) as websocket:
+        async with websockets.connect(uri, max_size=TRMM_WS_MAX_SIZE) as websocket:
             # Get Device groups to see if it exists
             await websocket.send(json.dumps({"action": "meshes"}))
 
diff --git a/api/tacticalrmm/core/mesh_utils.py b/api/tacticalrmm/core/mesh_utils.py
index a20069e1..cea3607b 100644
--- a/api/tacticalrmm/core/mesh_utils.py
+++ b/api/tacticalrmm/core/mesh_utils.py
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any
 import websockets
 
 from accounts.utils import is_superuser
-from tacticalrmm.constants import WS_MAX_SIZE
+from tacticalrmm.constants import TRMM_WS_MAX_SIZE
 from tacticalrmm.helpers import make_random_password
 from tacticalrmm.logger import logger
 
@@ -86,7 +86,7 @@ class MeshSync:
         self, *, payload: dict[str, Any], wait=True
     ) -> dict[str, Any] | None:
         async def _do(payload):
-            async with websockets.connect(self.uri, max_size=WS_MAX_SIZE) as ws:
+            async with websockets.connect(self.uri, max_size=TRMM_WS_MAX_SIZE) as ws:
                 await ws.send(json.dumps(payload))
                 if wait:
                     while 1:
diff --git a/api/tacticalrmm/core/utils.py b/api/tacticalrmm/core/utils.py
index df0d2c7e..23eed4a9 100644
--- a/api/tacticalrmm/core/utils.py
+++ b/api/tacticalrmm/core/utils.py
@@ -16,7 +16,7 @@ from tacticalrmm.constants import (
     AGENT_TBL_PEND_ACTION_CNT_CACHE_PREFIX,
     CORESETTINGS_CACHE_KEY,
     ROLE_CACHE_PREFIX,
-    WS_MAX_SIZE,
+    TRMM_WS_MAX_SIZE,
     AgentPlat,
     MeshAgentIdent,
 )
@@ -100,7 +100,7 @@ def get_mesh_ws_url() -> str:
 
 
 async def get_mesh_device_id(uri: str, device_group: str) -> None:
-    async with websockets.connect(uri, max_size=WS_MAX_SIZE) as ws:
+    async with websockets.connect(uri, max_size=TRMM_WS_MAX_SIZE) as ws:
         payload = {"action": "meshes", "responseid": "meshctrl"}
         await ws.send(json.dumps(payload))
 
diff --git a/api/tacticalrmm/tacticalrmm/constants.py b/api/tacticalrmm/tacticalrmm/constants.py
index 4292b9aa..c15ee099 100644
--- a/api/tacticalrmm/tacticalrmm/constants.py
+++ b/api/tacticalrmm/tacticalrmm/constants.py
@@ -1,6 +1,7 @@
 import zoneinfo
 from enum import Enum
 
+from django.conf import settings
 from django.db import models
 
 
@@ -32,7 +33,7 @@ AGENT_OUTAGES_LOCK = "agent-outages-task-lock-key"
 ORPHANED_WIN_TASK_LOCK = "orphaned-win-task-lock-key"
 SYNC_MESH_PERMS_TASK_LOCK = "sync-mesh-perms-lock-key"
 
-WS_MAX_SIZE = 100 * 2**20
+TRMM_WS_MAX_SIZE = getattr(settings, "TRMM_WS_MAX_SIZE", 100 * 2**20)
 
 
 class GoArch(models.TextChoices):