From c943ec827158d3cca8b860f5e59557f829b934b2 Mon Sep 17 00:00:00 2001
From: Sean Sinclair <146738689+sean-sinclair@users.noreply.github.com>
Date: Wed, 5 Feb 2025 10:28:38 +0100
Subject: [PATCH] fix: Check for and remove expired tokens before using them

---
 .gitignore                         |  1 +
 src/sumo/wrapper/_auth_provider.py | 59 +++++++++++++++---------------
 src/sumo/wrapper/sumo_client.py    |  7 ++--
 3 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/.gitignore b/.gitignore
index a1ed3c9..a10b609 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@ __pycache__/
 *~
 _venv
 venv
+.venv
 dist
 build
 *.egg-info/
diff --git a/src/sumo/wrapper/_auth_provider.py b/src/sumo/wrapper/_auth_provider.py
index 725ad03..fe9accb 100644
--- a/src/sumo/wrapper/_auth_provider.py
+++ b/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
diff --git a/src/sumo/wrapper/sumo_client.py b/src/sumo/wrapper/sumo_client.py
index cfdaae2..e9e3712 100644
--- a/src/sumo/wrapper/sumo_client.py
+++ b/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: