pipekit-sdk 2.0.1__tar.gz → 2.1.0__tar.gz

{pipekit_sdk-2.0.1 → pipekit_sdk-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pipekit-sdk
-Version: 2.0.1
+Version: 2.1.0
 Summary: Pipekit Python SDK
 Home-page: https://pipekit.io
 License: BSD-3-Clause
@@ -16,8 +16,10 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: hera (>=5.5.2)
 Requires-Dist: pydantic (<3)
+Requires-Dist: pyjwt (>=2.9.0,<3.0.0)
 Requires-Dist: responses (>=0.25.3,<0.26.0)
 Project-URL: Documentation, https://docs.pipekit.io/
 Description-Content-Type: text/markdown

{pipekit_sdk-2.0.1 → pipekit_sdk-2.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pipekit-sdk"
-version = "2.0.1"
+version = "2.1.0"
 description = "Pipekit Python SDK"
 authors = ["Pipekit <ci@pipekit.io>"]
 license = "BSD-3-Clause"
@@ -27,6 +27,7 @@ python = ">=3.8,<4"
 hera = ">=5.5.2"
 pydantic = "<3"
 responses = "^0.25.3"
+pyjwt = "^2.9.0"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "*"

{pipekit_sdk-2.0.1 → pipekit_sdk-2.1.0}/src/pipekit_sdk/service.py

@@ -2,19 +2,71 @@
 
 import json
 import os
-from typing import Any, Generator, Optional, cast
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any, Generator, List, Optional, cast
 
+import jwt
 import requests
 from hera.workflows import CronWorkflow, Workflow
 from hera.workflows.models import WorkflowCreateRequest
+from pydantic import BaseModel, ValidationError
 
 from pipekit_sdk._helpers import is_valid_uuid
 from pipekit_sdk.models.model import Cluster, Logs, Pipe, PipeRun
 
 
+class _LoginResponse(BaseModel):
+    """LoginResponse is a model for Login responses."""
+
+    token_type: str
+    access_token: str
+    expires_in: int
+    refresh_token: str
+
+
+class _AccessToken(BaseModel):
+    """AccessToken is a model for decoded access tokens."""
+
+    identityUUID: str
+    ClusterUUID: str
+    OrgUUID: str
+    isActive: bool
+    products: List[str]
+    iss: str
+    sub: str
+    exp: int
+
+
+class _ValidationResult(Enum):
+    OK = 0
+    INVALID_TOKEN = 1
+    EXPIRED_TOKEN = 2
+
+
 class PipekitService:
     """PipekitService is a wrapper around the Pipekit API. It provides a simple interface to interact with Pipekit."""
 
+    @staticmethod
+    def __from_home_dir() -> Optional[_LoginResponse]:
+        home = Path.home()
+        pipekit_token_file_path = os.path.join(home, ".pipekit", "token")
+        try:
+            with open(pipekit_token_file_path, "r") as f:
+                raw_data = f.read()
+                return _LoginResponse.model_validate_json(raw_data)
+        except ValidationError:
+            return None
+
+    @staticmethod
+    def __write_home_dir(login_resp: _LoginResponse):
+        home = Path.home()
+        pipekit_token_file_path = os.path.join(home, ".pipekit", "token")
+        with open(pipekit_token_file_path, "w") as f:
+            payload = login_resp.model_dump_json(indent=2)
+            f.write(payload)
+
     def __init__(
         self,
         username: str = "",
@@ -38,8 +90,60 @@ class PipekitService:
             self.access_token = token
             return
 
+        maybe_resp = self.__from_home_dir()
+        self.token_data = maybe_resp
+
+        if self.token_data is not None:
+            self.__token_login_flow()
+            return
+
         self.__login()
 
+    def __update_token(self) -> None:
+        login_url = f"{self.id_url}/api/id/v1/oauth/token"
+        assert self.token_data is not None, "Token data was None, this implies a bug"
+        login_request = {"grant_type": "refresh_token", "refresh_token": self.token_data.refresh_token}
+        login_response = requests.post(
+            login_url,
+            data=json.dumps(login_request),
+            headers={"Content-Type": "application/json"},
+            timeout=10,
+        )
+        if login_response.status_code // 100 != 2:
+            raise Exception(
+                "Couldn't refresh access token, obtained non 2xx status code, try logging in again via the cli"
+            )
+
+        login_response_model = _LoginResponse.model_validate(login_response.json())
+        self.__write_home_dir(login_response_model)
+        self.access_token = login_response_model.access_token
+
+    def __validate_token(self) -> _ValidationResult:
+        assert self.token_data is not None
+        try:
+            decoded_payload = jwt.decode(self.token_data.access_token, options={"verify_signature": False})
+            acc_tk = _AccessToken.model_validate(decoded_payload)
+            exp_time = datetime.fromtimestamp(acc_tk.exp)
+            now_time = datetime.now()
+
+            if exp_time <= now_time:
+                return _ValidationResult.EXPIRED_TOKEN
+        except jwt.PyJWTError:
+            return _ValidationResult.INVALID_TOKEN
+        return _ValidationResult.OK
+
+    def __token_login_flow(self) -> None:
+        assert self.token_data is not None, "token was None when it was expected to be valid"
+        res = self.__validate_token()
+        if res == _ValidationResult.EXPIRED_TOKEN:
+            self.__update_token()
+        elif res == _ValidationResult.INVALID_TOKEN:
+            raise RuntimeError(
+                "Token was invalid, if the `.pipekit/token` file is corrupted, please delete the file and login again"
+            )
+        else:
+            self.access_token = self.token_data.access_token
+
     def __login(self) -> None:
         login_request = {
             "username": self.username,
@@ -486,3 +590,7 @@ class PipekitService:
 
         for log in self.get_logs(run_uuid, pod_name=pod_name, container_name=container_name):
             self.__print_log_message(log)
+
+
+if __name__ == "__main__":
+    p = PipekitService()