otf-api 0.8.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

otf_api/auth/__init__.py

@@ -0,0 +1,4 @@
+from .user import OtfUser
+from .utils import HttpxCognitoAuth
+
+__all__ = ["OTF_AUTH_TYPE", "HttpxCognitoAuth", "OtfAuth", "OtfUser"]

otf_api/auth/auth.py

@@ -0,0 +1,234 @@
+import platform
+import typing
+from logging import getLogger
+from pathlib import Path
+from time import sleep
+from typing import Any, ClassVar
+
+from boto3 import Session
+from botocore import UNSIGNED
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from pycognito import AWSSRP, Cognito
+from pycognito.aws_srp import generate_hash_device
+
+from otf_api.utils import CacheableData
+
+if typing.TYPE_CHECKING:
+    from mypy_boto3_cognito_idp import CognitoIdentityProviderClient
+    from mypy_boto3_cognito_idp.type_defs import InitiateAuthResponseTypeDef
+
+LOGGER = getLogger(__name__)
+CLIENT_ID = "1457d19r0pcjgmp5agooi0rb1b"  # from android app
+USER_POOL_ID = "us-east-1_dYDxUeyL1"
+REGION = "us-east-1"
+COGNITO_IDP_URL = f"https://cognito-idp.{REGION}.amazonaws.com/"
+
+BOTO_CONFIG = Config(region_name=REGION, signature_version=UNSIGNED)
+CRED_CACHE = CacheableData("creds", Path("~/.otf-api"))
+
+DEVICE_KEYS = ["device_key", "device_group_key", "device_password"]
+TOKEN_KEYS = ["access_token", "id_token", "refresh_token"]
+
+
+class NoCredentialsError(Exception):
+    """Raised when no credentials are found."""
+
+
+class OtfCognito(Cognito):
+    """A subclass of the pycognito Cognito class that adds the device_key to the auth_params. Without this
+    being set the renew_access_token call will always fail with NOT_AUTHORIZED."""
+
+    user_pool_id: ClassVar[str] = USER_POOL_ID
+    client_id: ClassVar[str] = CLIENT_ID
+    user_pool_region: ClassVar[str] = REGION
+    client_secret: ClassVar[str] = ""
+
+    id_token: str
+    access_token: str
+    device_key: str
+    device_group_key: str
+    device_password: str
+    device_name: str
+
+    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,
+    ):
+        """
+
+        Args:
+            username (str, optional): User Pool username
+            password (str, optional): User Pool password
+            id_token (str, optional): ID Token returned by authentication
+            access_token (str, optional): Access Token returned by authentication
+            refresh_token (str, optional): Refresh Token returned by authentication
+        """
+
+        self.username = username
+        self.id_token = id_token  # type: ignore
+        self.access_token = access_token  # type: ignore
+        self.refresh_token = refresh_token
+
+        self.id_claims: dict[str, Any] = {}
+        self.access_claims: dict[str, Any] = {}
+        self.custom_attributes: dict[str, Any] = {}
+        self.base_attributes: dict[str, Any] = {}
+        self.pool_jwk: dict[str, Any] = {}
+        self.mfa_tokens: dict[str, Any] = {}
+        self.pool_domain_url: str | None = None
+
+        try:
+            dd = CRED_CACHE.get_cached_data(DEVICE_KEYS)
+        except Exception:
+            LOGGER.exception("Failed to read device key cache")
+            dd = {}
+
+        self.device_name = platform.node()
+        self.device_key = dd.get("device_key")  # type: ignore
+        self.device_group_key = dd.get("device_group_key")  # type: ignore
+        self.device_password = dd.get("device_password")  # type: ignore
+
+        self.client: CognitoIdentityProviderClient = Session().client(
+            "cognito-idp", config=BOTO_CONFIG, region_name=REGION
+        )
+
+        try:
+            token_cache = CRED_CACHE.get_cached_data(TOKEN_KEYS)
+        except Exception:
+            LOGGER.exception("Failed to read token cache")
+            token_cache = {}
+
+        if not (username and password) and not (id_token and access_token) and not token_cache:
+            raise NoCredentialsError("No credentials provided and no tokens cached, cannot authenticate")
+
+        if username and password:
+            self.login(password)
+        elif token_cache and not (id_token and access_token):
+            LOGGER.debug("Using cached tokens")
+            self.id_token = token_cache["id_token"]
+            self.access_token = token_cache["access_token"]
+            self.refresh_token = token_cache["refresh_token"]
+
+        self.check_token()
+        self.verify_tokens()
+        CRED_CACHE.write_to_cache(self.tokens)
+
+    @property
+    def tokens(self) -> dict[str, str]:
+        tokens = {
+            "access_token": self.access_token,
+            "id_token": self.id_token,
+            "refresh_token": self.refresh_token,
+        }
+        return {k: v for k, v in tokens.items() if v}
+
+    @property
+    def device_metadata(self) -> dict[str, str]:
+        dm = {
+            "device_key": self.device_key,
+            "device_group_key": self.device_group_key,
+            "device_password": self.device_password,
+        }
+        return {k: v for k, v in dm.items() if v}
+
+    def login(self, password: str) -> None:
+        """Called when logging in with a username and password. Will set the tokens and device metadata."""
+
+        LOGGER.debug("Logging in with username and password...")
+
+        aws = AWSSRP(
+            username=self.username,
+            password=password,
+            pool_id=USER_POOL_ID,
+            client_id=CLIENT_ID,
+            client=self.client,
+        )
+        try:
+            tokens: InitiateAuthResponseTypeDef = aws.authenticate_user()
+        except ClientError as e:
+            code = e.response["Error"]["Code"]
+            msg = e.response["Error"]["Message"]
+            if "UserLambdaValidationException" in msg or "UserLambdaValidation" in code:
+                sleep(5)
+                tokens = aws.authenticate_user()
+            else:
+                raise
+
+        self._set_tokens(tokens)
+        self._handle_device_setup()
+
+    def _handle_device_setup(self) -> None:
+        """Confirms the device with Cognito and caches the device metadata.
+
+        Devices are not remembered at this time, as OTF does not have MFA set up currently. Without MFA setup, there
+        is no benefit to remembering the device. Additionally, it does not appear that the OTF app remembers devices,
+        so this matches the behavior of the app.
+        """
+
+        if not self.device_key:
+            raise ValueError("Device key not set - device key is required by this Cognito pool")
+
+        self.device_password, device_secret_verifier_config = generate_hash_device(
+            self.device_group_key, self.device_key
+        )
+
+        self.client.confirm_device(
+            AccessToken=self.access_token,
+            DeviceKey=self.device_key,
+            DeviceSecretVerifierConfig=device_secret_verifier_config,
+            DeviceName=self.device_name,
+        )
+
+        try:
+            CRED_CACHE.write_to_cache(self.device_metadata)
+        except Exception:
+            LOGGER.exception("Failed to write device key cache")
+
+    ##### OVERRIDDEN METHODS #####
+
+    def renew_access_token(self) -> None:
+        """Sets a new access token on the User using the cached refresh token and device metadata.
+
+        Overridden to add the device key to the auth_params if it is set. Without this all calls to renew_access_token
+        will fail with NOT_AUTHORIZED. Also skips the call to _add_secret_hash since we don't have a client secret.
+
+        # https://docs.aws.amazon.com/cognito-user-identity-pools/latest/APIReference/API_InitiateAuth.html#CognitoUserPools-InitiateAuth-request-AuthFlow
+
+        """
+        if not self.device_key:
+            raise ValueError("Device key not set - device key is required by this Cognito pool")
+
+        if not self.refresh_token:
+            raise ValueError("No refresh token set - cannot renew access token")
+
+        refresh_response = self.client.initiate_auth(
+            ClientId=self.client_id,
+            AuthFlow="REFRESH_TOKEN_AUTH",
+            AuthParameters={"REFRESH_TOKEN": self.refresh_token, "DEVICE_KEY": self.device_key},
+        )
+        self._set_tokens(refresh_response)
+
+    def _set_tokens(self, tokens: "InitiateAuthResponseTypeDef") -> None:
+        """Helper function to verify and set token attributes based on a Cognito AuthenticationResult.
+
+        Overridden to cache the tokens and set/cache the device metadata.
+        """
+        auth_result = tokens["AuthenticationResult"]
+        device_metadata = auth_result.get("NewDeviceMetadata", {})
+
+        # tokens - refresh token defaults to existing value if not present
+        # note: verify_token also sets the token attribute
+        self.verify_token(auth_result["AccessToken"], "access_token", "access")
+        self.verify_token(auth_result["IdToken"], "id_token", "id")
+        self.refresh_token = auth_result.get("RefreshToken", self.refresh_token)
+        CRED_CACHE.write_to_cache(self.tokens)
+
+        # device metadata - default to existing values if not present
+        self.device_key = device_metadata.get("DeviceKey", self.device_key)
+        self.device_group_key = device_metadata.get("DeviceGroupKey", self.device_group_key)
+        CRED_CACHE.write_to_cache(self.device_metadata)

otf_api/auth/user.py

@@ -0,0 +1,66 @@
+from logging import getLogger
+
+import attrs
+
+from otf_api.auth.auth import CRED_CACHE, NoCredentialsError, OtfCognito
+from otf_api.auth.utils import HttpxCognitoAuth, can_provide_input, prompt_for_username_and_password
+
+LOGGER = getLogger(__name__)
+
+
+@attrs.define(init=False)
+class OtfUser:
+    """OtfUser is a thin wrapper around OtfCognito, meant to hide all of the gory details from end users."""
+
+    cognito_id: str
+    member_uuid: str
+    email_address: str
+    cognito: OtfCognito
+    httpx_auth: HttpxCognitoAuth
+
+    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,
+    ):
+        """Create a User instance.
+
+        Args:
+            username (str, optional): User Pool username
+            password (str, optional): User Pool password
+            id_token (str, optional): ID Token returned by authentication
+            access_token (str, optional): Access Token returned by authentication
+            refresh_token (str, optional): Refresh Token returned by authentication
+
+        Raises:
+            NoCredentialsError: If neither username/password nor id/access tokens are provided.
+        """
+        try:
+            self.cognito = OtfCognito(
+                username=username,
+                password=password,
+                id_token=id_token,
+                access_token=access_token,
+                refresh_token=refresh_token,
+            )
+        except NoCredentialsError:
+            if not can_provide_input():
+                LOGGER.error("Unable to prompt for credentials in a non-interactive shell")
+                raise
+
+            username, password = prompt_for_username_and_password()
+            self.cognito = OtfCognito(username=username, password=password)
+
+        self.cognito_id = self.cognito.access_claims["sub"]
+        self.member_uuid = self.cognito.id_claims["cognito:username"]
+        self.email_address = self.cognito.id_claims["email"]
+
+        self.httpx_auth = HttpxCognitoAuth(cognito=self.cognito)
+
+    @staticmethod
+    def clear_cache():
+        """Clear the cached credentials."""
+        CRED_CACHE.clear_cache()

