From 8337346cf5f9847e1ed8fd75d716f3dd74a2b876 Mon Sep 17 00:00:00 2001
From: Danglewood <85772166+deeleeramone@users.noreply.github.com>
Date: Mon, 18 Nov 2024 18:25:54 -0800
Subject: [PATCH] cleanup

---
 .../extensions/websockets/openbb_websockets/client.py    | 6 +-----
 .../extensions/websockets/openbb_websockets/helpers.py   | 9 +++++++--
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py
index 1e412fbdb35..671eb7fe545 100644
--- a/openbb_platform/extensions/websockets/openbb_websockets/client.py
+++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py
@@ -604,10 +604,6 @@ class WebSocketClient:
             else None
         )
 
-    def _get_auth_token(self):
-        """Get the authentication token."""
-        return self._decrypt_value(self._auth_token) if self._auth_token else None
-
     def start_broadcasting(
         self,
         host: str = "127.0.0.1",
@@ -647,7 +643,7 @@ class WebSocketClient:
             f"port={open_port}",
             f"results_file={self.results_file}",
             f"table_name={self.table_name}",
-            f"auth_token={self._get_auth_token()}",
+            f"auth_token={self._decrypt_value(self._auth_token)}",
         ]
         if kwargs:
             for kwarg in kwargs:
diff --git a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py
index d86a3523a9f..78a5b2b8b39 100644
--- a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py
+++ b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py
@@ -2,7 +2,6 @@
 
 import logging
 import re
-import sys
 from typing import Any, Optional
 
 from openbb_core.app.model.abstract.error import OpenBBError
@@ -112,7 +111,7 @@ async def check_auth(name: str, auth_token: Optional[str] = None) -> bool:
         return True
     if auth_token is None:
         raise UnauthorizedError(f"Client authorization token is required for {name}.")
-    if auth_token != client._get_auth_token():
+    if auth_token != client._decrypt_value(client._auth_token):
         raise UnauthorizedError(f"Invalid client authorization token for {name}.")
     return True
 
@@ -225,6 +224,9 @@ class StdOutSink:
 
     def write(self, message):
         """Write to stdout."""
+        # pylint: disable=import-outside-toplevel
+        import sys
+
         cleaned_message = AUTH_TOKEN_FILTER.sub(r"\1********", message)
         if cleaned_message != message:
             cleaned_message = f"{cleaned_message}\n"
@@ -232,6 +234,9 @@ class StdOutSink:
 
     def flush(self):
         """Flush stdout."""
+        # pylint: disable=import-outside-toplevel
+        import sys
+
         sys.__stdout__.flush()