diff --git a/CHANGELOG.md b/CHANGELOG.md index e99e173..252a085 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,29 @@ +# v16.0.0 +- Added `ApiTokenAuth` for authenticating with a API token. +- Added `DynamicApiTokenAuth` / `AsyncDynamicApiTokenAuth` for authenticating with API tokens with automatic rotation (sync + async). +- Added API token storage implementations: in-memory and JSON file (sync + async). + +Examples: +```python +from rtbhouse_sdk.client import ApiTokenAuth, Client + +auth = ApiTokenAuth(token="your_api_token") +api = Client(auth=auth) + +info = api.get_user_info() +api.close() +``` +or in async +```python +from rtbhouse_sdk.client import ApiTokenAuth, AsyncClient + +auth = ApiTokenAuth(token="your_api_token") +api = AsyncClient(auth=auth) + +info = await api.get_user_info() +api.close() +``` + # v15.0.0 - [breaking change] dropped support for python 3.9 (which is end-of-life), please use python 3.10+ diff --git a/poetry.lock b/poetry.lock index 66e91fe..d812817 100644 --- a/poetry.lock +++ b/poetry.lock @@ -107,6 +107,17 @@ d = ["aiohttp (>=3.10)"] jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"] uvloop = ["uvloop (>=0.15.2)"] +[[package]] +name = "cachetools" +version = "7.0.3" +description = "Extensible memoizing collections and decorators" +optional = false +python-versions = ">=3.10" +files = [ + {file = "cachetools-7.0.3-py3-none-any.whl", hash = "sha256:c128ffca156eef344c25fcd08a96a5952803786fa33097f5f2d49edf76f79d53"}, + {file = "cachetools-7.0.3.tar.gz", hash = "sha256:8c246313b95849964e54a909c03b327a87ab0428b068fac10da7b105ca275ef6"}, +] + [[package]] name = "certifi" version = "2026.1.4" @@ -992,6 +1003,17 @@ files = [ {file = "tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1"}, ] +[[package]] +name = "types-cachetools" +version = "6.2.0.20251022" +description = "Typing stubs for cachetools" +optional = false +python-versions = ">=3.9" +files = [ + {file = "types_cachetools-6.2.0.20251022-py3-none-any.whl", hash = "sha256:698eb17b8f16b661b90624708b6915f33dbac2d185db499ed57e4997e7962cad"}, + {file = "types_cachetools-6.2.0.20251022.tar.gz", hash = "sha256:f1d3c736f0f741e89ec10f0e1b0138625023e21eb33603a930c149e0318c0cef"}, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -1020,4 +1042,4 @@ typing-extensions = ">=4.12.0" [metadata] lock-version = "2.0" python-versions = ">=3.10, <4.0" -content-hash = "e51c4079d07bf1b47173e7747697435e4548edd1755e0edd3446774bcd26d188" +content-hash = "39616a78e3908e5224653b788f21a2bf75dae01895141a2b2f64a94b5644c559" diff --git a/pyproject.toml b/pyproject.toml index 9d021fe..8cd72f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "rtbhouse-sdk" -version = "15.0.0" +version = "16.0.0" description = "RTB House SDK" authors = ["RTB House Apps Team "] license = "BSD License" @@ -28,6 +28,8 @@ python = ">=3.10, <4.0" httpx = "^0.28.0" pydantic = ">=1.9, <3.0" +cachetools = "^7.0.3" +types-cachetools = "^6.2.0.20251022" [tool.poetry.group.dev.dependencies] pydantic = "^2.0.0" # required for tests diff --git a/rtbhouse_sdk/__init__.py b/rtbhouse_sdk/__init__.py index 84adc89..929e16d 100644 --- a/rtbhouse_sdk/__init__.py +++ b/rtbhouse_sdk/__init__.py @@ -1,3 +1,3 @@ """RTB House Python SDK.""" -__version__ = "15.0.0" +__version__ = "16.0.0" diff --git a/rtbhouse_sdk/__main__.py b/rtbhouse_sdk/__main__.py new file mode 100644 index 0000000..bcc8dde --- /dev/null +++ b/rtbhouse_sdk/__main__.py @@ -0,0 +1,6 @@ +"""Entry point for the RTB House Python SDK CLI.""" + +from .api_tokens.cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/rtbhouse_sdk/_utils.py b/rtbhouse_sdk/_utils.py index 1d216ad..a7172f9 100644 --- a/rtbhouse_sdk/_utils.py +++ b/rtbhouse_sdk/_utils.py @@ -1,6 +1,7 @@ """Utils used in SDK.""" import re +from datetime import datetime, timezone from pydantic.version import VERSION as PYDANTIC_VERSION @@ -28,3 +29,7 @@ def underscore(word: str) -> str: word = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", word) word = word.replace("-", "_") return word.lower() + + +def utcnow() -> datetime: + return datetime.now(timezone.utc) diff --git a/rtbhouse_sdk/api_tokens/__init__.py b/rtbhouse_sdk/api_tokens/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rtbhouse_sdk/api_tokens/cli.py b/rtbhouse_sdk/api_tokens/cli.py new file mode 100644 index 0000000..fe3efcb --- /dev/null +++ b/rtbhouse_sdk/api_tokens/cli.py @@ -0,0 +1,192 @@ +"""CLI for API token management.""" + +import argparse +import getpass +import os +import sys +from pathlib import Path +from typing import Any + +from ..client import ApiTokenAuth, Client +from ..exceptions import ApiRequestException +from ..schema import ApiToken as ApiTokenResponse +from .managers import ApiTokenExpiredException, ApiTokenManager +from .models import ApiToken +from .storages.base import ApiTokenStorageException +from .storages.json_file import DEFAULT_JSON_FILE_PATH, JsonFileApiTokenStorage + + +def _read_token_from_stdin() -> str: + if sys.stdin.isatty(): + token = getpass.getpass("Paste API token: ").strip() + return token + + token = sys.stdin.read().strip() + + return token + + +def _resolve_token( + *, + token_arg: str | None = None, + env_var: str | None = None, +) -> str: + print(token_arg, env_var) + if token_arg: + return token_arg.strip() + + if env_var: + env_val = os.getenv(env_var) + return env_val.strip() if env_val else "" + + return _read_token_from_stdin() + + +def _get_token(token: str) -> ApiTokenResponse: + with Client(auth=ApiTokenAuth(token=token)) as client: + return client.get_current_api_token() + + +def cmd_init_json(args: argparse.Namespace) -> int: + path = Path(args.path).expanduser() + + token = _resolve_token(token_arg=args.token, env_var=args.env_var) + if not token: + print("ERROR: Empty token. Aborting.", file=sys.stderr) + return 1 + + try: + api_token_response = _get_token(token) + except ApiRequestException as e: + print(f"ERROR: Could not verify token via API or token is invalid. Original error: {e}.", file=sys.stderr) + return 1 + + storage = JsonFileApiTokenStorage(path) + api_token = ApiToken( + token=token, + expires_at=api_token_response.expires_at, + ) + storage.save(api_token) + + print(f"OK: Token successfully initialized in {path}", file=sys.stdout) + return 0 + + +def cmd_keep_alive_json(args: argparse.Namespace) -> int: + path = Path(args.path).expanduser() + storage = JsonFileApiTokenStorage(path) + manager = ApiTokenManager(storage=storage) + + try: + manager.keep_alive(args.auto_rotate) + except (ApiTokenStorageException, ApiRequestException, ApiTokenExpiredException) as e: + print(f"ERROR: Keep-alive failed. Original error: {e}", file=sys.stderr) + return 1 + + print("OK: Token keep-alive successful", file=sys.stdout) + return 0 + + +def cmd_keep_alive(args: argparse.Namespace) -> int: + token = _resolve_token(token_arg=args.token, env_var=args.env_var) + if not token: + print("ERROR: Empty token. Aborting.", file=sys.stderr) + return 1 + + try: + api_token_response = _get_token(token) + except ApiRequestException as e: + print(f"ERROR: Could not verify token via API or token is invalid. Original error: {e}.", file=sys.stderr) + return 1 + + if api_token_response.can_rotate: + print("WARNING: token can be rotated. Consider rotating it before it expires.", file=sys.stderr) + + print("OK: Token keep-alive successful", file=sys.stdout) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="rtbhouse_sdk", + description="RTB House Python SDK CLI", + ) + sub = parser.add_subparsers(dest="command", required=True) + + api_token = sub.add_parser( + "api-token", + help="API token utilities", + ) + api_token_sub = api_token.add_subparsers( + dest="api_token_cmd", + required=True, + ) + build_cmd_init_json(api_token_sub) + build_cmd_keep_alive_json(api_token_sub) + build_cmd_keep_alive(api_token_sub) + + return parser + + +def build_cmd_init_json(parser: Any) -> None: + init_json = parser.add_parser( + "init-json", + help="Initialize API token JSON file storage", + ) + init_json.add_argument( + "--token", + default=None, + help="Token value (discouraged). Prefer stdin/env/pipe.", + ) + init_json.add_argument( + "--path", + default=DEFAULT_JSON_FILE_PATH, + help=f"Path to token JSON file (default: {DEFAULT_JSON_FILE_PATH})", + ) + init_json.add_argument( + "--env-var", + help="Environment variable to read token from.", + ) + init_json.set_defaults(func=cmd_init_json) + + +def build_cmd_keep_alive_json(parser: Any) -> None: + keep_alive_json = parser.add_parser( + "keep-alive-json", + help="Keep alive API token stored in JSON file storage by bumping its usage. Allows automatic rotation.", + ) + keep_alive_json.add_argument( + "--path", + default=DEFAULT_JSON_FILE_PATH, + help=f"Path to token JSON file (default: {DEFAULT_JSON_FILE_PATH})", + ) + keep_alive_json.add_argument( + "--auto-rotate", + action="store_true", + help="Enable automatic rotation if token is in rotation window", + ) + keep_alive_json.set_defaults(func=cmd_keep_alive_json) + + +def build_cmd_keep_alive(parser: Any) -> None: + keep_alive = parser.add_parser( + "keep-alive", + help="Keep alive API token by bumping its usage.", + ) + keep_alive.add_argument( + "--token", + default=None, + help="Token value (discouraged). Prefer stdin/env/pipe.", + ) + keep_alive.add_argument( + "--env-var", + default=None, + help="Environment variable to read token from.", + ) + keep_alive.set_defaults(func=cmd_keep_alive) + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + return int(args.func(args)) diff --git a/rtbhouse_sdk/api_tokens/managers.py b/rtbhouse_sdk/api_tokens/managers.py new file mode 100644 index 0000000..f48c264 --- /dev/null +++ b/rtbhouse_sdk/api_tokens/managers.py @@ -0,0 +1,204 @@ +"""Contains classes for managing API tokens""" + +import asyncio +import threading +import warnings +from collections.abc import AsyncIterator, Iterator +from contextlib import asynccontextmanager, contextmanager +from datetime import datetime, timedelta + +from .._utils import utcnow +from ..client import ApiTokenAuth, AsyncClient, AsyncDynamicApiTokenAuth, Client, DynamicApiTokenAuth +from ..exceptions import ApiRequestException +from .models import ApiToken +from .storages.base import ApiTokenStorage, AsyncApiTokenStorage + +ROTATION_WINDOW = timedelta(days=4) + + +EXPIRED_MSG = "API token expired. Please manually create a new one and configure it in storage." + + +class ApiTokenExpiredException(Exception): + """Exception raised when the API token has expired.""" + + +class ApiTokenManager(DynamicApiTokenAuth): + """ + Manages the lifecycle of an API token, including token retrieval, rotation and expiration handling. + """ + + _storage: ApiTokenStorage + _lock: threading.Lock + _expiration_margin: timedelta + + def __init__(self, storage: ApiTokenStorage) -> None: + super().__init__() + self._storage = storage + self._lock = threading.Lock() + self._expiration_margin = timedelta(minutes=1) + + @contextmanager + def with_client(self, token: str) -> Iterator[Client]: + client = Client(auth=ApiTokenAuth(token=token)) + try: + yield client + finally: + client.close() + + def get_token(self) -> str: + api_token = self._storage.load() + now = utcnow() + + _raise_if_expired(api_token, now, self._expiration_margin) + + if not _in_rotation_window(api_token, now): + # fast path + return api_token.token + + with self._lock: + # reload inside lock + api_token = self._storage.load() + now = utcnow() + + _raise_if_expired(api_token, now, self._expiration_margin) + + if not _in_rotation_window(api_token, now): + return api_token.token + + with self.with_client(api_token.token) as client: + try: + rotated = client.rotate_current_api_token() + + api_token = ApiToken( + token=rotated.token, + expires_at=rotated.expires_at, + ) + self._storage.save(api_token) + except ApiRequestException as e: + warnings.warn( + f"Attempted to rotate API token but failed. " + "Please check whether the token has already been rotated. " + f"Original error: {e}" + ) + + return api_token.token + + def keep_alive(self, auto_rotate: bool = False) -> None: + with self._lock: + api_token = self._storage.load() + now = utcnow() + + _raise_if_expired(api_token, now, self._expiration_margin) + + with self.with_client(api_token.token) as client: + # bump used at to now to keep the token alive + client.get_current_api_token() + + if not auto_rotate: + return + + if not _in_rotation_window(api_token, now): + return + + rotated = client.rotate_current_api_token() + + api_token = ApiToken( + token=rotated.token, + expires_at=rotated.expires_at, + ) + self._storage.save(api_token) + + +class AsyncApiTokenManager(AsyncDynamicApiTokenAuth): + """Asynchronous version of ApiTokenManager.""" + + _storage: AsyncApiTokenStorage + _lock: asyncio.Lock + _expiration_margin: timedelta + + def __init__(self, storage: AsyncApiTokenStorage) -> None: + super().__init__() + self._storage = storage + self._lock = asyncio.Lock() + self._expiration_margin = timedelta(minutes=1) + + @asynccontextmanager + async def with_client(self, token: str) -> AsyncIterator[AsyncClient]: + client = AsyncClient(auth=ApiTokenAuth(token=token)) + try: + yield client + finally: + await client.close() + + async def get_token(self) -> str: + api_token = await self._storage.load() + now = utcnow() + + _raise_if_expired(api_token, now, self._expiration_margin) + + if not _in_rotation_window(api_token, now): + # fast path + return api_token.token + + async with self._lock: + # reload inside lock + api_token = await self._storage.load() + now = utcnow() + + _raise_if_expired(api_token, now, self._expiration_margin) + + if not _in_rotation_window(api_token, now): + return api_token.token + + async with self.with_client(api_token.token) as client: + try: + rotated = await client.rotate_current_api_token() + + api_token = ApiToken( + token=rotated.token, + expires_at=rotated.expires_at, + ) + await self._storage.save(api_token) + except ApiRequestException as e: + warnings.warn( + f"Attempted to rotate API token but failed. " + "Please check whether the token has already been rotated. " + f"Original error: {e}" + ) + + return api_token.token + + async def keep_alive(self, auto_rotate: bool) -> None: + async with self._lock: + api_token = await self._storage.load() + now = utcnow() + + _raise_if_expired(api_token, now, self._expiration_margin) + + async with self.with_client(api_token.token) as client: + # bump used at to now to keep the token alive + await client.get_current_api_token() + + if not auto_rotate: + return + + if not _in_rotation_window(api_token, now): + return + + rotated = await client.rotate_current_api_token() + + api_token = ApiToken( + token=rotated.token, + expires_at=rotated.expires_at, + ) + await self._storage.save(api_token) + + +def _raise_if_expired(api_token: ApiToken, now: datetime, expiration_margin: timedelta) -> None: + if now >= api_token.expires_at - expiration_margin: + raise ApiTokenExpiredException(EXPIRED_MSG) + + +def _in_rotation_window(api_token: ApiToken, now: datetime) -> bool: + return now >= (api_token.expires_at - ROTATION_WINDOW) diff --git a/rtbhouse_sdk/api_tokens/models.py b/rtbhouse_sdk/api_tokens/models.py new file mode 100644 index 0000000..54fe0e9 --- /dev/null +++ b/rtbhouse_sdk/api_tokens/models.py @@ -0,0 +1,12 @@ +"""Contains models related to API tokens.""" + +from datetime import datetime + +from pydantic import BaseModel + + +class ApiToken(BaseModel): + """Basic model representing an API token.""" + + token: str + expires_at: datetime diff --git a/rtbhouse_sdk/api_tokens/storages/__init__.py b/rtbhouse_sdk/api_tokens/storages/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rtbhouse_sdk/api_tokens/storages/base.py b/rtbhouse_sdk/api_tokens/storages/base.py new file mode 100644 index 0000000..39071c7 --- /dev/null +++ b/rtbhouse_sdk/api_tokens/storages/base.py @@ -0,0 +1,33 @@ +"""Base classes for API token storage implementations.""" + +from abc import ABC, abstractmethod + +from ..models import ApiToken + + +class ApiTokenStorageException(Exception): + """Exception raised for errors in the API token storage operations.""" + + +class ApiTokenStorage(ABC): + """Abstract base class for API token storage implementations.""" + + @abstractmethod + def load(self) -> ApiToken: + pass + + @abstractmethod + def save(self, api_token: ApiToken) -> None: + pass + + +class AsyncApiTokenStorage(ABC): + """Abstract base class for asynchronous API token storage implementations.""" + + @abstractmethod + async def load(self) -> ApiToken: + pass + + @abstractmethod + async def save(self, api_token: ApiToken) -> None: + pass diff --git a/rtbhouse_sdk/api_tokens/storages/in_memory.py b/rtbhouse_sdk/api_tokens/storages/in_memory.py new file mode 100644 index 0000000..9f14bac --- /dev/null +++ b/rtbhouse_sdk/api_tokens/storages/in_memory.py @@ -0,0 +1,40 @@ +"""In-memory storage for API tokens.""" + +from ..models import ApiToken +from .base import ApiTokenStorage, ApiTokenStorageException, AsyncApiTokenStorage + + +class InMemoryApiTokenStorage(ApiTokenStorage): + """In-memory storage for API tokens.""" + + _api_token: ApiToken + + def __init__(self, api_token: ApiToken) -> None: + super().__init__() + self._api_token = api_token + + def load(self) -> ApiToken: + if self._api_token is None: + raise ApiTokenStorageException("No API token stored in memory") + return self._api_token + + def save(self, api_token: ApiToken) -> None: + self._api_token = api_token + + +class AsyncInMemoryApiTokenStorage(AsyncApiTokenStorage): + """Asynchronous in-memory storage for API tokens.""" + + _api_token: ApiToken + + def __init__(self, api_token: ApiToken) -> None: + super().__init__() + self._api_token = api_token + + async def load(self) -> ApiToken: + if self._api_token is None: + raise ApiTokenStorageException("No API token stored in memory") + return self._api_token + + async def save(self, api_token: ApiToken) -> None: + self._api_token = api_token diff --git a/rtbhouse_sdk/api_tokens/storages/json_file.py b/rtbhouse_sdk/api_tokens/storages/json_file.py new file mode 100644 index 0000000..c17687d --- /dev/null +++ b/rtbhouse_sdk/api_tokens/storages/json_file.py @@ -0,0 +1,97 @@ +"""JSON file storage for API tokens.""" + +import asyncio +import os +import tempfile +from pathlib import Path + +from cachetools import TTLCache, cachedmethod +from cachetools.keys import hashkey + +from ..models import ApiToken +from .base import ApiTokenStorage, ApiTokenStorageException, AsyncApiTokenStorage + +DEFAULT_JSON_FILE_PATH = "~/.rtbhouse/api_token.json" + + +class JsonFileApiTokenStorage(ApiTokenStorage): + """JSON file storage for API tokens.""" + + _path: Path + _load_cache: TTLCache[str, ApiToken] + + def __init__(self, path: str | Path = DEFAULT_JSON_FILE_PATH) -> None: + super().__init__() + self._path = Path(path).expanduser() + self._load_cache = TTLCache(maxsize=1, ttl=300) + + @cachedmethod(lambda self: self._load_cache, key=lambda self: hashkey()) + def load(self) -> ApiToken: + try: + text = self._path.read_text(encoding="utf-8") + except FileNotFoundError as e: + raise ApiTokenStorageException("JSON file does not exist. Please configure token first.") from e + except OSError as e: + raise ApiTokenStorageException("Failed to read API token JSON file") from e + + try: + return ApiToken.model_validate_json(text) + except ValueError as e: + raise ApiTokenStorageException("Invalid API token JSON file format") from e + + def save(self, api_token: ApiToken) -> None: + self._path.parent.mkdir( + parents=True, + exist_ok=True, + ) + + text = api_token.model_dump_json() + self._atomic_write_text(text) + + self._load_cache.clear() + + def _atomic_write_text(self, text: str) -> None: + descriptor, temp_path = tempfile.mkstemp( + prefix=self._path.name + ".", + dir=self._path.parent, + ) + try: + with os.fdopen( + descriptor, + "w", + encoding="utf-8", + ) as file: + file.write(text) + file.flush() + os.fsync(file.fileno()) + + try: + os.chmod(temp_path, 0o600) + except OSError: + pass + + os.replace(temp_path, self._path) + + finally: + try: + os.remove(temp_path) + except FileNotFoundError: + pass + except OSError: + pass + + +class AsyncJsonFileApiTokenStorage(AsyncApiTokenStorage): + """Asynchronous JSON file storage for API tokens.""" + + _sync: JsonFileApiTokenStorage + + def __init__(self, path: str | Path = DEFAULT_JSON_FILE_PATH) -> None: + super().__init__() + self._sync = JsonFileApiTokenStorage(path) + + async def load(self) -> ApiToken: + return self._sync.load() + + async def save(self, api_token: ApiToken) -> None: + return await asyncio.to_thread(self._sync.save, api_token) diff --git a/rtbhouse_sdk/client.py b/rtbhouse_sdk/client.py index f67b8d4..ee7a5dc 100644 --- a/rtbhouse_sdk/client.py +++ b/rtbhouse_sdk/client.py @@ -3,7 +3,8 @@ # pylint: disable=too-many-arguments import dataclasses import warnings -from collections.abc import AsyncIterable, Generator, Iterable +from abc import ABC, abstractmethod +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Generator, Iterable from datetime import date, timedelta from json import JSONDecodeError from types import TracebackType @@ -28,6 +29,17 @@ MAX_CURSOR_ROWS = 10000 +@dataclasses.dataclass +class ApiTokenAuth: + token: str + + +class DynamicApiTokenAuth(ABC): # pylint: disable=too-few-public-methods + @abstractmethod + def get_token(self) -> str: + pass + + @dataclasses.dataclass class BasicAuth: username: str @@ -61,7 +73,7 @@ class Client: def __init__( self, - auth: BasicAuth | BasicTokenAuth, + auth: ApiTokenAuth | DynamicApiTokenAuth | BasicAuth | BasicTokenAuth, timeout: timedelta = DEFAULT_TIMEOUT, ): self._httpx_client = httpx.Client( @@ -95,12 +107,33 @@ def _get(self, path: str, params: dict[str, Any] | None = None) -> Any: except (ValueError, KeyError) as exc: raise ApiException("Invalid response format") from exc + def _post(self, path: str, data: dict[str, Any] | None = None, params: dict[str, Any] | None = None) -> Any: + response = self._httpx_client.post( + path, + json=data, + params=params, + ) + _validate_response(response) + try: + resp_json = response.json() + return resp_json["data"] + except (ValueError, KeyError) as exc: + raise ApiException("Invalid response format") from exc + def _get_dict(self, path: str, params: dict[str, Any] | None = None) -> dict[str, Any]: data = self._get(path, params) if not isinstance(data, dict): raise ValueError("Result is not a dict") return data + def _post_dict( + self, path: str, data: dict[str, Any] | None = None, params: dict[str, Any] | None = None + ) -> dict[str, Any]: + data = self._post(path, data, params) + if not isinstance(data, dict): + raise ValueError("Result is not a dict") + return data + def _get_list_of_dicts(self, path: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]: data = self._get(path, params) if not isinstance(data, list) or not all(isinstance(item, dict) for item in data): @@ -232,6 +265,20 @@ def get_summary_stats( data = self._get_list_of_dicts(f"/advertisers/{adv_hash}/summary-stats", params) return [schema.Stats(**st) for st in data] + def get_current_api_token(self) -> schema.ApiToken: + data = self._get_dict("/tokens/current") + return schema.ApiToken(**data) + + def rotate_current_api_token(self) -> schema.RotatedApiToken: + data = self._post_dict("/tokens/current/rotate") + return schema.RotatedApiToken(**data) + + +class AsyncDynamicApiTokenAuth(ABC): # pylint: disable=too-few-public-methods + @abstractmethod + async def get_token(self) -> str: + pass + class AsyncClient: """ @@ -247,7 +294,7 @@ class AsyncClient: def __init__( self, - auth: BasicAuth | BasicTokenAuth, + auth: ApiTokenAuth | AsyncDynamicApiTokenAuth | BasicAuth | BasicTokenAuth, timeout: timedelta = DEFAULT_TIMEOUT, ) -> None: self._httpx_client = httpx.AsyncClient( @@ -281,12 +328,33 @@ async def _get(self, path: str, params: dict[str, Any] | None = None) -> Any: except (ValueError, KeyError) as exc: raise ApiException("Invalid response format") from exc + async def _post(self, path: str, data: dict[str, Any] | None = None, params: dict[str, Any] | None = None) -> Any: + response = await self._httpx_client.post( + path, + json=data, + params=params, + ) + _validate_response(response) + try: + resp_json = response.json() + return resp_json["data"] + except (ValueError, KeyError) as exc: + raise ApiException("Invalid response format") from exc + async def _get_dict(self, path: str, params: dict[str, Any] | None = None) -> dict[str, Any]: data = await self._get(path, params) if not isinstance(data, dict): raise ValueError("Result is not a dict") return data + async def _post_dict( + self, path: str, data: dict[str, Any] | None = None, params: dict[str, Any] | None = None + ) -> dict[str, Any]: + data = await self._post(path, data, params) + if not isinstance(data, dict): + raise ValueError("Result is not a dict") + return data + async def _get_list_of_dicts(self, path: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]: data = await self._get(path, params) if not isinstance(data, list) or not all(isinstance(item, dict) for item in data): @@ -421,6 +489,58 @@ async def get_summary_stats( data = await self._get_list_of_dicts(f"/advertisers/{adv_hash}/summary-stats", params) return [schema.Stats(**st) for st in data] + async def get_current_api_token(self) -> schema.ApiToken: + data = await self._get_dict("/tokens/current") + return schema.ApiToken(**data) + + async def rotate_current_api_token(self) -> schema.RotatedApiToken: + data = await self._post_dict("/tokens/current/rotate") + return schema.RotatedApiToken(**data) + + +class _HttpxApiTokenAuth(httpx.Auth): + """API token auth backend.""" + + def __init__(self, token: str) -> None: + self._token = token + + def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + request.headers["Authorization"] = f"Bearer {self._token}" + yield request + + +class _HttpxDynamicApiTokenAuth(httpx.Auth): + """API token auth backend.""" + + def __init__(self, token_provider: Callable[[], str]) -> None: + self._token_provider = token_provider + + def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + token = self._token_provider() + request.headers["Authorization"] = f"Bearer {token}" + yield request + + async def async_auth_flow(self, _: httpx.Request) -> AsyncGenerator[ # type: ignore[override] + httpx.Request, + httpx.Response, + ]: + raise RuntimeError("This auth backend does not support async mode") + + +class _AsyncHttpxDynamicApiTokenAuth(httpx.Auth): + """API token auth backend.""" + + def __init__(self, token_provider: Callable[[], Awaitable[str]]) -> None: + self._token_provider = token_provider + + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + token = await self._token_provider() + request.headers["Authorization"] = f"Bearer {token}" + yield request + + def sync_auth_flow(self, _: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + raise RuntimeError("This auth backend does not support sync mode") + class _HttpxBasicTokenAuth(httpx.Auth): """Basic token auth backend.""" @@ -443,12 +563,22 @@ def _build_headers() -> dict[str, str]: } -def _choose_auth_backend(auth: BasicAuth | BasicTokenAuth) -> httpx.Auth: - if isinstance(auth, BasicAuth): - return httpx.BasicAuth(auth.username, auth.password) - if isinstance(auth, BasicTokenAuth): - return _HttpxBasicTokenAuth(auth.token) - raise ValueError("Unknown auth method") +def _choose_auth_backend( + auth: ApiTokenAuth | DynamicApiTokenAuth | AsyncDynamicApiTokenAuth | BasicAuth | BasicTokenAuth, +) -> httpx.Auth: + match auth: + case ApiTokenAuth(token=token): + return _HttpxApiTokenAuth(token) + case DynamicApiTokenAuth(): + return _HttpxDynamicApiTokenAuth(auth.get_token) + case AsyncDynamicApiTokenAuth(): + return _AsyncHttpxDynamicApiTokenAuth(auth.get_token) + case BasicAuth(username=username, password=password): + return httpx.BasicAuth(username, password) + case BasicTokenAuth(token=token): + return _HttpxBasicTokenAuth(token) + case _: + raise ValueError("Unknown auth method") def _validate_response(response: httpx.Response) -> None: diff --git a/rtbhouse_sdk/schema.py b/rtbhouse_sdk/schema.py index bb4ad42..997f02d 100644 --- a/rtbhouse_sdk/schema.py +++ b/rtbhouse_sdk/schema.py @@ -272,3 +272,13 @@ class Stats(CamelizedBaseModel): cpvisit: float | None = None user_frequency: float | None = None user_reach: float | None = None + + +class RotatedApiToken(CamelizedBaseModel): + token: str + expires_at: datetime + + +class ApiToken(CamelizedBaseModel): + expires_at: datetime + can_rotate: bool