otf-api 0.2.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

otf_api/auth.py

@@ -0,0 +1,314 @@
+import typing
+from typing import Any
+
+import jwt
+import pendulum
+from loguru import logger
+from pycognito import AWSSRP, Cognito, MFAChallengeException
+from pycognito.exceptions import TokenVerificationException
+from pydantic import Field
+from pydantic.config import ConfigDict
+
+from otf_api.models.base import OtfItemBase
+
+if typing.TYPE_CHECKING:
+    from boto3.session import Session
+    from botocore.config import Config
+
+CLIENT_ID = "65knvqta6p37efc2l3eh26pl5o"  # from otlive
+USER_POOL_ID = "us-east-1_dYDxUeyL1"
+
+
+class OtfCognito(Cognito):
+    _device_key: str | None = None
+
+    def __init__(
+        self,
+        user_pool_id: str,
+        client_id: str,
+        user_pool_region: str | None = None,
+        username: str | None = None,
+        id_token: str | None = None,
+        refresh_token: str | None = None,
+        access_token: str | None = None,
+        client_secret: str | None = None,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+        session: "Session|None" = None,
+        botocore_config: "Config|None" = None,
+        boto3_client_kwargs: dict[str, Any] | None = None,
+        device_key: str | None = None,
+    ):
+        super().__init__(
+            user_pool_id,
+            client_id,
+            user_pool_region=user_pool_region,
+            username=username,
+            id_token=id_token,
+            refresh_token=refresh_token,
+            access_token=access_token,
+            client_secret=client_secret,
+            access_key=access_key,
+            secret_key=secret_key,
+            session=session,
+            botocore_config=botocore_config,
+            boto3_client_kwargs=boto3_client_kwargs,
+        )
+        self.device_key = device_key
+
+    @property
+    def device_key(self) -> str | None:
+        return self._device_key
+
+    @device_key.setter
+    def device_key(self, value: str | None):
+        if not value:
+            logger.info("Clearing device key")
+            self._device_key = value
+            return
+
+        redacted_value = value[:4] + "*" * (len(value) - 8) + value[-4:]
+        logger.info(f"Setting device key: {redacted_value}")
+        self._device_key = value
+
+    def _set_tokens(self, tokens: dict[str, Any]):
+        """Set the tokens and device metadata from the response.
+
+        Args:
+            tokens (dict): The response from the Cognito service.
+        """
+        super()._set_tokens(tokens)
+
+        if new_metadata := tokens["AuthenticationResult"].get("NewDeviceMetadata"):
+            self.device_key = new_metadata["DeviceKey"]
+
+    def authenticate(self, password: str, client_metadata: dict[str, Any] | None = None, device_key: str | None = None):
+        """
+        Authenticate the user using the SRP protocol. Overridden to add `confirm_device` call.
+
+        Args:
+            password (str): The user's password
+            client_metadata (dict, optional): Any additional client metadata to send to Cognito
+        """
+        aws = AWSSRP(
+            username=self.username,
+            password=password,
+            pool_id=self.user_pool_id,
+            client_id=self.client_id,
+            client=self.client,
+            client_secret=self.client_secret,
+        )
+        try:
+            tokens = aws.authenticate_user(client_metadata=client_metadata)
+        except MFAChallengeException as mfa_challenge:
+            self.mfa_tokens = mfa_challenge.get_tokens()
+            raise mfa_challenge
+
+        # Set the tokens and device metadata
+        self._set_tokens(tokens)
+
+        if not device_key:
+            # Confirm the device so we can use the refresh token
+            aws.confirm_device(tokens)
+        else:
+            self.device_key = device_key
+            try:
+                self.renew_access_token()
+            except TokenVerificationException:
+                logger.error("Failed to renew access token. Confirming device.")
+                self.device_key = None
+                aws.confirm_device(tokens)
+
+    def check_token(self, renew: bool = True) -> bool:
+        """
+        Checks the exp attribute of the access_token and either refreshes
+        the tokens by calling the renew_access_tokens method or does nothing
+        :param renew: bool indicating whether to refresh on expiration
+        :return: bool indicating whether access_token has expired
+        """
+        if not self.access_token:
+            raise AttributeError("Access Token Required to Check Token")
+        now = pendulum.now()
+        dec_access_token = jwt.decode(self.access_token, options={"verify_signature": False})
+
+        exp = pendulum.DateTime.fromtimestamp(dec_access_token["exp"])
+        if now > exp.subtract(minutes=15):
+            expired = True
+            if renew:
+                self.renew_access_token()
+        else:
+            expired = False
+        return expired
+
+    def renew_access_token(self):
+        """Sets a new access token on the User using the cached refresh token and device metadata."""
+        auth_params = {"REFRESH_TOKEN": self.refresh_token}
+        self._add_secret_hash(auth_params, "SECRET_HASH")
+
+        if self.device_key:
+            logger.info("Using device key for refresh token")
+            auth_params["DEVICE_KEY"] = self.device_key
+
+        refresh_response = self.client.initiate_auth(
+            ClientId=self.client_id, AuthFlow="REFRESH_TOKEN_AUTH", AuthParameters=auth_params
+        )
+        self._set_tokens(refresh_response)
+
+    @classmethod
+    def from_token(
+        cls, access_token: str, id_token: str, refresh_token: str | None = None, device_key: str | None = None
+    ) -> "OtfCognito":
+        """Create an OtfCognito instance from an id token.
+
+        Args:
+            access_token (str): The access token.
+            id_token (str): The id token.
+            refresh_token (str, optional): The refresh token. Defaults to None.
+            device_key (str, optional): The device key. Defaults
+
+        Returns:
+            OtfCognito: The user instance
+        """
+        cognito = OtfCognito(
+            USER_POOL_ID,
+            CLIENT_ID,
+            access_token=access_token,
+            id_token=id_token,
+            refresh_token=refresh_token,
+            device_key=device_key,
+        )
+        cognito.verify_tokens()
+        cognito.check_token()
+        return cognito
+
+    @classmethod
+    def login(cls, username: str, password: str) -> "OtfCognito":
+        """Create an OtfCognito instance from a username and password.
+
+        Args:
+            username (str): The username to login with.
+            password (str): The password to login with.
+
+        Returns:
+            OtfCognito: The logged in user.
+        """
+        cognito_user = OtfCognito(USER_POOL_ID, CLIENT_ID, username=username)
+        cognito_user.authenticate(password)
+        cognito_user.check_token()
+        return cognito_user
+
+
+class IdClaimsData(OtfItemBase):
+    sub: str
+    email_verified: bool
+    iss: str
+    cognito_username: str = Field(alias="cognito:username")
+    given_name: str
+    locale: str
+    home_studio_id: str = Field(alias="custom:home_studio_id")
+    aud: str
+    event_id: str
+    token_use: str
+    auth_time: int
+    exp: int
+    is_migration: str = Field(alias="custom:isMigration")
+    iat: int
+    family_name: str
+    email: str
+    koji_person_id: str = Field(alias="custom:koji_person_id")
+
+    @property
+    def member_uuid(self) -> str:
+        return self.cognito_username
+
+    @property
+    def full_name(self) -> str:
+        return f"{self.given_name} {self.family_name}"
+
+
+class AccessClaimsData(OtfItemBase):
+    sub: str
+    device_key: str
+    iss: str
+    client_id: str
+    event_id: str
+    token_use: str
+    scope: str
+    auth_time: int
+    exp: int
+    iat: int
+    jti: str
+    username: str
+
+    @property
+    def member_uuid(self) -> str:
+        return self.username
+
+
+class OtfUser(OtfItemBase):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    cognito: OtfCognito
+
+    def __init__(
+        self,
+        username: str | None = None,
+        password: str | None = None,
+        id_token: str | None = None,
+        access_token: str | None = None,
+        refresh_token: str | None = None,
+        device_key: str | None = None,
+        cognito: OtfCognito | None = None,
+    ):
+        """Create a User instance.
+
+        Args:
+            username (str, optional): The username to login with. Defaults to None.
+            password (str, optional): The password to login with. Defaults to None.
+            id_token (str, optional): The id token. Defaults to None.
+            access_token (str, optional): The access token. Defaults to None.
+            refresh_token (str, optional): The refresh token. Defaults to None.
+            device_key (str, optional): The device key. Defaults to None.
+            cognito (OtfCognito, optional): A Cognito instance. Defaults to None.
+
+        Raises:
+            ValueError: Must provide either username and password or id token
+
+
+        """
+        if cognito:
+            cognito = cognito
+        elif username and password:
+            cognito = OtfCognito.login(username, password)
+        elif access_token and id_token:
+            cognito = OtfCognito.from_token(access_token, id_token, refresh_token, device_key)
+        else:
+            raise ValueError("Must provide either username and password or id token.")
+
+        super().__init__(cognito=cognito)
+
+    @property
+    def member_id(self) -> str:
+        return self.id_claims_data.cognito_username
+
+    @property
+    def member_uuid(self) -> str:
+        return self.access_claims_data.sub
+
+    @property
+    def access_claims_data(self) -> AccessClaimsData:
+        return AccessClaimsData(**self.cognito.access_claims)
+
+    @property
+    def id_claims_data(self) -> IdClaimsData:
+        return IdClaimsData(**self.cognito.id_claims)
+
+    def get_tokens(self) -> dict[str, str]:
+        return {
+            "id_token": self.cognito.id_token,
+            "access_token": self.cognito.access_token,
+            "refresh_token": self.cognito.refresh_token,
+        }
+
+    @property
+    def device_key(self) -> str:
+        return self.cognito.device_key

otf_api/cli/__init__.py

@@ -0,0 +1,4 @@
+from otf_api.cli.app import base_app
+from otf_api.cli.bookings import bookings_app, classes_app
+
+__all__ = ["base_app", "bookings_app", "classes_app"]

otf_api/cli/_utilities.py

@@ -0,0 +1,60 @@
+import functools
+import inspect
+import traceback
+from collections.abc import Awaitable, Callable
+from typing import Any, NoReturn, ParamSpec, TypeGuard, TypeVar
+
+import typer
+from click.exceptions import ClickException
+
+T = TypeVar("T")
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def is_async_fn(func: Callable[P, R] | Callable[P, Awaitable[R]]) -> TypeGuard[Callable[P, Awaitable[R]]]:
+    """
+    Returns `True` if a function returns a coroutine.
+
+    See https://github.com/microsoft/pyright/issues/2142 for an example use
+    """
+    while hasattr(func, "__wrapped__"):
+        func = func.__wrapped__
+
+    return inspect.iscoroutinefunction(func)
+
+
+def exit_with_error(message: Any, code: int = 1, **kwargs: Any) -> NoReturn:
+    """
+    Utility to print a stylized error message and exit with a non-zero code
+    """
+    from otf_api.cli.app import base_app
+
+    kwargs.setdefault("style", "red")
+    base_app.console.print(message, **kwargs)
+    raise typer.Exit(code)
+
+
+def exit_with_success(message: Any, **kwargs: Any) -> NoReturn:
+    """
+    Utility to print a stylized success message and exit with a zero code
+    """
+    from otf_api.cli.app import base_app
+
+    kwargs.setdefault("style", "green")
+    base_app.console.print(message, **kwargs)
+    raise typer.Exit(0)
+
+
+def with_cli_exception_handling(fn: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any | None:
+        try:
+            return fn(*args, **kwargs)
+        except (typer.Exit, typer.Abort, ClickException):
+            raise  # Do not capture click or typer exceptions
+        except Exception:
+            traceback.print_exc()
+            exit_with_error("An exception occurred.")
+
+    return wrapper

otf_api/cli/app.py

@@ -0,0 +1,172 @@
+import asyncio
+import functools
+import sys
+import typing
+from collections.abc import Callable
+from enum import Enum
+
+import typer
+from loguru import logger
+from rich.console import Console
+from rich.theme import Theme
+
+import otf_api
+from otf_api.cli._utilities import is_async_fn, with_cli_exception_handling
+
+if typing.TYPE_CHECKING:
+    from otf_api import Otf
+
+
+class OutputType(str, Enum):
+    json = "json"
+    table = "table"
+    interactive = "interactive"
+
+
+def version_callback(value: bool) -> None:
+    if value:
+        print(otf_api.__version__)
+        raise typer.Exit()
+
+
+OPT_USERNAME: str = typer.Option(None, envvar=["OTF_EMAIL", "OTF_USERNAME"], help="Username for the OTF API")
+OPT_PASSWORD: str = typer.Option(envvar="OTF_PASSWORD", help="Password for the OTF API", hide_input=True)
+OPT_OUTPUT: OutputType = typer.Option(None, envvar="OTF_OUTPUT", show_default=False, help="Output format")
+OPT_LOG_LEVEL: str = typer.Option("CRITICAL", help="Log level", envvar="OTF_LOG_LEVEL")
+OPT_VERSION: bool = typer.Option(
+    None, "--version", callback=version_callback, help="Display the current version.", is_eager=True
+)
+
+
+def register_main_callback(app: "AsyncTyper") -> None:
+    @app.callback()  # type: ignore
+    @with_cli_exception_handling
+    def main_callback(
+        ctx: typer.Context,  # noqa
+        version: bool = OPT_VERSION,  # noqa
+        username: str = OPT_USERNAME,
+        password: str = OPT_PASSWORD,
+        output: OutputType = OPT_OUTPUT,
+        log_level: str = OPT_LOG_LEVEL,
+    ) -> None:
+        app.setup_console()
+        app.set_username(username)
+        app.password = password
+        app.output = output or OutputType.table
+        app.set_log_level(log_level)
+
+        # When running on Windows we need to ensure that the correct event loop policy is
+        # in place or we will not be able to spawn subprocesses. Sometimes this policy is
+        # changed by other libraries, but here in our CLI we should have ownership of the
+        # process and be able to safely force it to be the correct policy.
+        # https://github.com/PrefectHQ/prefect/issues/8206
+        if sys.platform == "win32":
+            asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
+
+
+class AsyncTyper(typer.Typer):
+    """
+    Wraps commands created by `Typer` to support async functions and handle errors.
+    """
+
+    console: Console
+
+    def __init__(self, *args: typing.Any, **kwargs: typing.Any):
+        super().__init__(*args, **kwargs)
+
+        theme = Theme({"prompt.choices": "bold blue"})
+        self.console = Console(highlight=False, theme=theme, color_system="auto")
+
+        # TODO: clean these up later, just don't want warnings everywhere that these could be None
+        self.api: Otf = None  # type: ignore
+        self.username: str = None  # type: ignore
+        self.password: str = None  # type: ignore
+        self.output: OutputType = None  # type: ignore
+        self.logger = logger
+        self.log_level = "CRITICAL"
+
+    def set_username(self, username: str | None = None) -> None:
+        if username:
+            self.username = username
+            return
+
+        raise ValueError("Username not provided and not found in cache")
+
+    def set_log_level(self, level: str) -> None:
+        self.log_level = level
+        logger.remove()
+        logger.add(sys.stderr, level=self.log_level.upper())
+
+    def print(self, *args: typing.Any, **kwargs: typing.Any) -> None:
+        if self.output == "json":
+            self.console.print_json(*args, **kwargs)
+        else:
+            self.console.print(*args, **kwargs)
+
+    def add_typer(
+        self, typer_instance: "AsyncTyper", *args: typing.Any, aliases: list[str] | None = None, **kwargs: typing.Any
+    ) -> typing.Any:
+        aliases = aliases or []
+        for alias in aliases:
+            super().add_typer(typer_instance, *args, name=alias, no_args_is_help=True, hidden=True, **kwargs)
+
+        return super().add_typer(typer_instance, *args, no_args_is_help=True, **kwargs)
+
+    def command(
+        self, name: str | None = None, *args: typing.Any, aliases: list[str] | None = None, **kwargs: typing.Any
+    ) -> Callable[[typing.Any], typing.Any]:
+        """
+        Create a new command. If aliases are provided, the same command function
+        will be registered with multiple names.
+        """
+
+        aliases = aliases or []
+
+        def wrapper(fn: Callable[[typing.Any], typing.Any]) -> Callable[[typing.Any], typing.Any]:
+            # click doesn't support async functions, so we wrap them in
+            # asyncio.run(). This has the advantage of keeping the function in
+            # the main thread, which means signal handling works for e.g. the
+            # server and workers. However, it means that async CLI commands can
+            # not directly call other async CLI commands (because asyncio.run()
+            # can not be called nested). In that (rare) circumstance, refactor
+            # the CLI command so its business logic can be invoked separately
+            # from its entrypoint.
+            if is_async_fn(fn):
+                _fn = fn
+
+                @functools.wraps(fn)
+                def fn(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+                    return asyncio.run(_fn(*args, **kwargs))  # type: ignore
+
+                fn.aio = _fn  # type: ignore
+
+            fn = with_cli_exception_handling(fn)
+
+            # register fn with its original name
+            command_decorator = super(AsyncTyper, self).command(name=name, *args, **kwargs)
+            original_command = command_decorator(fn)
+
+            # register fn for each alias, e.g. @marvin_app.command(aliases=["r"])
+            for alias in aliases:
+                super(AsyncTyper, self).command(
+                    name=alias,
+                    *args,
+                    **{k: v for k, v in kwargs.items() if k != "aliases"},
+                )(fn)
+
+            return typing.cast(Callable[[typing.Any], typing.Any], original_command)
+
+        return wrapper
+
+    def setup_console(self, soft_wrap: bool = True, prompt: bool = True) -> None:
+        self.console = Console(
+            highlight=False,
+            color_system="auto",
+            theme=Theme({"prompt.choices": "bold blue"}),
+            soft_wrap=not soft_wrap,
+            force_interactive=prompt,
+        )
+
+
+base_app = AsyncTyper(add_completion=True, no_args_is_help=True)
+register_main_callback(base_app)