From a12f0feb668d133e3e5b2dd485501b961f96aa7b Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Fri, 26 Jan 2024 07:26:50 +0000
Subject: [PATCH] rework nats

---
 .../commands/create_natsapi_conf.py           |  6 +--
 api/tacticalrmm/tacticalrmm/helpers.py        | 38 ++++++++++++++++---
 api/tacticalrmm/tacticalrmm/utils.py          |  4 ++
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/api/tacticalrmm/core/management/commands/create_natsapi_conf.py b/api/tacticalrmm/core/management/commands/create_natsapi_conf.py
index 9d74632f45..3335cd475a 100644
--- a/api/tacticalrmm/core/management/commands/create_natsapi_conf.py
+++ b/api/tacticalrmm/core/management/commands/create_natsapi_conf.py
@@ -4,7 +4,7 @@
 from django.conf import settings
 from django.core.management.base import BaseCommand
 
-from tacticalrmm.helpers import get_nats_internal_protocol, get_nats_ports
+from tacticalrmm.helpers import get_nats_url
 
 
 class Command(BaseCommand):
@@ -20,11 +20,9 @@ def handle(self, *args, **kwargs):
         else:
             ssl = "disable"
 
-        nats_std_port, _ = get_nats_ports()
-        proto = get_nats_internal_protocol()
         config = {
             "key": settings.SECRET_KEY,
-            "natsurl": f"{proto}://{settings.ALLOWED_HOSTS[0]}:{nats_std_port}",
+            "natsurl": get_nats_url(),
             "user": db["USER"],
             "pass": db["PASSWORD"],
             "host": db["HOST"],
diff --git a/api/tacticalrmm/tacticalrmm/helpers.py b/api/tacticalrmm/tacticalrmm/helpers.py
index e8aec60303..5e9fed254d 100644
--- a/api/tacticalrmm/tacticalrmm/helpers.py
+++ b/api/tacticalrmm/tacticalrmm/helpers.py
@@ -1,3 +1,4 @@
+import os
 import random
 import secrets
 import string
@@ -43,10 +44,37 @@ def get_nats_ports() -> tuple[int, int]:
 
 
 def get_nats_internal_protocol() -> str:
-    if getattr(settings, "TRMM_INSECURE", False):
-        return "nats"
+    if getattr(settings, "USE_NATS_STANDARD", False):
+        return "tls"
 
-    return "tls"
+    return "nats"
+
+
+def get_nats_hosts() -> tuple[str, str]:
+    std_host = "0.0.0.0"
+    ws_host = "0.0.0.0"
+
+    if not settings.DOCKER_BUILD:
+        std_host, ws_host = "localhost", "localhost"
+
+    if "NATS_STD_HOST" in os.environ:
+        std_host = os.getenv("NATS_STD_HOST")
+    elif hasattr(settings, "NATS_STD_HOST"):
+        std_host = settings.NATS_STD_HOST
+
+    if "NATS_WS_HOST" in os.environ:
+        ws_host = os.getenv("NATS_WS_HOST")
+    elif hasattr(settings, "NATS_WS_HOST"):
+        ws_host = settings.NATS_WS_HOST
+
+    return std_host, ws_host
+
+
+def get_nats_url() -> str:
+    host, _ = get_nats_hosts()
+    proto = get_nats_internal_protocol()
+    port, _ = get_nats_ports()
+    return f"{proto}://{host}:{port}"
 
 
 def date_is_in_past(*, datetime_obj: "datetime", agent_tz: str) -> bool:
@@ -72,10 +100,8 @@ def rand_range(min: int, max: int) -> float:
 
 
 def setup_nats_options() -> dict[str, Any]:
-    nats_std_port, _ = get_nats_ports()
-    proto = get_nats_internal_protocol()
     opts = {
-        "servers": f"{proto}://{settings.ALLOWED_HOSTS[0]}:{nats_std_port}",
+        "servers": get_nats_url(),
         "user": "tacticalrmm",
         "name": "trmm-django",
         "password": settings.SECRET_KEY,
diff --git a/api/tacticalrmm/tacticalrmm/utils.py b/api/tacticalrmm/tacticalrmm/utils.py
index 84d07fc2a4..ea80fa75bc 100644
--- a/api/tacticalrmm/tacticalrmm/utils.py
+++ b/api/tacticalrmm/tacticalrmm/utils.py
@@ -33,6 +33,7 @@
 )
 from tacticalrmm.helpers import (
     get_certs,
+    get_nats_hosts,
     get_nats_internal_protocol,
     get_nats_ports,
     notify_error,
@@ -206,13 +207,16 @@ def reload_nats() -> None:
             )
 
     cert_file, key_file = get_certs()
+    nats_std_host, nats_ws_host = get_nats_hosts()
     nats_std_port, nats_ws_port = get_nats_ports()
 
     config = {
         "authorization": {"users": users},
         "max_payload": 67108864,
+        "host": nats_std_host,
         "port": nats_std_port,  # internal only
         "websocket": {
+            "host": nats_ws_host,
             "port": nats_ws_port,
             "no_tls": True,  # TLS is handled by nginx, so not needed here
         },