diff --git a/src/sumo/wrapper/_auth_provider.py b/src/sumo/wrapper/_auth_provider.py index e28aa19..1f83f1f 100644 --- a/src/sumo/wrapper/_auth_provider.py +++ b/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), diff --git a/src/sumo/wrapper/_decorators.py b/src/sumo/wrapper/_decorators.py index d53d700..628fad1 100644 --- a/src/sumo/wrapper/_decorators.py +++ b/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 diff --git a/src/sumo/wrapper/sumo_client.py b/src/sumo/wrapper/sumo_client.py index 2a85339..c8780a7 100644 --- a/src/sumo/wrapper/sumo_client.py +++ b/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.