hero-sdk 0.14.0__py3-none-any.whl

hero/__init__.py

@@ -0,0 +1,9 @@
+import os
+import json
+
+from . lib import errors
+from . lib.config import set_environment
+from . lib.event_trigger import event_trigger
+from . lib.loaders import get_env_variable, load_environment, load_runtime_config
+
+from .client import HeroClient

hero/client.py

@@ -0,0 +1,194 @@
+import jwt
+import base64
+from requests import Session
+from jwt import (
+    PyJWKClient,
+    ExpiredSignatureError,
+    InvalidTokenError,
+    InvalidSignatureError,
+    DecodeError,
+)
+
+from .url_map import URL_MAP
+from .lib import (
+    get_conf_from_collection,
+    get_env,
+    get_client_credentials,
+)
+from .services import (
+    AuthService,
+    DataRepoService,
+    TaskEngineService,
+    MLModelRegistry,
+    AssistantService,
+)
+from .services import SearchService
+
+from .lib.errors import (
+    TokenDecodeError,
+    TokenInvalidSignatureError,
+    TokenInvalidError,
+    TokenGeneralError,
+)
+
+
+class HeroClient:
+    """
+    Primary client for interfacing with the HERO API. Provides access to all services.
+    """
+
+    def __init__(self, client_id: str = None, client_secret: str = None):
+        """
+        Creates the Hero client.
+        """
+        self._scopes = []
+
+        self.env = get_env()
+        self.api = Session()
+        region = "us-west-2"
+
+        if client_id is None or client_secret is None:
+            client_id, client_secret = get_client_credentials()
+
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._cognito_api_url = get_conf_from_collection(
+            URL_MAP, "HERO_COGNITO_API_URL"
+        )
+        self._cognito_user_pool_id = get_conf_from_collection(URL_MAP, "USER_POOL_ID")
+        self._cognito_auth_url = f"https://cognito-idp.{region}.amazonaws.com/{region}_{self._cognito_user_pool_id}"
+        self._jwks_url = f"{self._cognito_auth_url}/.well-known/jwks.json"
+        self._jwk_client = PyJWKClient(self._jwks_url)
+        self._access_token = None
+
+    def _fetch_token(self):
+        """
+        Login to the Cognito user pool. Requires a client with a client secret and authorization to assign requested scopes.
+
+        Returns a JWT access token.
+        """
+        app_client_id_secret = f"{self._client_id}:{self._client_secret}".encode(
+            "utf-8"
+        )
+        # Request access_token following client credentials grant flow
+        basic_auth = f"Basic {base64.urlsafe_b64encode(app_client_id_secret).decode()}"
+
+        response = self.api.post(
+            self._cognito_api_url,
+            data=f'grant_type=client_credentials&scope={" ".join(self._scopes)}&client_id={self._client_id}',
+            headers={
+                "Authorization": basic_auth,
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            verify=False,
+        )
+
+        if response.status_code != 200:
+            raise Exception(
+                f"Failed to fetch access token: {response.status_code} - {response.text}, {self._scopes}"
+            )
+
+        self._access_token = response.json()["access_token"]
+
+    def _decode_token(self, token):
+        """
+        Decodes a JWT token and returns the payload, verifying its signature and expiration.
+
+        Raises TokenInvalidSignatureError if the signature is invalid.
+        Raises TokenDecodeError if the token is malformed.
+        Raises TokenInvalidError if the token is invalid.
+        Raises TokenGeneralError for any other unexpected errors.
+        """
+        try:
+            signing_key = self._jwk_client.get_signing_key_from_jwt(token).key
+            # Token is valid and not expired
+            return jwt.decode(
+                token,
+                key=signing_key,
+                algorithms=["RS256"],
+                options={"verify_exp": True, "verify_iat": False},
+            )
+        except ExpiredSignatureError:
+            # token expired, we need to refresh it
+            self._fetch_token()
+            # try to decode again
+            return jwt.decode(
+                token, algorithms=["RS256"], options={"verify_signature": False}
+            )
+        except InvalidSignatureError:
+            # The signature doesn't match — token could be tampered with
+            raise TokenInvalidSignatureError("Invalid token signature")
+        except DecodeError:
+            # Token is malformed (e.g., bad base64, wrong structure)
+            raise TokenDecodeError("Malformed token")
+        except InvalidTokenError as e:
+            # Other invalid token cases
+            raise TokenInvalidError(f"Invalid token: {str(e)}")
+        except Exception as e:
+            # Any other unexpected errors
+            raise TokenGeneralError("Unexpected error")
+
+    def add_scope(self, scope):
+        """
+        Adds a scope to the client.
+        """
+        if scope in self._scopes:
+            return
+        self._scopes.append(scope)
+
+    def authenticate(self):
+        """
+        Authenticates the client with the Cognito user pool.
+        """
+        self._fetch_token()
+
+    def get_token(self):
+        """
+        Returns the access token if it is valid.
+
+        Note: we could add re-auth capabilities here Or manage elsewhere
+        (i.e. resilient sessions, etc)
+        """
+        access_token_decoded = self._decode_token(self._access_token)
+        if access_token_decoded:
+            return self._access_token
+        else:
+            self._fetch_token()
+            return self._access_token
+
+    def Auth(self):
+        """
+        Returns a AuthService instance.
+        """
+        return AuthService(self)
+
+    def DataRepo(self, application_id=None):
+        """
+        Returns a DataRepoService instance.
+        """
+        return DataRepoService(self, application_id)
+
+    def TaskEngine(self, application_id=None):
+        """
+        Returns a TaskEngineService instance.
+        """
+        return TaskEngineService(self, application_id)
+
+    def MLModelRegistry(self, application_id=None):
+        """
+        Returns a MLModelRegistry instance.
+        """
+        self.authInstance = self.Auth()
+        return MLModelRegistry(self, application_id)
+
+    def Search(self):
+        """
+        Returns a SearchService instance.
+        """
+        return SearchService(self)
+
+    def Assistant(self, application_id):
+        """
+        Returns a AssistantService instance.
+        """
+        return AssistantService(self, application_id)

hero/lib/__init__.py

@@ -0,0 +1,19 @@
+from .service_base import ServiceBase
+from .config import (
+    get_resilient_session,
+    get_client_credentials,
+    get_service_id,
+    get_conf_from_collection,
+    get_env,
+    set_environment,
+)
+from .decorators import (
+    decorate_all,
+    log_errors,
+    track_calls,
+    retry_method,
+)
+from .errors import HeroRetryError
+from .helpers import set_log_level_from_env
+from .attr_mapper import HeroObject
+from .credentials import HeroCredentialsProvider

hero/lib/attr_mapper.py

@@ -0,0 +1,82 @@
+from collections.abc import MutableMapping
+
+
+def _to_attr(value):
+    if isinstance(value, dict):
+        return HeroObject(value)
+    if isinstance(value, list):
+        return [_to_attr(v) for v in value]
+    if isinstance(value, tuple):
+        return tuple(_to_attr(v) for v in value)
+    return value
+
+
+class HeroObject(MutableMapping):
+    """
+    Mapping that supports attribute-style access and recursive wrapping.
+    Example:
+        run = HeroObject({"info": {"run_id": "123"}, "data": {"tags": [{"key": "k","value": "v"}]}})
+        run.info.run_id -> "123"
+        run["info"]["run_id"] -> "123"
+    """
+
+    __slots__ = ("_store",)
+
+    def __init__(self, data=None, **kwargs):
+        object.__setattr__(self, "_store", {})
+        if data:
+            for k, v in data.items():
+                self._store[k] = _to_attr(v)
+        for k, v in kwargs.items():
+            self._store[k] = _to_attr(v)
+
+    # Mapping protocol
+    def __getitem__(self, key):
+        return self._store[key]
+
+    def __setitem__(self, key, value):
+        self._store[key] = _to_attr(value)
+
+    def __delitem__(self, key):
+        del self._store[key]
+
+    def __iter__(self):
+        return iter(self._store)
+
+    def __len__(self):
+        return len(self._store)
+
+    # Attribute access
+    def __getattr__(self, name):
+        try:
+            return self._store[name]
+        except KeyError as e:
+            raise AttributeError(name) from e
+
+    def __setattr__(self, name, value):
+        # keep internal slot private; everything else goes in the mapping
+        if name == "_store":
+            object.__setattr__(self, name, value)
+        else:
+            self._store[name] = _to_attr(value)
+
+    def __delattr__(self, name):
+        try:
+            del self._store[name]
+        except KeyError as e:
+            raise AttributeError(name) from e
+
+    def to_dict(self):
+        def _to_plain(v):
+            if isinstance(v, HeroObject):
+                return {k: _to_plain(v[k]) for k in v}
+            if isinstance(v, list):
+                return [_to_plain(i) for i in v]
+            if isinstance(v, tuple):
+                return tuple(_to_plain(i) for i in v)
+            return v
+
+        return {k: _to_plain(v) for k, v in self._store.items()}
+
+    def __repr__(self):
+        return f"HeroObject({self.to_dict()!r})"

hero/lib/config.py

@@ -0,0 +1,71 @@
+import os
+import pathlib
+import json
+import logging
+import pathlib
+
+log = logging.getLogger("hero:config")
+
+
+def get_env():
+    return os.environ.get("HERO_ENV", "dev")
+
+
+def get_service_id(key):
+    env = get_env()
+    return f"{env}-{os.environ[key]}"
+
+
+def get_conf_from_collection(collection, key):
+    return os.environ.get(key, collection[get_env()][key])
+
+
+def get_resilient_session():
+    return os.environ.get("HERO_RESILIENT_SESSION", "False").lower() in ("true")
+
+
+def get_client_credentials():
+    """Returns the client credentials tuple (client_id, client_secret) from the environment variables HERO_CLIENT_ID and HERO_CLIENT_SECRET"""
+    client_credentials = (
+        os.environ["HERO_CLIENT_ID"],
+        os.environ["HERO_CLIENT_SECRET"],
+    )
+    return client_credentials
+
+
+def set_environment(application_id, path=None):
+    """
+    This function will set the HERO environment variables from a file stored in ~/.hero/.credentials.json
+
+    Ensure you have your hero credentials saved in `~/.hero/credentials.json` format.
+
+        {
+            "dev-aeroportal": {
+                "HERO_CLIENT_ID": "1c5ngb6o6lvtdfkus0sflstdq4",
+                "HERO_CLIENT_SECRET": "******",
+                "[ANOTHER_KEY]": "***"
+            }
+        }
+
+
+    Parameters
+    ----------
+    application_id: name in the credentials.json file (e.g. dev-aeroportal)
+
+    """
+    if path is None:
+        path = str(pathlib.Path(os.environ.get("HOME")) / ".hero" / "credentials.json")
+
+    try:
+        credentials = json.loads(open(path, "r").read())
+        these_credentials = credentials[application_id]
+        for key, value in these_credentials.items():
+            os.environ[key] = value
+
+    except FileNotFoundError as e:
+        log.error(f"Unable to load {path}")
+    except KeyError as e:
+        log.error(f"Unable to read key {e}")
+    except Exception as e:
+        log.error(str(e))
+        log.error(f"Unable to set credentials from {path}")

hero/lib/credentials.py

@@ -0,0 +1,88 @@
+from ..client import HeroClient
+from typing import Optional, Dict
+import os
+
+class HeroCredentialsProvider():
+    """
+    Provide AWS credentials through Hero's Token Vending Machine.
+    """
+    def __init__(self):
+        self._hero_client = None
+        self._auth = None
+        self._initialize_hero_client()
+    
+    def _initialize_hero_client(self):
+        try:
+            self._hero_client = HeroClient()
+            self._auth = self._hero_client.Auth()
+            self._hero_client.authenticate()
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize Hero Client: {e}")
+    
+    def get_aws_credentials(
+            self,
+            app_id: str,
+            app_type: str,
+            action: str,
+            resource_id: Optional[str] = None,
+            resource_type: Optional[str] = None,
+            aws_region: str = "us-west-2"
+    ):
+        """
+        Retrieve temporary AWS credentials from Hero's TVM.
+
+        Parameters
+        ----------
+        app_id : str, required
+            The ID of the application
+        app_type: str, required
+            The type of application
+        action : str, required
+            The action to perform (e.g., 'readFile', 'executeQuery')
+        resource_id : str, optional
+            ID of the specific resource for fine-grained access control
+        resource_type : str, optional
+            Type of resource for fine-grained access control
+        aws_region : str, optional
+            AWS region for credentials ( default: 'us-west-2')
+        
+        Returns
+        -------
+        dict
+            Dictionary containing access_key_id, secret_access_key, 
+            session_token, region, and expiration
+        """
+        try:
+            hero_response = self._auth.get_tvm_session(
+                app_id=app_id,
+                app_type=app_type,
+                resource_id=resource_id,
+                resource_type=resource_type,
+                action=action
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to retrieve credentials from Hero TVM: {e}")
+        
+        credentials = {
+            'access_key_id': hero_response.get('AccessKeyId'),
+            'secret_access_key': hero_response.get('SecretAccessKey'),
+            'session_token': hero_response.get('SessionToken'),
+            'region': aws_region,
+            'expiration': hero_response.get('Expiration'),
+        }
+
+        return credentials
+    
+    def inject_into_env(self, credentials: Dict[str, str]):
+        """
+        Inject AWS Credentials into the environment variables
+
+        Parameters
+        ----------
+        credentials : dict
+            Credentials dictionary from get_aws_credentials()
+        """
+        os.environ['AWS_ACCESS_KEY_ID'] = credentials['access_key_id']
+        os.environ['AWS_SECRET_ACCESS_KEY'] = credentials['secret_access_key']
+        os.environ['AWS_SESSION_TOKEN'] = credentials['session_token']
+        os.environ['AWS_REGION'] = credentials['region']

hero/lib/decorators.py

@@ -0,0 +1,88 @@
+import os
+import logging
+from functools import wraps
+from tenacity import stop_after_attempt, wait_fixed, wait_exponential
+
+from .errors import HeroRetryError
+
+log = logging.getLogger("hero:service")
+
+
+def decorate_all(decorator):
+
+    def decorate(cls):
+        for name, attr in cls.__dict__.items():
+            if name.startswith("__"):
+                continue
+            if isinstance(attr, staticmethod):
+                setattr(cls, name, staticmethod(decorator(attr.__func__)))
+            elif isinstance(attr, classmethod):
+                setattr(cls, name, classmethod(decorator(attr.__func__)))
+            elif callable(attr):
+                setattr(cls, name, decorator(attr))
+        return cls
+
+    return decorate
+
+
+def log_errors(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception:
+            if os.getenv("HERO_LOG_ALL_ERRORS") == "True":
+                log.error(f"Hero Service Error in {func.__name__}", exc_info=True)
+            raise
+
+    return wrapper
+
+
+def track_calls(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        self._calls += 1
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
+def retry_method(func, errFunc):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        # attempts order: kwargs -> EVN -> default
+
+        attempts = int(
+            kwargs.get(
+                "attempts",
+                os.environ.get("HERO_RETRY_ATTEMPTS", self.default_attempts),
+            )
+        )
+
+        wait_schedule = str(
+            kwargs.get(
+                "wait",
+                os.environ.get("HERO_RETRY_WAIT", self.default_wait),
+            )
+        )
+
+        # print(str(func.__name__), wait_schedule)
+        wait = wait_fixed(1)
+        if wait_schedule == "exp":
+            wait = wait_exponential(multiplier=1, min=1, max=60)
+
+        try:
+            local_instance = errFunc.retry_with(
+                stop=stop_after_attempt(attempts), wait=wait
+            )
+            results = local_instance.__call__(self, func, *args, **kwargs)
+            return results
+
+        except Exception as e:
+            raise HeroRetryError(
+                str(e),
+                local_instance.retry.statistics.get("attempt_number"),
+                local_instance.retry.statistics.get("idle_for"),
+            )
+
+    return wrapper