"""Shared AWS session, JSON helpers, and credential resolution for tools/aws/."""
from __future__ import annotations
import jsonutil as json
import logging
import os
from typing import Any, TYPE_CHECKING
from tools.manage_api_keys import missing_api_key_error
logger = logging.getLogger(__name__)
DEFAULT_REGION = "us-east-1"
if TYPE_CHECKING:
pass
async def _aws_user_key(ctx) -> str | None:
"""Resolve the stored AWS credential blob for the current user.
Returns the raw (still JSON-encoded) AWS API-key value that the user stored
via the API-key management system, or ``None`` when the tool context lacks a
Redis client or user id. This is the first step in turning a ``ToolContext``
into a usable boto3 session.
Reads from Redis only indirectly: it delegates to
``tools.manage_api_keys.get_user_api_key``, which looks up the per-user key,
then falls back to the channel pool and the global pool, decrypting the value
when necessary. It passes ``ctx.user_id``, ``ctx.redis``, ``ctx.channel_id``,
and ``ctx.config`` through to that resolver. No value is returned when no
context, no Redis client, or no user id is present.
Called by :func:`_get_session` within this module as the credential-lookup
step; no other internal callers were found.
Args:
ctx: The tool ``ToolContext``. Expected to expose ``redis`` (Redis
client), ``user_id`` (str), and optionally ``channel_id`` and
``config``; missing attributes are tolerated via ``getattr``.
Returns:
str | None: The raw stored AWS credential string (a JSON blob) if found,
otherwise ``None``.
"""
if ctx and getattr(ctx, "redis", None) and getattr(ctx, "user_id", None):
from tools.manage_api_keys import get_user_api_key
return await get_user_api_key(
ctx.user_id, "aws",
redis_client=ctx.redis,
channel_id=getattr(ctx, "channel_id", None),
config=getattr(ctx, "config", None),
)
return None
def _parse_creds(raw: str | None) -> dict | None:
"""Parse a stored credential string into a validated AWS-credentials dict.
Decodes the JSON blob returned by :func:`_aws_user_key` and returns it only
when it contains both ``access_key_id`` and ``secret_access_key`` keys, so
that downstream session construction never sees a malformed or partial
credential set. Any decode failure or wrong type yields ``None`` rather than
propagating.
Interacts only with the local JSON helper (``jsonutil`` imported as
``json``); it has no I/O or other side effects.
Called by :func:`_get_session` to turn the raw stored value into structured
credentials before building a session; no other internal callers were found.
Args:
raw: The raw credential string (expected to be a JSON object), or
``None``.
Returns:
dict | None: The parsed credentials dict when it is valid JSON and
contains both required keys, otherwise ``None``.
"""
if not raw:
return None
try:
c = json.loads(raw)
if "access_key_id" in c and "secret_access_key" in c:
return c
except (json.JSONDecodeError, TypeError):
pass
return None
def _session(creds: dict):
"""Build a boto3 ``Session`` from a parsed user-credentials dict.
Constructs an explicit-credentials boto3 session using the
``access_key_id``/``secret_access_key`` from the user's stored AWS key,
defaulting the region to :data:`DEFAULT_REGION` when the credentials omit
one. boto3 is imported lazily inside the function to keep the dependency off
the import path for callers that never need an authenticated session.
Called by :func:`_get_session` when :func:`_parse_creds` yields a valid
credentials dict; no other internal callers were found.
Args:
creds: A credentials dict with required keys ``access_key_id`` and
``secret_access_key`` and an optional ``region``.
Returns:
boto3.session.Session: A session bound to the supplied credentials and
region.
Raises:
KeyError: If ``access_key_id`` or ``secret_access_key`` is absent
(callers pass dicts already validated by :func:`_parse_creds`).
"""
import boto3
return boto3.Session(
aws_access_key_id=creds["access_key_id"],
aws_secret_access_key=creds["secret_access_key"],
region_name=creds.get("region", DEFAULT_REGION),
)
def _env_session():
"""Build a boto3 ``Session`` from process environment variables, if present.
Provides the host/operator-level fallback used when the user has no stored
AWS credentials: it reads ``AWS_ACCESS_KEY_ID`` and ``AWS_SECRET_ACCESS_KEY``
from the environment and, when both are set, constructs a session whose
region comes from ``AWS_DEFAULT_REGION`` (defaulting to
:data:`DEFAULT_REGION`). boto3 is imported lazily and only when both
variables are populated.
Reads only ``os.environ``; it performs no network or Redis I/O. Called by
:func:`_get_session` as the second credential source after the per-user key;
no other internal callers were found.
Returns:
boto3.session.Session | None: A session built from the environment
credentials, or ``None`` when either variable is empty/unset.
"""
ak = os.environ.get("AWS_ACCESS_KEY_ID", "")
sk = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
if ak and sk:
import boto3
return boto3.Session(
aws_access_key_id=ak,
aws_secret_access_key=sk,
region_name=os.environ.get("AWS_DEFAULT_REGION", DEFAULT_REGION),
)
return None
async def _get_session(ctx) -> Any:
"""Resolve an authenticated boto3 ``Session`` for an AWS tool invocation.
This is the single credential-resolution entry point shared by every AWS
handler. It tries the per-user stored key first, then falls back to process
environment credentials, returning ``None`` when neither source yields
usable credentials so callers can surface a "missing API key" message.
It chains the module's helpers: :func:`_aws_user_key` to fetch the raw
stored value (which reaches Redis and the API-key pools), :func:`_parse_creds`
to validate it, :func:`_session` to build a user-credentialed session, and
:func:`_env_session` for the environment fallback. It has no side effects of
its own beyond those reads.
Called widely across the AWS tool surface: by ``tools/aws/helpers.py``
(``aws_method``), and by the per-service handler modules (``s3``,
``route53``, ``lambda_mod``, ``ec2``). It is also re-exported through
``tools/aws_tools.py`` and consumed by ``tools/object_storage_tools.py``,
``tools/cloud_rag.py``, and ``tools/backup_tools.py``; tests patch
``tools.aws.base._get_session`` to simulate the no-credentials path.
Args:
ctx: The tool ``ToolContext`` carrying ``redis``/``user_id`` and related
attributes used for per-user key lookup.
Returns:
Any: A boto3 ``Session`` (user-key or environment-derived), or ``None``
when no credentials are available.
"""
raw = await _aws_user_key(ctx)
creds = _parse_creds(raw)
if creds:
return _session(creds)
s = _env_session()
if s:
return s
return None
def _json_serial(obj):
"""Serialize boto3 response values that JSON cannot encode natively.
Acts as the ``default`` callback for :func:`_dumps`, converting the
non-JSON-native types boto3 commonly returns: ``datetime``/``date`` become
ISO-8601 strings and ``bytes`` are decoded as UTF-8 with replacement for
invalid sequences. Any other unsupported type triggers a ``TypeError`` so
the encoder fails loudly rather than emitting garbage.
Pure transformation with no side effects. Called by :func:`_dumps` in this
module and, via re-export, by ``tools/oci_tools.py`` which reuses it as its
own ``json.dumps`` default.
Args:
obj: The object the JSON encoder could not serialize.
Returns:
str: An ISO-8601 timestamp for date/datetime inputs, or a decoded string
for bytes inputs.
Raises:
TypeError: If ``obj`` is not a ``datetime``, ``date``, or ``bytes``.
"""
import datetime
if isinstance(obj, (datetime.datetime, datetime.date)):
return obj.isoformat()
if isinstance(obj, bytes):
return obj.decode("utf-8", errors="replace")
raise TypeError(f"Type {type(obj)} not serializable")
def _dumps(data) -> str:
"""Pretty-print an AWS response object as a JSON string for tool output.
Produces the human-readable JSON that AWS tool handlers return to the model:
two-space indented, non-ASCII preserved (``ensure_ascii=False``), and using
:func:`_json_serial` as the fallback encoder so boto3 ``datetime``/``bytes``
values serialize cleanly.
Delegates to ``jsonutil.dumps`` (imported as ``json``) and to
:func:`_json_serial`; no other side effects. Called by ``tools/aws/helpers.py``
(``aws_method``) to format successful responses, and broadly across the cloud
tool surface via the ``tools/aws_tools.py`` re-export, including
``tools/object_storage_tools.py`` and ``tools/oci_tools.py``.
Args:
data: Any JSON-serializable structure (typically a boto3 response dict).
Returns:
str: The indented JSON representation of ``data``.
"""
return json.dumps(data, indent=2, ensure_ascii=False, default=_json_serial)
def _err(msg: str) -> str:
"""Wrap an error message as a compact JSON error payload.
Gives AWS tool handlers a uniform machine-readable error shape,
``{"error": msg}``, returned in place of a result whenever credentials are
missing or a boto3 call raises.
Delegates to ``jsonutil.dumps`` (imported as ``json``); no side effects.
Called by ``tools/aws/helpers.py`` (``aws_method`` for the missing-key and
exception paths) and across the per-service handler modules (``s3``,
``route53``, ``lambda_mod``, ``ec2``), plus the ``tools/aws_tools.py``
re-export consumed by ``tools/object_storage_tools.py``.
Args:
msg: The human-readable error description to embed.
Returns:
str: A JSON string of the form ``{"error": "<msg>"}``.
"""
return json.dumps({"error": msg})
[docs]
def strip_rm(resp: Any) -> Any:
"""Drop the boilerplate ``ResponseMetadata`` block from a boto3 response.
Removes the ``ResponseMetadata`` key (HTTP status, request id, retry counts)
that boto3 attaches to every response, trimming noise before the response is
serialized and shown to the model. The dict is mutated in place and returned;
non-dict responses pass through unchanged.
Pure local mutation with no I/O. Called by ``tools/aws/helpers.py``
(``aws_method``) on dict responses just before :func:`_dumps`; no other
internal callers were found.
Args:
resp: A boto3 response, typically a dict but any type is accepted.
Returns:
Any: The same object, with ``ResponseMetadata`` removed when it was a
dict.
"""
if isinstance(resp, dict):
resp.pop("ResponseMetadata", None)
return resp