bug fixes

Author: tmunzer-AIDE

src/mistapi/__api_request.py

@@ -253,7 +253,8 @@ def _request_with_retry(
                     "apirequest:%s:Exception occurred", method_name, exc_info=True
                 )
                 break
-        self._count += 1
+        with self._token_lock:
+            self._count += 1
         return APIResponse(url=url, response=resp, proxy_error=proxy_failed)
 
     def mist_get(self, uri: str, query: dict[str, str] | None = None) -> APIResponse:

src/mistapi/__api_response.py

@@ -11,6 +11,8 @@
 This module manages API responses
 """
 
+import re
+
 from requests import Response
 from requests.structures import CaseInsensitiveDict
 
@@ -85,7 +87,11 @@ def _check_next(self) -> None:
                             separator = "&" if "?" in uri else "?"
                             self.next = f"{uri}{separator}page={page + 1}"
                         else:
-                            self.next = uri.replace(f"page={page}", f"page={page + 1}")
+                            self.next = re.sub(
+                                rf"(?<=[?&])page={page}(?=&|$)",
+                                f"page={page + 1}",
+                                uri,
+                            )
                         logger.debug(f"apiresponse:_check_next:set next to {self.next}")
                 except ValueError:
                     logger.error(

src/mistapi/__api_session.py

@@ -580,7 +580,7 @@ def _get_api_token_data(self, apitoken) -> tuple[str | None, list | None]:
         if data_json.get("email"):
             token_type = "user"  # nosec bandit B105
 
-        for priv in data_json.get("privileges"):
+        for priv in data_json.get("privileges", []):
             tmp = {
                 "scope": priv.get("scope"),
                 "role": priv.get("role"),
@@ -715,7 +715,7 @@ def _process_login(self, retry: bool = True) -> str | None:
                     "email/password cleaned up. Restarting authentication function"
                 )
                 if retry:
-                    return self._process_login(retry)
+                    return self._process_login(retry=False)
         except requests.exceptions.ProxyError as proxy_error:
             LOGGER.critical("apisession:_process_login:proxy not valid...")
             CONSOLE.critical("Proxy not valid...\r\n")
@@ -935,9 +935,15 @@ def _set_authenticated(self, authentication_status: bool) -> None:
                     LOGGER.error(
                         "apirequest:mist_post_file: Exception occurred", exc_info=True
                     )
-                self._csrftoken = self._session.cookies["csrftoken" + cookies_ext]
-                self._session.headers.update({"X-CSRFToken": self._csrftoken})
-                LOGGER.info("apisession:_set_authenticated:CSRF Token stored")
+                csrf_cookie = self._session.cookies.get("csrftoken" + cookies_ext)
+                if csrf_cookie:
+                    self._csrftoken = csrf_cookie
+                    self._session.headers.update({"X-CSRFToken": self._csrftoken})
+                    LOGGER.info("apisession:_set_authenticated:CSRF Token stored")
+                else:
+                    LOGGER.error(
+                        "apisession:_set_authenticated:CSRF Token cookie not found"
+                    )
         elif authentication_status is False:
             self._authenticated = False
             LOGGER.info(
@@ -1093,7 +1099,7 @@ def _getself(self) -> bool:
                     for key, val in resp.data.items():
                         if key == "privileges":
                             self.privileges = Privileges(resp.data["privileges"])
-                        if key == "tags":
+                        elif key == "tags":
                             for tag in resp.data["tags"]:
                                 self.tags.append(tag)
                         else:

src/mistapi/device_utils/__tools/bgp.py

@@ -40,6 +40,8 @@ def summary(
         UUID of the site where the device is located.
     device_id : str
         UUID of the device to show BGP summary on.
+    timeout : int, default 5
+        Time in seconds to wait for data before closing the connection.
     on_message : Callable, optional
         Callback invoked with each extracted raw message as it arrives.
 

src/mistapi/device_utils/ex.py

@@ -47,10 +47,12 @@
 )
 from mistapi.device_utils.__tools.shell import interactive_shell as interactiveShell
 
-# Tools (ping, monitor traffic)
+# Tools (ping, traceroute, monitor traffic)
+from mistapi.device_utils.__tools.miscellaneous import TracerouteProtocol
 from mistapi.device_utils.__tools.miscellaneous import monitor_traffic as monitorTraffic
 from mistapi.device_utils.__tools.miscellaneous import ping
 from mistapi.device_utils.__tools.miscellaneous import top_command as topCommand
+from mistapi.device_utils.__tools.miscellaneous import traceroute
 
 # Policy functions
 from mistapi.device_utils.__tools.policy import clear_hit_count as clearHitCount
@@ -88,4 +90,6 @@
     "monitorTraffic",
     "ping",
     "topCommand",
+    "traceroute",
+    "TracerouteProtocol",
 ]

src/mistapi/device_utils/srx.py

@@ -37,10 +37,12 @@
 )
 from mistapi.device_utils.__tools.shell import interactive_shell as interactiveShell
 
-# Tools (ping, monitor traffic)
+# Tools (ping, traceroute, monitor traffic)
+from mistapi.device_utils.__tools.miscellaneous import TracerouteProtocol
 from mistapi.device_utils.__tools.miscellaneous import monitor_traffic as monitorTraffic
 from mistapi.device_utils.__tools.miscellaneous import ping
 from mistapi.device_utils.__tools.miscellaneous import top_command as topCommand
+from mistapi.device_utils.__tools.miscellaneous import traceroute
 
 # OSPF functions
 from mistapi.device_utils.__tools.ospf import show_database as retrieveOspfDatabase
@@ -52,6 +54,7 @@
 from mistapi.device_utils.__tools.port import bounce as bouncePort
 
 # Route functions
+from mistapi.device_utils.__tools.routes import RouteProtocol
 from mistapi.device_utils.__tools.routes import show as retrieveRoutes
 
 # Sessions functions
@@ -61,6 +64,8 @@
 __all__ = [
     # Classes/Enums
     "Node",
+    "RouteProtocol",
+    "TracerouteProtocol",
     # ARP
     "retrieveArpTable",
     # BGP
@@ -88,4 +93,5 @@
     "monitorTraffic",
     "ping",
     "topCommand",
+    "traceroute",
 ]

src/mistapi/device_utils/ssr.py

@@ -30,8 +30,10 @@
 from mistapi.device_utils.__tools.dhcp import release_dhcp_leases as releaseDhcpLeases
 from mistapi.device_utils.__tools.dhcp import retrieve_dhcp_leases as retrieveDhcpLeases
 
-# Tools (ping only - no monitor_traffic for SSR)
+# Tools (ping, traceroute - no monitor_traffic for SSR)
+from mistapi.device_utils.__tools.miscellaneous import TracerouteProtocol
 from mistapi.device_utils.__tools.miscellaneous import ping
+from mistapi.device_utils.__tools.miscellaneous import traceroute
 
 # DNS functions
 # from mistapi.utils.dns import test_resolution as test_dns_resolution
@@ -45,6 +47,7 @@
 from mistapi.device_utils.__tools.port import bounce as bouncePort
 
 # Route functions
+from mistapi.device_utils.__tools.routes import RouteProtocol
 from mistapi.device_utils.__tools.routes import show as retrieveRoutes
 
 # Service Path functions
@@ -59,6 +62,8 @@
 __all__ = [
     # Classes/Enums
     "Node",
+    "RouteProtocol",
+    "TracerouteProtocol",
     # ARP
     "retrieveArpTable",
     # BGP
@@ -84,4 +89,5 @@
     "clearSessions",
     # Tools
     "ping",
+    "traceroute",
 ]

src/mistapi/websockets/__ws_client.py

@@ -74,7 +74,7 @@ def __init__(
         self._on_message_cb: Callable[[dict], None] | None = None
         self._on_error_cb: Callable[[Exception], None] | None = None
         self._on_open_cb: Callable[[], None] | None = None
-        self._on_close_cb: Callable[[int, str], None] | None = None
+        self._on_close_cb: Callable[[int | None, str | None], None] | None = None
 
     # ------------------------------------------------------------------
     # Auth / URL helpers
@@ -141,7 +141,7 @@ def on_open(self, callback: Callable[[], None]) -> None:
         """Register a callback invoked when the connection is established."""
         self._on_open_cb = callback
 
-    def on_close(self, callback: Callable[[int, str], None]) -> None:
+    def on_close(self, callback: Callable[[int | None, str | None], None]) -> None:
         """Register a callback invoked when the connection closes."""
         self._on_close_cb = callback
 
@@ -176,8 +176,8 @@ def _handle_error(self, ws: websocket.WebSocketApp, error: Exception) -> None:
     def _handle_close(
         self,
         ws: websocket.WebSocketApp,
-        close_status_code: int,
-        close_msg: str,
+        close_status_code: int | None,
+        close_msg: str | None,
     ) -> None:
         self._connected.clear()
         self._last_close_code = close_status_code
@@ -208,7 +208,7 @@ def connect(self, run_in_background: bool = True) -> None:
             If True, runs the WebSocket loop in a daemon thread (non-blocking).
             If False, blocks the calling thread until disconnected.
         """
