Adding logic to allow the plugin to attempt to get a oneshot token from moonraker.

Author: QuinnDamerell

.vscode/settings.json

@@ -185,6 +185,7 @@
         "octowebstreamwshelper",
         "oeapi",
         "oestreamboundary",
+        "oneshot",
         "openwrt",
         "opkg",
         "oprint",

moonraker_octoeverywhere/moonrakerclient.py

@@ -113,6 +113,9 @@ def __init__(self, logger:logging.Logger, config:Config, moonrakerConfigFilePath
         # This is found and set when we try to connect and we fail due to an unauthed socket.
         # It can also be set by the user in the config.
         self.MoonrakerApiKey = self.Config.GetStr(Config.MoonrakerSection, Config.MoonrakerApiKey, None, keepInConfigIfNone=True)
+        # Same idea, but sometimes we need to get the oneshot_token to access the system.
+        # This code is only valid for 5 seconds, so it will be retrieved when we need it.
+        self.OneshotToken = None
 
         # Setup the WS vars and a websocket worker thread.
         # Don't run it until StartRunningIfNotAlready is called!
@@ -507,8 +510,14 @@ def _WebSocketWorkerThread(self):
                 # We know if the WS is connected, the host and port must be correct.
                 self._UpdateMoonrakerHostAndPort()
 
-                # Create a websocket client and start it connecting.
+                # Build the URL, include the oneshot token if we have one.
                 url = "ws://"+self.MoonrakerHostAndPort+"/websocket"
+                if self.OneshotToken is not None:
+                    url += "?token="+self.OneshotToken
+                    # Since the one shot token expires after 5 seconds, we need to clear it out.
+                    self.OneshotToken = None
+
+                # Create a websocket client and start it connecting.
                 self.Logger.info("Connecting to moonraker: "+url)
                 with self.WebSocketLock:
                     self.WebSocket = Client(url,
@@ -542,13 +551,13 @@ def _WebSocketWorkerThread(self):
 
             # This will only happen if the websocket closes or there was an error.
             # Sleep for a bit so we don't spam the system with attempts.
+            # Note that if we failed auth but got a oneshot token it will expire in 5 seconds, so we need to retry quickly.
             time.sleep(2.0)
 
 
     # Based on the docs: https://moonraker.readthedocs.io/en/latest/web_api/#websocket-setup
     # After the websocket is open, we need to do this sequence to make sure the system is healthy and ready.
     def _AfterOpenReadyWaiter(self, targetWsObjRef):
-
         logCounter = 0
         self.Logger.info("Moonraker client waiting for klippy ready...")
         try:
@@ -589,16 +598,15 @@ def _AfterOpenReadyWaiter(self, targetWsObjRef):
                 if result.HasError():
                     # Check if the error is Unauthorized, in which case we need to try to get credentials.
                     if result.ErrorCode == -32602 or result.ErrorStr == "Unauthorized":
-                        self.Logger.info("Our websocket connection to moonraker needs auth, trying to get the API key...")
-                        self.MoonrakerApiKey = MoonrakerCredentialManager.Get().TryToGetApiKey()
-                        if self.MoonrakerApiKey is None:
-                            self.Logger.error("Our websocket connection to moonraker needs auth and we failed to get the API key.")
-                            raise Exception("Websocket unauthorized.")
-                        else:
+                        if self._TryToGetWebsocketAuth():
+                            # On success, shut down the websocket so we do the reconnect logic.
                             self.Logger.info("Successfully got the API key, restarting the websocket to try again using it.")
-                            # Shut down the websocket so we do the reconnect logic.
                             self._RestartWebsocket()
                             return
+                        # Since we know we will keep failing, sleep for a while to avoid spamming the logs and so the user sees this error.
+                        self.Logger.error("!!!! Moonraker auth is required, so you must re-run the OctoEverywhere installer or generate an API key in Mainsail or Fluidd and set it the octoeverywhere.conf. The octoeverywhere.conf config file can be found in /data for docker or ~/.octoeverywhere*/ for CLI installs")
+                        time.sleep(10)
+                        raise Exception("Websocket unauthorized.")
 
                     # Handle the timeout without throwing, since this happens sometimes when the system is down.
                     if result.ErrorCode == JsonRpcResponse.OE_ERROR_TIMEOUT:
@@ -645,6 +653,34 @@ def _AfterOpenReadyWaiter(self, targetWsObjRef):
             self._RestartWebsocket()
 
 
+    # Attempts to update the moonraker api key or one shot token.
+    # Returns true if it's able to get one of them.
+    def _TryToGetWebsocketAuth(self) -> bool:
+        # First, try to get an API key
+        # If we are running locally, we should be able to connect to the unix socket and always get it.
+        self.Logger.info("Our websocket connection to moonraker needs auth, trying to get the API key...")
+        newApiKey = MoonrakerCredentialManager.Get().TryToGetApiKey()
+        if newApiKey is not None:
+            # If we got a new API key, use it now.
+            self.Logger.info("Successfully got a new API key from Moonraker.")
+            self.MoonrakerApiKey = newApiKey
+            return True
+
+        # If we didn't get an new API key, try to get a oneshot token.
+        # Note that we might already have an existing Moonraker API key, so we will include it incase.
+        #
+        # For some systems we need auth for the websocket, but no auth for the oneshot token API, which makes no sense.
+        # But if that's the case, we try to get a one_shot token.
+        self.OneshotToken = MoonrakerCredentialManager.Get().TryToGetOneshotToken(self.MoonrakerApiKey)
+        if self.OneshotToken is None:
+            # If we got a one shot token, use it.
+            self.Logger.info("Successfully got a new oneshot token from Moonraker.")
+            return True
+
+        # If we got here, we failed to get either.
+        return False
+
+
     # Kills the current websocket connection. Our logic will auto try to reconnect.
     def _RestartWebsocket(self):
         with self.WebSocketLock:

moonraker_octoeverywhere/moonrakercredentailmanager.py

@@ -5,9 +5,13 @@
 import logging
 import time
 
+from typing import Optional
+
 import configparser
 
 from octoeverywhere.sentry import Sentry
+from octoeverywhere.Proto.PathTypes import PathTypes
+from octoeverywhere.octohttprequest import OctoHttpRequest
 
 # A class that handles trying to get user credentials from Moonraker if needed.
 #
@@ -46,12 +50,41 @@ def __init__(self, logger:logging.Logger, moonrakerConfigFilePath:str, isCompani
         self.IsCompanionMode = isCompanionMode
 
 
-    def TryToGetApiKey(self) -> str or None:
+    # Attempts to get the API key from moonraker. If it fails, it will return None.
+    def TryToGetOneshotToken(self, apiKey:str=None) -> Optional[str]:
+        try:
+            # If we got an API key, try to set it.
+            headers = {}
+            if apiKey is not None:
+                headers["X-Api-Key"] = apiKey
+
+            # Make the call
+            result = OctoHttpRequest.MakeHttpCall(self.Logger, "/access/oneshot_token", PathTypes.Relative, "GET", headers)
+            if result is None:
+                raise Exception("Failed to get the oneshot token from moonraker.")
+            if result.StatusCode != 200:
+                raise Exception("Failed to get the oneshot token from moonraker. "+str(result.StatusCode))
+
+            # Read the response.
+            result.ReadAllContentFromStreamResponse(self.Logger)
+            buf = result.FullBodyBuffer
+            if buf is None:
+                raise Exception("Failed to get the oneshot token from moonraker. No content.")
+
+            # Decode & parse the response.
+            jsonMsg = json.loads(buf.decode(encoding="utf-8"))
+            token = jsonMsg.get("result", None)
+            if token is None:
+                raise Exception("Failed to get the oneshot token from moonraker. No result.")
+            return str(token)
+        except Exception as e:
+            Sentry.Exception("TryToGetOneshotToken failed to get the token.", e)
+        return None
+
+
+    def TryToGetApiKey(self) -> Optional[str]:
         # If this is an companion plugin, we dont' have the moonraker config file nor can we access the UNIX socket.
         if self.IsCompanionMode:
-            self.Logger.error("!!!! Moonraker auth is required, so you must generate an API key in Mainsail or Fluidd and set it the octoeverywhere.conf file in the companion root folder. The config file can be found in /data for docker or ~/.octoeverywhere-companion for CLI installs")
-            # Since we know we will keep failing, sleep for a while to avoid spamming the logs and so the user sees this error.
-            time.sleep(10)
             return None
 
         # First, we need to find the unix socket to connect to
@@ -124,7 +157,7 @@ def TryToGetApiKey(self) -> str or None:
             return None
 
 
-    def _TryToFindUnixSocket(self) -> str or None:
+    def _TryToFindUnixSocket(self) -> Optional[str]:
 
         # First, try to parse the moonraker config to find the klipper socket path, since the moonraker socket should be similar.
         try:
@@ -185,7 +218,7 @@ def _GetParentDirectory(self, path):
         return os.path.abspath(os.path.join(path, os.pardir))
 
 
-    def _ReadSingleJsonObject(self, sock) -> str or None:
+    def _ReadSingleJsonObject(self, sock) -> Optional[str]:
         # Since sock.recv blocks, we must read each char one by one so we know when the message ends.
         # This is messy, but since it only happens very occasionally, it's fine.
         message = bytearray()

setup.py

@@ -30,7 +30,7 @@
 
 # The plugin's version. Can be overwritten within OctoPrint's internal data via __plugin_version__ in the plugin module
 # Note that this single version string is used by all of the plugins in OctoEverywhere!
-plugin_version = "4.1.0"
+plugin_version = "4.1.3"
 
 # The plugin's description. Can be overwritten within OctoPrint's internal data via __plugin_description__ in the plugin
 # module