seer-pas-sdk 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

seer_pas_sdk/auth/auth.py

@@ -1,5 +1,9 @@
+import atexit
+from datetime import datetime
 import requests
 import jwt
+from ..common import get_version
+from ..common.errors import ServerError
 
 
 class Auth:
@@ -24,6 +28,7 @@ class Auth:
 
         self.username = username
         self.__password = password
+        self.version = get_version()
 
         if instance not in Auth._instances:
             if instance.startswith("https://") or instance.startswith(
@@ -35,7 +40,10 @@ class Auth:
                 raise ValueError("Invalid PAS instance.")
         else:
             self.url = Auth._instances[instance]
-
+        self.refresh_token = None
+        self.id_token = None
+        self.access_token = None
+        self.refresh_token_expiry = 0
         self.instance = instance
 
         # Null initialize multi tenant attributes
@@ -46,7 +54,9 @@ class Auth:
             self.active_role,
         ) = [None] * 4
 
-    def login(self):
+        atexit.register(self._logout)
+
+    def _login(self):
         """
         Logs into the PAS instance using the mapped URL and the login credentials (username and password) provided in the constructor.
 
@@ -55,7 +65,11 @@ class Auth:
         dict
             A dictionary containing the login response from the PAS instance.
         """
-        response = requests.post(
+        s = requests.Session()
+        s.headers.update(
+            {"x-seer-source": "sdk", "x-seer-id": f"{self.version}/login"}
+        )
+        response = s.post(
             f"{self.url}auth/login",
             json={"username": self.username, "password": self.__password},
         )
@@ -65,35 +79,180 @@ class Auth:
                 "Check if the credentials are correct or if the backend is running or not."
             )
 
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            raise ServerError("Could not login to the PAS instance")
+
+        res = response.json()
+        mfa_challenge = res.get("challenge", None)
+        if mfa_challenge:
+            mfa_args = {}
+            if any(
+                x not in res
+                for x in ["challenge", "session", "challengeParameters"]
+            ):
+                raise ServerError(
+                    "Unexpected return from server during MFA challenge. Please check the PAS SDK version and update the PAS SDK if necessary."
+                )
+            mfa_args["challengeName"] = res["challenge"]
+            mfa_args["session"] = res["session"]
+            mfa_args["username"] = res["challengeParameters"].get(
+                "USER_ID_FOR_SRP", ""
+            )
+            if not mfa_args["username"]:
+                raise ServerError(
+                    "Unexpected return from server during MFA challenge. Please check the PAS SDK version and update the PAS SDK if necessary."
+                )
+            print(
+                "Multi-factor authentication (MFA) is enabled for your account."
+            )
+            mfa_code = input(
+                "Please enter the code from your authenticator app:"
+            )
+            mfa_response = s.post(
+                f"{self.url}auth/confirmMFA",
+                json={
+                    "username": mfa_args["username"],
+                    "mfaCode": mfa_code,
+                    "challengeName": mfa_args["challengeName"],
+                    "session": mfa_args["session"],
+                },
+            )
+            if mfa_response.status_code != 200:
+                raise ServerError(
+                    "Could not confirm MFA for the PAS instance."
+                )
+            mfa_response = mfa_response.json()
+            if "AuthenticationResult" not in mfa_response:
+                raise ServerError(
+                    "Unexpected return from server during MFA confirmation. Please check the PAS SDK version and update the PAS SDK if necessary."
+                )
+            mfa_response = mfa_response["AuthenticationResult"]
+            if any(
+                x not in mfa_response
+                for x in [
+                    "AccessToken",
+                    "IdToken",
+                    "ExpiresIn",
+                    "RefreshToken",
+                ]
+            ):
+                raise ServerError(
+                    "Unexpected return from server during MFA confirmation. Please check the PAS SDK version and update the PAS SDK if necessary."
+                )
+            res["id_token"] = mfa_response["IdToken"]
+            res["access_token"] = mfa_response["AccessToken"]
+            res["expiresIn"] = mfa_response["ExpiresIn"]
+            res["refresh_token"] = mfa_response["RefreshToken"]
+
+        decoded_token = jwt.decode(
+            res["id_token"], options={"verify_signature": False}
+        )
+        self.base_tenant_id = decoded_token["custom:tenantId"]
+        self.base_role = decoded_token["custom:role"]
+
+        if not self.active_tenant_id:
+            self.active_tenant_id = self.base_tenant_id
+
+        if not self.active_role:
+            self.active_role = self.base_role
+
+        self.id_token = res["id_token"]
+        self.access_token = res["access_token"]
+        self.refresh_token = res.get("refresh_token", None)
+        self.token_expiry = int(datetime.now().timestamp()) + res.get(
+            "expiresIn", 0
+        )
+        return res["id_token"], res["access_token"]
+
+    def _logout(self):
+        if not self.has_valid_token():
+            return True
+        s = requests.Session()
+        s.headers.update(
+            {"x-seer-source": "sdk", "x-seer-id": f"{self.version}/logout"}
+        )
+        response = s.post(
+            f"{self.url}auth/logout",
+            json={
+                "username": self.username,
+                "refreshtoken": self.refresh_token,
+            },
+        )
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            raise ServerError("Could not logout from the PAS instance")
+        self.refresh_token = None
+        self.token_expiry = 0
+        print(
+            f"User {self.username} logged out successfully. Thank you for using the PAS SDK."
+        )
+        return True
+
+    def _refresh_token(self):
+        """Refreshes the token using the refresh token.
+
+        Raises:
+            ServerError: If the token could not be refreshed.
+
+        Returns:
+            dict: The response from the server containing the new token.
+        """
+        s = requests.Session()
+        s.headers.update(
+            {
+                "x-seer-source": "sdk",
+                "x-seer-id": f"{self.version}/refresh_token",
+            }
+        )
+        response = s.post(
+            f"{self.url}auth/refreshtoken",
+            json={"refreshtoken": self.refresh_token},
+        )
         if response.status_code == 200:
             return response.json()
+        else:
+            raise ServerError("Could not refresh token")
 
     def get_token(self):
         """
-        Gets the token from the login response.
+        Gets the current token. If the token is expired, it refreshes the token using the refresh token.
 
         Returns
         -------
         str
             The token from the login response.
         """
-
-        res = self.login()
+        if self.has_valid_token():
+            return self.id_token, self.access_token
+        else:
+            res = self._refresh_token()
 
         if "id_token" not in res or "access_token" not in res:
             raise ValueError(
                 "Check if the credentials are correct or if the backend is running or not."
             )
-        decoded_token = jwt.decode(
-            res["id_token"], options={"verify_signature": False}
+        self.token_expiry = int(datetime.now().timestamp()) + res.get(
+            "expiresIn", 0
         )
-        self.base_tenant_id = decoded_token["custom:tenantId"]
-        self.base_role = decoded_token["custom:role"]
 
-        if not self.active_tenant_id:
-            self.active_tenant_id = self.base_tenant_id
+        self.id_token = res["id_token"]
+        self.access_token = res["access_token"]
+        return res["id_token"], res["access_token"]
 
-        if not self.active_role:
-            self.active_role = self.base_role
+    def has_valid_token(self):
+        """
+        Check if the id and access tokens are valid.
 
-        return res["id_token"], res["access_token"]
+        Returns
+        -------
+        bool
+            True if the refresh token is valid, False otherwise.
+        """
+        return (
+            self.access_token is not None
+            and self.id_token is not None
+            and self.token_expiry > int(datetime.now().timestamp())
+        )

seer_pas_sdk/common/__init__.py

@@ -13,7 +13,8 @@ import json
 import zipfile
 import tempfile
 
-from ..auth.auth import Auth
+import subprocess
+from importlib.metadata import version, PackageNotFoundError
 
 from .groupanalysis import *
 
@@ -98,7 +99,7 @@ def dict_to_df(data):
 
 
 # Most cases appear to be a .tsv file.
-def url_to_df(url, is_tsv=True):
+def download_df(url, is_tsv=True, dtype={}):
     """
     Fetches a TSV/CSV file from a URL and returns as a Pandas DataFrame.
 
@@ -110,6 +111,9 @@ def url_to_df(url, is_tsv=True):
     is_tsv : bool
         True if the file is a TSV file, False if it is a CSV file.
 
+    dtype : dict
+        Data type conversion when intaking columns. e.g. {'a': str, 'b': np.float64}
+
     Returns
     -------
     pandas.core.frame.DataFrame
@@ -122,7 +126,7 @@ def url_to_df(url, is_tsv=True):
 
     Examples
     --------
-    >>> csv = url_to_df("link_to_csv_file")
+    >>> csv = download_df("link_to_csv_file")
     >>> print(csv)
     >>>     Sample ID  Sample name  Well location  MS file name
         0           1  SampleName1              1  SDKTest1.raw
@@ -137,9 +141,9 @@ def url_to_df(url, is_tsv=True):
         return pd.DataFrame()
     url_content = io.StringIO(requests.get(url).content.decode("utf-8"))
     if is_tsv:
-        csv = pd.read_csv(url_content, sep="\t")
+        csv = pd.read_csv(url_content, sep="\t", dtype=dtype)
     else:
-        csv = pd.read_csv(url_content)
+        csv = pd.read_csv(url_content, dtype=dtype)
     return csv
 
 
@@ -720,3 +724,40 @@ def rename_d_zip_file(source, destination):
                     zip_ref.write(file_path, arcname)
 
     print(f"Renamed {source} to {destination}")
+
+
+def get_version():
+    """
+    Returns the version of the seer-pas-sdk package.
+    """
+    v = ""
+    try:
+        v = version("seer-pas-sdk")
+        v = "v" + v
+        if "dev" in v:
+            v = v[: v.index("dev") - 1].strip()
+            v += "-dev"
+    except Exception as e:
+        pass
+
+    if not v:
+        try:
+            v = subprocess.check_output(
+                ["git", "describe", "--tags"], text=True
+            ).strip()
+            if "-" in v:
+                v = v.split("-")[0]
+            v += "-dev"
+        except subprocess.CalledProcessError:
+            v = "unknown"
+    return f"{v}"
+
+
+def title_case_to_snake_case(s: str):
+    """
+    Converts a title case string to snake case.
+
+    """
+    s = s.replace(" ", "_").replace(".", "_").casefold()
+
+    return s