Merge pull request #66 from equinor/new-auth-class

new auth class for SumoClient

Author: adnejacobsen

requirements.txt

@@ -1,4 +1,4 @@
-msal~=1.5.1
+msal~=1.17.0
 requests~=2.24.0
 pytest~=6.1.1
 PyYAML>=5.4

src/sumo/wrapper/_new_auth.py

@@ -0,0 +1,106 @@
+import atexit
+import msal
+import os
+import sys
+import json
+from .config import AUTHORITY_HOST_URI
+
+HOME_DIR = os.path.expanduser("~")
+
+class NewAuth:
+    def __init__(
+        self,
+        client_id,
+        resource_id,
+        tenant_id,
+        interactive=False,
+        refresh_token=None
+    ):
+        self.interactive = interactive
+        self.scope = resource_id + "/.default"
+        self.refresh_token = refresh_token
+
+        self.token_path = os.path.join(
+            HOME_DIR, ".sumo", str(resource_id) + ".token"
+        )
+
+        self.cache = None
+
+        if not self.refresh_token:
+            self.cache = self.__load_cache()
+            atexit.register(self.__save_cache)
+            
+        self.msal = msal.PublicClientApplication(
+            client_id=client_id,
+            authority=f"{AUTHORITY_HOST_URI}/{tenant_id}", 
+            token_cache=self.cache
+        )
+
+
+    def get_token(self):
+        accounts = self.msal.get_accounts()
+        result = None
+
+        if accounts:
+            result = self.msal.acquire_token_silent([self.scope], account=accounts[0])
+
+        if not result:
+            if self.refresh_token:
+                result = self.msal.acquire_token_by_refresh_token(self.refresh_token, [self.scope])
+
+                if "error" in result:
+                    raise ValueError(
+                        "Failed to acquire token by refresh token. Err: %s" % json.dumps(result, indent=4)
+                    )
+            else:
+                if self.interactive:
+                    result = self.msal.acquire_token_interactive([self.scope])
+
+                    if "error" in result:
+                        raise ValueError(
+                            "Failed to acquire token interactively. Err: %s" % json.dumps(result, indent=4)
+                        )
+                else:
+                    flow = self.msal.initiate_device_flow([self.scope])
+
+                    if "error" in flow:
+                        raise ValueError(
+                            "Failed to create device flow. Err: %s" % json.dumps(flow, indent=4)
+                        )
+
+                    print(flow["message"])
+                    result = self.msal.acquire_token_by_device_flow(flow)
+
+                    if "error" in result:
+                        raise ValueError(
+                            "Failed to acquire token by device flow. Err: %s" % json.dumps(result, indent=4)
+                        )
+
+        return result["access_token"]
+
+
+    def __load_cache(self):
+        cache = msal.SerializableTokenCache()
+
+        if os.path.isfile(self.token_path):
+            with open(self.token_path, "r") as file:
+                cache.deserialize(file.read())
+
+        return cache
+
+    
+    def __save_cache(self):
+        if self.cache.has_state_changed:
+            old_mask = os.umask(0o077)
+
+            dir_path = os.path.dirname(self.token_path)
+            os.makedirs(dir_path, exist_ok=True)
+
+            with open(self.token_path, "w") as file:
+                file.write(self.cache.serialize())
+
+            if not sys.platform.lower().startswith("win"):
+                os.chmod(self.token_path, 0o600)
+                os.chmod(dir_path, 0o700)
+
+            os.umask(old_mask)

src/sumo/wrapper/_sumo_client.py

@@ -1,63 +1,65 @@
 import requests
-import logging
+import jwt
+import time
 
-from .config import APP_REGISTRATION, TENANT_ID, AUTHORITY_HOST_URI
-from ._auth import Auth
+from .config import APP_REGISTRATION, TENANT_ID
+from ._new_auth import NewAuth
 from ._request_error import AuthenticationError, TransientError, PermanentError
 
-logging.basicConfig(
-    format="%(asctime)s %(levelname)-8s %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S"
-)
-
-
 class SumoClient:
     def __init__(
         self,
         env,
-        access_token=None,
-        logging_level='INFO',
-        write_back=False
+        token=None,
+        interactive=False
     ):
-        self.env = env
-        self.user_provided_access_token = access_token
-
-        self.logger = logging.getLogger(__name__)
-        self.logger.setLevel(level=logging_level)
-
         if env not in APP_REGISTRATION:
             raise ValueError(f"Invalid environment: {env}")
 
-        self.client_id = APP_REGISTRATION[env]['CLIENT_ID']
-        self.resource_id = APP_REGISTRATION[env]['RESOURCE_ID']
-        self.authority_uri = AUTHORITY_HOST_URI + '/' + TENANT_ID
-
-        if not self.user_provided_access_token:
-            self.auth = Auth(
-                client_id=self.client_id,
-                resource_id=self.resource_id,
-                authority=self.authority_uri,
-                writeback=write_back
-            )
-
-            self.access_token = self.auth.get_token()
+        self.access_token = None
+        self.access_token_expires = None
+        self.refresh_token = None
+
+        if token:
+            payload = self.__decode_token(token)
+
+            if payload:
+                self.access_token = token
+                self.access_token_expires = payload["exp"]
+            else:
+                self.refresh_token = token
+
+        self.auth = NewAuth(
+            client_id=APP_REGISTRATION[env]['CLIENT_ID'],
+            resource_id=APP_REGISTRATION[env]['RESOURCE_ID'],
+            tenant_id=TENANT_ID,
+            interactive=interactive,
+            refresh_token=self.refresh_token
+        )
 
         if env == "localhost":
             self.base_url = f"http://localhost:8084/api/v1"
         else:
             self.base_url = f"https://main-sumo-{env}.radix.equinor.com/api/v1"
 
+
+    def __decode_token(self, token):
+        try:
+            payload = jwt.decode(token, options={"verify_signature": False})
+            return payload
+        except:
+            return None
+
+
     def _retrieve_token(self):
-        if self.user_provided_access_token:
-            self.logger.debug("User provided token exists, returning token")
-            return self.user_provided_access_token
-        else:
-            if self.auth.is_token_expired():
-                self.logger.debug("Token is expired, regenerating")
-                self.access_token = self.auth.get_token()
+        if self.access_token:
+            if self.access_token_expires <= int(time.time()):
+                raise ValueError("Access_token has expired")
+            else:
+                return self.access_token
+
+        return self.auth.get_token()
 
-        self.logger.debug("returning self.access_token from _retrieve_token")
-        return self.access_token
 
     def _process_params(self, params_dict):
         prefixed_params = {}
@@ -174,9 +176,6 @@ def _raise_request_error_exception(self, code, message):
         Raise the proper authentication error according to the code received from sumo.
         """
 
-        self.logger.debug("code: %s", code)
-        self.logger.debug("message: %s", message)
-
         if 503 <= code <= 504 or code == 404 or code == 500:
             raise TransientError(code, message)
         elif 401 <= code <= 403:

tests/test_sumo_thin_client.py

@@ -13,7 +13,7 @@
 
 class Connection:
     def __init__(self):
-        self.api = SumoClient(env="dev", logging_level="DEBUG")
+        self.api = SumoClient(env="dev")
 
 
 def _upload_parent_object(C, json):