-        if self._thread is not None and self._thread.is_alive():
+        if self._connected.is_set() or (self._thread is not None and self._thread.is_alive()):
             raise RuntimeError("Already connected; call disconnect() first")
         self._user_disconnect.clear()
         self._reconnect_attempts = 0
@@ -270,10 +270,7 @@ def _run_forever_safe(self) -> None:
         # Final close: put sentinel and call callback
         self._queue.put(None)
         if self._on_close_cb:
-            # websocket-client may provide None for close code/message; normalize
-            code = self._last_close_code if self._last_close_code is not None else -1
-            msg = self._last_close_msg if self._last_close_msg is not None else ""
-            self._on_close_cb(code, msg)
+            self._on_close_cb(self._last_close_code, self._last_close_msg)
 
     def disconnect(self) -> None:
         """Close the WebSocket connection."""
@@ -291,7 +288,12 @@ def receive(self) -> Generator[dict, None, None]:
 
         Intended for use after connect(run_in_background=True).
         """
-        if not self._connected.wait(timeout=10):
+        if self._auto_reconnect:
+            while not self._connected.is_set() and not self._user_disconnect.is_set():
+                self._connected.wait(timeout=1)
+            if self._user_disconnect.is_set() and not self._connected.is_set():
+                return
+        elif not self._connected.wait(timeout=10):
             return
         while True:
             try: