rework mesh sync #182

api/tacticalrmm/core/mesh_utils.py

@@ -1,5 +1,7 @@
 import asyncio
 import json
+import re
+import traceback
 from typing import TYPE_CHECKING, Any
 
 import websockets
@@ -31,94 +33,157 @@ def build_mesh_display_name(
     return ret
 
 
-def mesh_action(
-    *, payload: dict[str, Any], uri: str, wait=True
-) -> dict[str, Any] | None:
-    async def _do(payload, uri: str):
-        async with websockets.connect(uri, max_size=WS_MAX_SIZE) as ws:
-            await ws.send(json.dumps(payload))
-            if wait:
-                async for message in ws:
-                    r = json.loads(message)
-                    if r["action"] == payload["action"]:
-                        return r
-            else:
-                return None
-
-    payload["responseid"] = "meshctrl"
-    logger.debug(payload)
-
-    return asyncio.run(_do(payload, uri))
-
-
-def update_mesh_displayname(*, user_info: dict[str, Any], uri: str) -> None:
-    payload = {
-        "action": "edituser",
-        "id": user_info["_id"],
-        "realname": user_info["full_name"],
-    }
-    mesh_action(payload=payload, uri=uri, wait=False)
-    logger.debug(
-        f"Updating user {user_info['username']} display name to: {user_info['full_name']}"
-    )
-
-
-def add_user_to_mesh(*, user_info: dict[str, Any], uri: str) -> None:
-    payload = {
-        "action": "adduser",
-        "username": user_info["username"],
-        "email": user_info["email"],
-        "pass": make_random_password(len=30),
-        "resetNextLogin": False,
-        "randomPassword": False,
-        "removeEvents": False,
-    }
-    mesh_action(payload=payload, uri=uri, wait=False)
-    logger.debug(f"Adding user {user_info['username']} to mesh")
-    if user_info["full_name"]:
-        update_mesh_displayname(user_info=user_info, uri=uri)
-
-
-def delete_user_from_mesh(*, mesh_user_id: str, uri: str) -> None:
-    logger.debug(f"Deleting {mesh_user_id} from mesh")
-    payload = {
-        "action": "deleteuser",
-        "userid": mesh_user_id,
-    }
-    mesh_action(payload=payload, uri=uri, wait=False)
-
-
-def add_agent_to_user(*, user_id: str, node_id: str, hostname: str, uri: str) -> None:
-    logger.debug(f"Adding agent {hostname} to {user_id}")
-    payload = {
-        "action": "adddeviceuser",
-        "nodeid": node_id,
-        "userids": [user_id],
-        "rights": 72,
-        "remove": False,
-    }
-    mesh_action(payload=payload, uri=uri, wait=False)
-
-
-def remove_agent_from_user(*, user_id: str, node_id: str, uri: str) -> None:
-    logger.debug(f"Removing agent {node_id} from {user_id}")
-    payload = {
-        "action": "adddeviceuser",
-        "nodeid": node_id,
-        "userids": [user_id],
-        "rights": 0,
-        "remove": True,
-    }
-    mesh_action(payload=payload, uri=uri, wait=False)
-
-
 def has_mesh_perms(*, user: "User") -> bool:
     if user.is_superuser or is_superuser(user):
         return True
 
     return user.role and getattr(user.role, "can_use_mesh")
 
 
-def get_mesh_users(*, uri: str) -> dict[str, Any] | None:
-    payload = {"action": "users"}
-    return mesh_action(payload=payload, uri=uri, wait=True)
+def transform_trmm(obj):
+    ret = []
+    try:
+        for node in obj:
+            node_id = node["node_id"]
+            user_ids = [link["_id"] for link in node["links"]]
+            ret.append({"node_id": node_id, "user_ids": user_ids})
+    except Exception:
+        logger.debug(traceback.format_exc)
+    return ret
+
+
+def transform_mesh(obj):
+    pattern = re.compile(r".*___\d+")
+    ret = []
+    try:
+        for _, nodes in obj.items():
+            for node in nodes:
+                node_id = node["_id"]
+                user_ids = [
+                    user_id
+                    for user_id in node["links"].keys()
+                    if pattern.match(user_id)
+                ]
+                ret.append({"node_id": node_id, "user_ids": user_ids})
+    except KeyError:
+        # will trigger on initial sync cuz no mesh users yet
+        pass
+    except Exception:
+        logger.debug(traceback.format_exc)
+    return ret
+
+
+class MeshSync:
+    def __init__(self, uri: str):
+        self.uri = uri
+        self.mesh_users = self.get_trmm_mesh_users()  # full list
+
+    def mesh_action(
+        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:
+                await ws.send(json.dumps(payload))
+                if wait:
+                    while 1:
+                        try:
+                            message = await asyncio.wait_for(ws.recv(), 120)
+                            r = json.loads(message)
+                            if r["action"] == payload["action"]:
+                                return r
+                        except asyncio.TimeoutError:
+                            logger.error("Timeout reached.")
+                            return None
+                else:
+                    return None
+
+        payload["responseid"] = "meshctrl"
+        logger.debug(payload)
+
+        return asyncio.run(_do(payload))
+
+    def get_unique_mesh_users(
+        self, trmm_agents_list: list[dict[str, Any]]
+    ) -> list[str]:
+        userids = [i["links"] for i in trmm_agents_list]
+        all_ids = [item["_id"] for sublist in userids for item in sublist]
+        return list(set(all_ids))
+
+    def get_trmm_mesh_users(self):
+        payload = {"action": "users"}
+        ret = {
+            i["_id"]: i
+            for i in self.mesh_action(payload=payload, wait=True)["users"]
+            if re.search(r".*___\d+", i["_id"])
+        }
+        return ret
+
+    def add_users_to_node(self, *, node_id: str, user_ids: list[str]):
+
+        payload = {
+            "action": "adddeviceuser",
+            "nodeid": node_id,
+            "usernames": [s.replace("user//", "") for s in user_ids],
+            "rights": 72,
+            "remove": False,
+        }
+        self.mesh_action(payload=payload, wait=False)
+
+    def delete_users_from_node(self, *, node_id: str, user_ids: list[str]):
+        payload = {
+            "action": "adddeviceuser",
+            "nodeid": node_id,
+            "userids": user_ids,
+            "rights": 0,
+            "remove": True,
+        }
+        self.mesh_action(payload=payload, wait=False)
+
+    def update_mesh_displayname(self, *, user_info: dict[str, Any]) -> None:
+        payload = {
+            "action": "edituser",
+            "id": user_info["_id"],
+            "realname": user_info["full_name"],
+        }
+        self.mesh_action(payload=payload, wait=False)
+
+    def add_user_to_mesh(self, *, user_info: dict[str, Any]) -> None:
+        payload = {
+            "action": "adduser",
+            "username": user_info["username"],
+            "email": user_info["email"],
+            "pass": make_random_password(len=30),
+            "resetNextLogin": False,
+            "randomPassword": False,
+            "removeEvents": False,
+        }
+        self.mesh_action(payload=payload, wait=False)
+        if user_info["full_name"]:
+            self.update_mesh_displayname(user_info=user_info)
+
+    def delete_user_from_mesh(self, *, mesh_user_id: str) -> None:
+        payload = {
+            "action": "deleteuser",
+            "userid": mesh_user_id,
+        }
+        self.mesh_action(payload=payload, wait=False)
+
+    def add_agent_to_user(self, *, user_id: str, node_id: str) -> None:
+        payload = {
+            "action": "adddeviceuser",
+            "nodeid": node_id,
+            "userids": [user_id],
+            "rights": 72,
+            "remove": False,
+        }
+        self.mesh_action(payload=payload, wait=False)
+
+    def remove_agent_from_user(self, *, user_id: str, node_id: str) -> None:
+        payload = {
+            "action": "adddeviceuser",
+            "nodeid": node_id,
+            "userids": [user_id],
+            "rights": 0,
+            "remove": True,
+        }
+        self.mesh_action(payload=payload, wait=False)