From 74f476bc0eae90453025939f41c4fbbac527e55e Mon Sep 17 00:00:00 2001
From: Raymond Wiker <rayw@equinor.com>
Date: Fri, 17 Jan 2025 15:32:38 +0100
Subject: [PATCH 1/3] feat: remove stale shared keys.

---
 src/sumo/wrapper/_auth_provider.py | 32 +++++++++++++++++++++++++++---
 src/sumo/wrapper/sumo_client.py    |  2 ++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/sumo/wrapper/_auth_provider.py b/src/sumo/wrapper/_auth_provider.py
index 5b93e15..844273f 100644
--- a/src/sumo/wrapper/_auth_provider.py
+++ b/src/sumo/wrapper/_auth_provider.py
@@ -5,6 +5,7 @@
 import sys
 import time
 from datetime import datetime, timedelta
+from urllib.parse import parse_qs
 
 import jwt
 import msal
@@ -36,13 +37,12 @@ def get_token_dir():
 def get_token_path(resource_id, suffix, case_uuid=None):
     if case_uuid is not None:
         return os.path.join(
-            os.path.expanduser("~"),
-            ".sumo",
+            os.path.expanduser("~/.sumo"),
             str(resource_id) + "+" + str(case_uuid) + suffix,
         )
     else:
         return os.path.join(
-            os.path.expanduser("~"), ".sumo", str(resource_id) + suffix
+            os.path.expanduser("~/.sumo"), str(resource_id) + suffix
         )
 
 
@@ -91,6 +91,32 @@ def store_shared_access_key_for_case(self, case_uuid, token):
         ) as f:
             f.write(token)
         protect_token_cache(self._resource_id, ".sharedkey", case_uuid)
+        return
+    
+    def cleanup_shared_keys(self):
+        tokendir = os.path.join(os.path.expanduser("~/.sumo"))
+        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(
diff --git a/src/sumo/wrapper/sumo_client.py b/src/sumo/wrapper/sumo_client.py
index 383b198..72c1036 100644
--- a/src/sumo/wrapper/sumo_client.py
+++ b/src/sumo/wrapper/sumo_client.py
@@ -92,6 +92,8 @@ def __init__(
             case_uuid=case_uuid,
         )
 
+        self.auth.cleanup_shared_keys()
+
         if env == "localhost":
             self.base_url = "http://localhost:8084/api/v1"
         else:

From 8711a943bf076dfbf3d62fdea7f532ff13ab50b6 Mon Sep 17 00:00:00 2001
From: Raymond Wiker <rayw@equinor.com>
Date: Fri, 17 Jan 2025 15:54:17 +0100
Subject: [PATCH 2/3] chore: ruff format

---
 src/sumo/wrapper/_auth_provider.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/sumo/wrapper/_auth_provider.py b/src/sumo/wrapper/_auth_provider.py
index 844273f..f695e97 100644
--- a/src/sumo/wrapper/_auth_provider.py
+++ b/src/sumo/wrapper/_auth_provider.py
@@ -92,7 +92,7 @@ def store_shared_access_key_for_case(self, case_uuid, token):
             f.write(token)
         protect_token_cache(self._resource_id, ".sharedkey", case_uuid)
         return
-    
+
     def cleanup_shared_keys(self):
         tokendir = os.path.join(os.path.expanduser("~/.sumo"))
         for f in os.listdir(tokendir):
@@ -105,7 +105,9 @@ def cleanup_shared_keys(self):
                             token = file.read()
                             pq = parse_qs(token)
                             se = pq["se"][0]
-                            end = datetime.strptime(se, "%Y-%m-%dT%H:%M:%S.%fZ")
+                            end = datetime.strptime(
+                                se, "%Y-%m-%dT%H:%M:%S.%fZ"
+                            )
                             now = datetime.utcnow()
                             if now > end:
                                 os.unlink(ff)

From da868de8b5fc41f553b480c88dbf13a73a0a41ab Mon Sep 17 00:00:00 2001
From: Raymond Wiker <rayw@equinor.com>
Date: Fri, 17 Jan 2025 16:03:37 +0100
Subject: [PATCH 3/3] chore: cleanup/refactor, and only look for stale
 sharedkey files if the token directory exists.

---
 src/sumo/wrapper/_auth_provider.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/sumo/wrapper/_auth_provider.py b/src/sumo/wrapper/_auth_provider.py
index f695e97..725ad03 100644
--- a/src/sumo/wrapper/_auth_provider.py
+++ b/src/sumo/wrapper/_auth_provider.py
@@ -31,19 +31,17 @@ def _maybe_nfs_exception(exception):
 
 
 def get_token_dir():
-    return os.path.join(os.path.expanduser("~"), ".sumo")
+    return os.path.expanduser("~/.sumo")
 
 
 def get_token_path(resource_id, suffix, case_uuid=None):
     if case_uuid is not None:
         return os.path.join(
-            os.path.expanduser("~/.sumo"),
+            get_token_dir(),
             str(resource_id) + "+" + str(case_uuid) + suffix,
         )
     else:
-        return os.path.join(
-            os.path.expanduser("~/.sumo"), str(resource_id) + suffix
-        )
+        return os.path.join(get_token_dir(), str(resource_id) + suffix)
 
 
 class AuthProvider:
@@ -94,7 +92,9 @@ def store_shared_access_key_for_case(self, case_uuid, token):
         return
 
     def cleanup_shared_keys(self):
-        tokendir = os.path.join(os.path.expanduser("~/.sumo"))
+        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):