otf_api/auth/utils.py

@@ -0,0 +1,129 @@
+import os
+import sys
+from collections.abc import Generator
+from functools import partial
+from getpass import getpass
+from logging import getLogger
+from typing import Any
+
+import httpx
+from httpx import Request
+from pycognito import Cognito
+
+LOGGER = getLogger(__name__)
+USERNAME_PROMPT = "Enter your Orangetheory Fitness username/email: "
+PASSWORD_PROMPT = "Enter your Orangetheory Fitness password: "
+
+
+def _show_error_message(message: str) -> None:
+    try:
+        from rich import print  # type: ignore
+
+        print(message, style="bold red")
+    except ImportError:
+        print(message)
+
+
+def _get_password_input(prompt: str) -> str:
+    try:
+        from rich import get_console  # type: ignore
+
+        otf_input = partial(get_console().input, password=True)
+        prompt = f"[bold blue]{prompt}[/bold blue]"
+    except ImportError:
+        otf_input = getpass
+
+    return otf_input(prompt)
+
+
+def _get_input(prompt: str) -> str:
+    try:
+        from rich import get_console  # type: ignore
+
+        otf_input = get_console().input
+        prompt = f"[bold blue]{prompt}[/bold blue]"
+    except ImportError:
+        otf_input = input
+
+    return otf_input(prompt)
+
+
+def _prompt_for_username() -> str:
+    def get_username() -> str:
+        username = _get_input(USERNAME_PROMPT)
+
+        if not username:
+            _show_error_message("Username is required")
+            return ""
+
+        if "@" not in username or username.endswith("@"):
+            _show_error_message("Username should be a valid email address")
+            return ""
+
+        return username
+
+    while not (username := get_username()):
+        pass
+
+    return username
+
+
+def _prompt_for_password() -> str:
+    def get_password() -> str:
+        password = _get_password_input(PASSWORD_PROMPT)
+
+        if not password:
+            _show_error_message("Password is required")
+            return ""
+
+        return password
+
+    while not (password := get_password()):
+        pass
+
+    return password
+
+
+def prompt_for_username_and_password() -> tuple[str, str]:
+    """Prompt for a username and password.
+
+    Returns:
+        tuple[str, str]: A tuple containing the username and password.
+    """
+
+    username = os.getenv("OTF_EMAIL") or _prompt_for_username()
+    password = os.getenv("OTF_PASSWORD") or _prompt_for_password()
+
+    return username, password
+
+
+def can_provide_input() -> bool:
+    """Check if the script is running in an interactive shell.
+
+    Returns:
+        bool: True if the script is running in an interactive shell.
+    """
+    return os.isatty(sys.stdin.fileno()) and os.isatty(sys.stdout.fileno())
+
+
+class HttpxCognitoAuth(httpx.Auth):
+    http_header: str = "Authorization"
+    http_header_prefix: str = "Bearer "
+
+    def __init__(self, cognito: Cognito):
+        """HTTPX Authentication extension for Cognito User Pools.
+
+        Args:
+            cognito (Cognito): A Cognito instance.
+        """
+
+        self.cognito = cognito
+
+    def auth_flow(self, request: Request) -> Generator[Request, Any, None]:
+        self.cognito.check_token(renew=True)
+
+        token = self.cognito.id_token
+
+        request.headers[self.http_header] = self.http_header_prefix + token
+
+        yield request

otf_api/exceptions.py

@@ -1,4 +1,25 @@
-class BookingError(Exception):
+from httpx import Request, Response
+
+
+class OtfException(Exception):
+    """Base class for all exceptions in this package."""
+
+
+class OtfRequestError(OtfException):
+    """Raised when an error occurs while making a request to the OTF API."""
+
+    response: Response
+    request: Request
+
+    def __init__(self, message: str, response: Response, request: Request):
+        super().__init__(message)
+        self.response = response
+        self.request = request
+
+
+class BookingError(OtfException):
+    """Base class for booking-related errors, with an optional booking UUID attribute."""
+
     booking_uuid: str | None
 
     def __init__(self, message: str, booking_uuid: str | None = None):
@@ -6,13 +27,25 @@ class BookingError(Exception):
         self.booking_uuid = booking_uuid
 
 
-class AlreadyBookedError(BookingError): ...
+class AlreadyBookedError(BookingError):
+    """Raised when attempting to book a class that is already booked."""
+
+
+class ConflictingBookingError(BookingError):
+    """Raised when attempting to book a class that conflicts with an existing booking."""
+
+
+class BookingAlreadyCancelledError(BookingError):
+    """Raised when attempting to cancel a booking that is already cancelled."""
 
 
-class BookingAlreadyCancelledError(BookingError): ...
+class OutsideSchedulingWindowError(OtfException):
+    """Raised when attempting to book a class outside the scheduling window."""
 
 
-class OutsideSchedulingWindowError(Exception): ...
+class BookingNotFoundError(OtfException):
+    """Raised when a booking is not found."""
 
 
-class BookingNotFoundError(Exception): ...
+class ResourceNotFoundError(OtfException):
+    """Raised when a resource is not found."""

otf_api/filters.py

@@ -0,0 +1,97 @@
+from datetime import date, datetime, time
+
+from pydantic import BaseModel, field_validator
+
+from otf_api.models import ClassType, DoW, OtfClass
+
+
+class ClassFilter(BaseModel):
+    """ClassFilter is used to filter classes, to separate the filtering logic from the API client.
+
+    The `class_type`, `day_of_week`, and `start_time` fields can either be a single value or a list of values.
+    If a single value is provided, it will be converted to a list.
+
+    The `class_type` and `day_of_week` fields can be provided as strings or as the corresponding Enum values. The
+    class will attempt to match the string to the Enum value, regardless of case.
+
+    The arguments are applied as an AND filter, meaning that all filters must match for a class to be included. If
+    a filter is not provided, it is not applied. If multiple values are provided for a filter the values are treated
+    as an OR filter, meaning that a class will be included if it matches any of the values.
+
+    All arguments are optional and default to None.
+
+    Args:
+        start_date (date): Filter classes that start on or after this date.
+        end_date (date): Filter classes that start on or before this date.
+        class_type (list[ClassType]): Filter classes by class type.
+        day_of_week (list[DoW]): Filter classes by day of the week.
+        start_time (list[time]): Filter classes by start time.
+    """
+
+    start_date: date | None = None
+    end_date: date | None = None
+    class_type: list[ClassType] | None = None
+    day_of_week: list[DoW] | None = None
+    start_time: list[time] | None = None
+
+    def filter_classes(self, classes: list[OtfClass]) -> list[OtfClass]:
+        """Filters a list of classes based on the filter arguments.
+
+        Args:
+            classes (list[OtfClass]): A list of classes to filter.
+
+        Returns:
+            list[OtfClass]: The filtered list of classes.
+        """
+        # in case these are set after the class is created
+        if self.start_date and isinstance(self.start_date, datetime):
+            self.start_date = self.start_date.date()
+
+        if self.end_date and isinstance(self.end_date, datetime):
+            self.end_date = self.end_date.date()
+
+        if self.start_date:
+            classes = [c for c in classes if c.starts_at.date() >= self.start_date]
+
+        if self.end_date:
+            classes = [c for c in classes if c.starts_at.date() <= self.end_date]
+
+        if self.class_type:
+            classes = [c for c in classes if c.class_type in self.class_type]
+
+        if self.day_of_week:
+            classes = [c for c in classes if c.day_of_week in self.day_of_week]
+
+        if self.start_time:
+            classes = [c for c in classes if c.starts_at.time() in self.start_time]
+
+        return classes
+
+    @field_validator("class_type", "day_of_week", "start_time", mode="before")
+    @classmethod
+    def _single_item_to_list(cls, v):
+        if v and not isinstance(v, list):
+            return [v]
+        return v
+
+    @field_validator("day_of_week", mode="before")
+    @classmethod
+    def _day_of_week_str_to_enum(cls, v):
+        if v and isinstance(v, str):
+            return [DoW(v.title())]
+
+        if v and isinstance(v, list) and not all(isinstance(i, DoW) for i in v):
+            return [DoW(i.title()) for i in v]
+
+        return v
+
+    @field_validator("class_type", mode="before")
+    @classmethod
+    def _class_type_str_to_enum(cls, v):
+        if v and isinstance(v, str):
+            return [ClassType.get_case_insensitive(v)]
+
+        if v and isinstance(v, list) and not all(isinstance(i, ClassType) for i in v):
+            return [ClassType.get_case_insensitive(i) for i in v]
+
+        return v

otf_api/logging.py

@@ -0,0 +1,19 @@
+import logging
+import os
+
+LOG_LEVEL = os.getenv("OTF_LOG_LEVEL", "INFO")
+
+LOG_FMT = "{asctime} - {module}.{funcName}:{lineno} - {levelname} - {message}"
+DATE_FMT = "%Y-%m-%d %H:%M:%S%z"
+
+logger = logging.getLogger("otf_api")
+
+# 2) Set the logger level to INFO (or whatever you need).
+logger.setLevel(LOG_LEVEL)
+
+# 3) Create a handler (e.g., console) and set its formatter.
+handler = logging.StreamHandler()
+handler.setFormatter(logging.Formatter(fmt=LOG_FMT, datefmt=DATE_FMT, style="{"))
+
+# 4) Add this handler to your package logger.
+logger.addHandler(handler)