fix: Check for and remove expired tokens before using them (#228)

Author: sean-sinclair

.gitignore

@@ -4,6 +4,7 @@ __pycache__/
 *~
 _venv
 venv
+.venv
 dist
 build
 *.egg-info/

src/sumo/wrapper/_auth_provider.py

@@ -4,7 +4,7 @@
 import stat
 import sys
 import time
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from urllib.parse import parse_qs
 
 import jwt
@@ -91,35 +91,6 @@ def store_shared_access_key_for_case(self, case_uuid, token):
         protect_token_cache(self._resource_id, ".sharedkey", case_uuid)
         return
 
-    def cleanup_shared_keys(self):
-        tokendir = get_token_dir()
-        if not os.path.exists(tokendir):
-            return
-        for f in os.listdir(tokendir):
-            ff = os.path.join(tokendir, f)
-            if os.path.isfile(ff):
-                (name, ext) = os.path.splitext(ff)
-                if ext.lower() == ".sharedkey":
-                    try:
-                        with open(ff, "r") as file:
-                            token = file.read()
-                            pq = parse_qs(token)
-                            se = pq["se"][0]
-                            end = datetime.strptime(
-                                se, "%Y-%m-%dT%H:%M:%S.%fZ"
-                            )
-                            now = datetime.utcnow()
-                            if now > end:
-                                os.unlink(ff)
-                                pass
-                            pass
-                        pass
-                    except Exception:
-                        pass
-                pass
-            pass
-        return
-
     def has_case_token(self, case_uuid):
         return os.path.exists(
             get_token_path(self._resource_id, ".sharedkey", case_uuid)
@@ -487,3 +458,31 @@ def get_auth_provider(
         ]
     ):
         return AuthProviderManaged(resource_id)
+
+
+def cleanup_shared_keys():
+    tokendir = get_token_dir()
+    if not os.path.exists(tokendir):
+        return
+    for f in os.listdir(tokendir):
+        ff = os.path.join(tokendir, f)
+        if os.path.isfile(ff):
+            (name, ext) = os.path.splitext(ff)
+            if ext.lower() == ".sharedkey":
+                try:
+                    with open(ff, "r") as file:
+                        token = file.read()
+                        pq = parse_qs(token)
+                        se = pq["se"][0]
+                        end = datetime.strptime(se, "%Y-%m-%dT%H:%M:%S.%fZ")
+                        now = datetime.now(timezone.utc)
+                        if now.timestamp() > end.timestamp():
+                            os.unlink(ff)
+                            pass
+                        pass
+                    pass
+                except Exception:
+                    pass
+            pass
+        pass
+    return

src/sumo/wrapper/sumo_client.py

@@ -6,7 +6,7 @@
 import httpx
 import jwt
 
-from ._auth_provider import get_auth_provider
+from ._auth_provider import cleanup_shared_keys, get_auth_provider
 from ._blob_client import BlobClient
 from ._decorators import (
     raise_for_status,
@@ -94,6 +94,9 @@ def __init__(
                 refresh_token = token
                 pass
             pass
+
+        cleanup_shared_keys()
+
         self.auth = get_auth_provider(
             client_id=APP_REGISTRATION[env]["CLIENT_ID"],
             authority=f"{AUTHORITY_HOST_URI}/{TENANT_ID}",
@@ -105,8 +108,6 @@ def __init__(
             case_uuid=case_uuid,
         )
 
-        self.auth.cleanup_shared_keys()
-
         if env == "localhost":
             self.base_url = "http://localhost:8084/api/v1"
         else: