feat: delete sharedkey if invalid and attempt other auth method (#253)

* feat: delete sharedkey if invalid and attempt other auth method

* feat/delete shared key if not valid

* refactor/optimize imports

* feat: remove verification of instance in the decorator and use delete_token response instead

Author: atadj

src/sumo/wrapper/_auth_provider.py

@@ -99,6 +99,9 @@ def has_case_token(self, case_uuid):
             get_token_path(self._resource_id, ".sharedkey", case_uuid)
         )
 
+    def delete_token(self):
+        return False
+
     pass
 
 
@@ -400,9 +403,10 @@ class AuthProviderSumoToken(AuthProvider):
     def __init__(self, resource_id, case_uuid=None):
         super().__init__(resource_id)
         protect_token_cache(resource_id, ".sharedkey", case_uuid)
-        token_path = get_token_path(resource_id, ".sharedkey", case_uuid)
-        with open(token_path, "r") as f:
+        self.token_path = get_token_path(resource_id, ".sharedkey", case_uuid)
+        with open(self.token_path, "r") as f:
             self._token = f.readline().strip()
+
         return
 
     def get_token(self):
@@ -411,6 +415,11 @@ def get_token(self):
     def get_authorization(self):
         return {"X-SUMO-Token": self._token}
 
+    def delete_token(self):
+        if os.path.exists(self.token_path):
+            os.unlink(self.token_path)
+        return True
+
 
 @tn.retry(
     retry=tn.retry_if_exception(_maybe_nfs_exception),

src/sumo/wrapper/_decorators.py

@@ -4,11 +4,13 @@
 
 def raise_for_status(func):
     @wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(self, *args, **kwargs):
         # FIXME: in newer versions of httpx, raise_for_status() is chainable,
         # so we could simply write
         # return func(*args, **kwargs).raise_for_status()
-        response = func(*args, **kwargs)
+        response = func(self, *args, **kwargs)
+        if response.status_code == 401:
+            self._handle_invalid_shared_key()
         response.raise_for_status()
         return response
 
@@ -17,11 +19,13 @@ def wrapper(*args, **kwargs):
 
 def raise_for_status_async(func):
     @wraps(func)
-    async def wrapper(*args, **kwargs):
+    async def wrapper(self, *args, **kwargs):
         # FIXME: in newer versions of httpx, raise_for_status() is chainable,
         # so we could simply write
         # return func(*args, **kwargs).raise_for_status()
-        response = await func(*args, **kwargs)
+        response = await func(self, *args, **kwargs)
+        if response.status_code == 401:
+            self._handle_invalid_shared_key()
         response.raise_for_status()
         return response
 

src/sumo/wrapper/sumo_client.py

@@ -208,6 +208,13 @@ def blob_client(self) -> BlobClient:
             self._retry_strategy,
         )
 
+    def _handle_invalid_shared_key(self):
+        """Handle the invalid shared key by deleting it."""
+        if self.auth.delete_token():
+            print(
+                "Invalid shared key detected and deleted, run again to reset automatically"
+            )
+
     @raise_for_status
     def get(self, path: str, params: Optional[Dict] = None) -> httpx.Response:
         """Performs a GET-request to the Sumo API.