imbi-plugin-aws 0.1.0__py3-none-any.whl

imbi_plugin_aws/__init__.py

@@ -0,0 +1,21 @@
+"""AWS plugins for Imbi.
+
+Ships three plugin entry points from a single distribution:
+
+* ``aws-iam-ic`` -- :class:`IdentityPlugin` for AWS IAM Identity Center.
+* ``aws-ssm`` -- :class:`ConfigurationPlugin` backed by SSM Parameter
+  Store.
+* ``aws-cloudwatch-logs`` -- :class:`LogsPlugin` backed by CloudWatch
+  Logs Insights.
+
+The data plugins consume credentials from a flat ``dict[str, str]``
+populated by either the ``ServiceApplication`` static-key path or the
+identity-hydrated path that calls ``aws-iam-ic.materialize()``;
+neither plugin needs to know which source applied.
+"""
+
+from imbi_plugin_aws.cloudwatch import CloudWatchLogsPlugin
+from imbi_plugin_aws.identity import AwsIamIcPlugin
+from imbi_plugin_aws.ssm import SsmPlugin
+
+__all__ = ['AwsIamIcPlugin', 'CloudWatchLogsPlugin', 'SsmPlugin']

imbi_plugin_aws/_helpers.py

@@ -0,0 +1,31 @@
+"""Shared utilities for the AWS data plugins."""
+
+from __future__ import annotations
+
+from imbi_common.plugins.base import PluginContext
+
+
+def template_vars(ctx: PluginContext) -> dict[str, str | None]:
+    """Whitelisted template variables for plugin option expansion."""
+    return {
+        'project_slug': ctx.project_slug,
+        'org_slug': ctx.org_slug,
+        'team_slug': ctx.team_slug,
+        'environment': ctx.environment,
+        'project_id': ctx.project_id,
+    }
+
+
+def assignment_region(ctx: PluginContext) -> str | None:
+    region = ctx.assignment_options.get('region')
+    return str(region) if region else None
+
+
+def assignment_timeout(
+    ctx: PluginContext, *, default: float, key: str = 'timeout_seconds'
+) -> float:
+    raw = ctx.assignment_options.get(key)
+    try:
+        return float(raw) if raw is not None else default
+    except (TypeError, ValueError):
+        return default

imbi_plugin_aws/account_resolution.py

@@ -0,0 +1,138 @@
+"""``MAPS_TO`` traversal helper for the AWS plugin.
+
+At call time, ``materialize`` (or any sibling plugin's handler) needs to
+pick a single :class:`imbi_plugin_aws.models.AwsAccount` from the
+actor's context.  Operators model the mapping declaratively as
+``(:Environment | :Project | :ProjectType | :Organization)-[:MAPS_TO]->
+(:AwsAccount)`` edges; this module walks them in selector order and
+returns the first match.
+"""
+
+from __future__ import annotations
+
+import logging
+import typing
+
+from imbi_common.plugins.base import PluginContext
+
+from imbi_plugin_aws.errors import AccountNotResolvedError
+from imbi_plugin_aws.models import AwsAccount
+
+LOGGER = logging.getLogger(__name__)
+
+DEFAULT_SELECTOR: list[str] = [
+    'project',
+    'environment',
+    'project_type',
+    'organization',
+]
+
+_LABEL_FOR_ANCHOR: dict[str, str] = {
+    'project': 'Project',
+    'environment': 'Environment',
+    'project_type': 'ProjectType',
+    'organization': 'Organization',
+}
+
+
+async def resolve_account(
+    db: typing.Any,
+    ctx: PluginContext,
+    options: dict[str, typing.Any],
+) -> AwsAccount:
+    """Return the first ``AwsAccount`` reachable via ``MAPS_TO`` from
+    the actor's context anchors, in selector order.
+
+    Tag filters narrow the candidate set per anchor: each candidate's
+    ``tags`` dict must be a superset of every entry in
+    ``options['tag_filters']``.
+
+    Raises:
+        AccountNotResolvedError: No anchor produced a matching account.
+    """
+    selector = list(options.get('account_selector', DEFAULT_SELECTOR))
+    tag_filters: dict[str, str] = options.get('tag_filters', {}) or {}
+
+    anchors_checked: list[str] = []
+    for anchor in selector:
+        label = _LABEL_FOR_ANCHOR.get(anchor)
+        if label is None:
+            continue
+        anchor_id = _anchor_id_from_ctx(anchor, ctx)
+        if not anchor_id:
+            continue
+        anchors_checked.append(anchor)
+        candidates = await _candidates_for_anchor(db, label, anchor_id)
+        for candidate in candidates:
+            if _matches_tag_filters(candidate.tags, tag_filters):
+                return candidate
+
+    raise AccountNotResolvedError(
+        selector=selector,
+        anchors_checked=anchors_checked,
+    )
+
+
+def _anchor_id_from_ctx(anchor: str, ctx: PluginContext) -> str | None:
+    """Look up the id of ``anchor`` from the request's
+    :class:`PluginContext`.
+
+    The host populates ``project_id`` / ``project_slug`` / ``org_slug``
+    today; ``environment`` is a slug, not an id, and ``project_type`` is
+    not yet on the context — those anchors are no-ops in Phase 1 and
+    will be reachable once the host extends ``PluginContext``.
+    """
+    if anchor == 'project':
+        return ctx.project_id or None
+    # Environment / project_type / organization arrive as slugs from
+    # the host, not ids; the calling plugin must perform the slug→id
+    # lookup before passing them in.  Phase 1 only resolves via
+    # project_id; the other anchors return None.
+    return None
+
+
+async def _candidates_for_anchor(
+    db: typing.Any,
+    label: str,
+    anchor_id: str,
+) -> list[AwsAccount]:
+    """Walk ``MAPS_TO`` from a labeled node id to its ``AwsAccount``
+    targets.
+
+    Uses the same Cypher template the host uses elsewhere — ``{}``
+    placeholders are SQL-format-style, ``{{`` / ``}}`` escape Cypher
+    map literals.
+    """
+    query = (
+        'MATCH (n:' + label + ' {{id: {anchor_id}}})'
+        '-[:MAPS_TO]->(a:AwsAccount)'
+        ' RETURN a{{.*}} AS account'
+    )
+    rows = await db.execute(query, {'anchor_id': anchor_id}, ['account'])
+    accounts: list[AwsAccount] = []
+    for row in rows:
+        raw = row.get('account')
+        if raw is None:
+            continue
+        # ``parse_agtype`` is the host's helper.  Imported lazily to
+        # avoid coupling this module to imbi_common's import order
+        # in tests that stub the db.
+        try:
+            from imbi_common.graph import parse_agtype
+        except ImportError:
+            data = raw
+        else:
+            data = parse_agtype(raw)
+        try:
+            accounts.append(AwsAccount.model_validate(data))
+        except Exception:  # noqa: BLE001
+            LOGGER.warning(
+                'AwsAccount row %r failed validation', data, exc_info=True
+            )
+    return accounts
+
+
+def _matches_tag_filters(
+    tags: dict[str, str], tag_filters: dict[str, str]
+) -> bool:
+    return all(tags.get(k) == v for k, v in tag_filters.items())

imbi_plugin_aws/aws_session.py

@@ -0,0 +1,324 @@
+"""Shared AWS credentials + JSON-protocol client for the SSM and
+CloudWatch Logs plugins.
+
+We deliberately stay on ``httpx`` (already used by the IAM IC identity
+plugin) and implement a minimal AWS SigV4 signer so the test suite can
+keep using ``respx`` for mocking.  Both SSM and CloudWatch Logs speak
+the JSON-1.1 wire protocol, so a single helper covers both surfaces.
+"""
+
+from __future__ import annotations
+
+import datetime
+import hashlib
+import hmac
+import json
+import typing
+import urllib.parse
+
+import httpx
+from imbi_common.plugins.base import PluginContext
+from imbi_common.plugins.errors import (
+    PluginAuthenticationFailed,
+    PluginCredentialsMissing,
+    PluginTimeoutError,
+    PluginUnavailableError,
+)
+
+ServiceName = typing.Literal['ssm', 'logs']
+
+_SIGNING_ALGORITHM = 'AWS4-HMAC-SHA256'
+_AWS_REQUEST = 'aws4_request'
+
+_TARGET_PREFIX: dict[ServiceName, str] = {
+    'ssm': 'AmazonSSM',
+    'logs': 'Logs_20140328',
+}
+
+# AWS error codes that indicate the caller's credentials are no longer
+# accepted by the upstream service.  Refreshing the IAM IC session and
+# retrying the call once is a reasonable recovery, so the host's retry
+# layer is given a chance via :class:`PluginAuthenticationFailed`.
+_AUTH_ERROR_CODES: frozenset[str] = frozenset(
+    {
+        'ExpiredToken',
+        'ExpiredTokenException',
+        'InvalidClientTokenId',
+        'UnauthorizedException',
+        'UnrecognizedClientException',
+    }
+)
+
+
+class AwsCredentials(typing.NamedTuple):
+    """A validated set of AWS credentials + region."""
+
+    access_key_id: str
+    secret_access_key: str
+    session_token: str | None
+    region: str
+
+
+_AWS_IDENTITY_KEYS = (
+    'aws_access_key_id',
+    'aws_secret_access_key',
+    'aws_session_token',
+    'aws_region',
+    'aws_account_id',
+)
+
+
+def resolve_credentials(
+    credentials: dict[str, str],
+    *,
+    region: str | None = None,
+    ctx: PluginContext | None = None,
+) -> AwsCredentials:
+    """Validate the credentials dict and resolve the effective region.
+
+    Two credential sources are supported and combined here:
+
+    * ``credentials`` — static keys stored on the ``Plugin`` node by
+      the operator (``api_token``-style auth).
+    * ``ctx.identity.extra`` — STS keys minted by an identity plugin
+      (currently ``aws-iam-ic``).  The host attaches these to
+      :attr:`PluginContext.identity` via ``hydrate_identity`` and the
+      data plugin must consult them, since the host does not merge
+      identity creds into the ``credentials`` dict.
+
+    Identity-supplied keys take precedence when present (a connected
+    user's STS keys are always preferable to static fallbacks).
+    ``region`` argument wins over both — it's the assignment option.
+
+    Raises:
+        PluginCredentialsMissing: when either AWS key is missing or
+            when the region cannot be resolved.
+    """
+    merged: dict[str, str] = {k: v for k, v in credentials.items() if v}
+    if ctx is not None and ctx.identity is not None:
+        for key in _AWS_IDENTITY_KEYS:
+            value = ctx.identity.extra.get(key)
+            if isinstance(value, str) and value:
+                merged[key] = value
+
+    access_key = merged.get('aws_access_key_id') or ''
+    secret_key = merged.get('aws_secret_access_key') or ''
+    if not access_key and not secret_key:
+        raise PluginCredentialsMissing(
+            'AWS credentials missing: aws_access_key_id and '
+            'aws_secret_access_key are required'
+        )
+    if not access_key or not secret_key:
+        raise PluginCredentialsMissing(
+            'AWS credentials malformed: both aws_access_key_id and '
+            'aws_secret_access_key must be provided'
+        )
+    resolved_region = region or merged.get('aws_region')
+    if not resolved_region:
+        raise PluginCredentialsMissing(
+            'AWS region missing: pass region explicitly or set '
+            'aws_region in credentials'
+        )
+    session_token = merged.get('aws_session_token') or None
+    return AwsCredentials(
+        access_key_id=access_key,
+        secret_access_key=secret_key,
+        session_token=session_token,
+        region=resolved_region,
+    )
+
+
+def _hmac(key: bytes, data: str) -> bytes:
+    return hmac.new(key, data.encode('utf-8'), hashlib.sha256).digest()
+
+
+def _signing_key(
+    secret_access_key: str, date_stamp: str, region: str, service: str
+) -> bytes:
+    k_date = _hmac(f'AWS4{secret_access_key}'.encode(), date_stamp)
+    k_region = _hmac(k_date, region)
+    k_service = _hmac(k_region, service)
+    return _hmac(k_service, _AWS_REQUEST)
+
+
+def sign_request(
+    *,
+    method: str,
+    url: str,
+    headers: dict[str, str],
+    body: bytes,
+    service: str,
+    credentials: AwsCredentials,
+    now: datetime.datetime | None = None,
+) -> dict[str, str]:
+    """Return signed request headers (SigV4) for an AWS JSON request.
+
+    The returned mapping is the *complete* set of headers to send: the
+    caller's ``headers`` plus ``Host``, ``X-Amz-Date``,
+    ``X-Amz-Security-Token`` (when applicable), and ``Authorization``.
+    """
+    parsed = urllib.parse.urlsplit(url)
+    host = parsed.netloc
+    canonical_uri = parsed.path or '/'
+    canonical_query = parsed.query
+    timestamp = (now or datetime.datetime.now(datetime.UTC)).strftime(
+        '%Y%m%dT%H%M%SZ'
+    )
+    date_stamp = timestamp[:8]
+    payload_hash = hashlib.sha256(body).hexdigest()
+
+    signed: dict[str, str] = {**headers, 'host': host, 'x-amz-date': timestamp}
+    if credentials.session_token:
+        signed['x-amz-security-token'] = credentials.session_token
+
+    canonical_headers_pairs = sorted(
+        (name.lower(), value.strip()) for name, value in signed.items()
+    )
+    canonical_headers = ''.join(
+        f'{name}:{value}\n' for name, value in canonical_headers_pairs
+    )
+    signed_header_names = ';'.join(name for name, _ in canonical_headers_pairs)
+    canonical_request = '\n'.join(
+        [
+            method.upper(),
+            canonical_uri,
+            canonical_query,
+            canonical_headers,
+            signed_header_names,
+            payload_hash,
+        ]
+    )
+
+    credential_scope = (
+        f'{date_stamp}/{credentials.region}/{service}/{_AWS_REQUEST}'
+    )
+    string_to_sign = '\n'.join(
+        [
+            _SIGNING_ALGORITHM,
+            timestamp,
+            credential_scope,
+            hashlib.sha256(canonical_request.encode()).hexdigest(),
+        ]
+    )
+
+    signature = hmac.new(
+        _signing_key(
+            credentials.secret_access_key,
+            date_stamp,
+            credentials.region,
+            service,
+        ),
+        string_to_sign.encode(),
+        hashlib.sha256,
+    ).hexdigest()
+    authorization = (
+        f'{_SIGNING_ALGORITHM} '
+        f'Credential={credentials.access_key_id}/{credential_scope}, '
+        f'SignedHeaders={signed_header_names}, '
+        f'Signature={signature}'
+    )
+    signed['Authorization'] = authorization
+    return signed
+
+
+def _service_endpoint(service: ServiceName, region: str) -> str:
+    return f'https://{service}.{region}.amazonaws.com/'
+
+
+def _extract_error_code(payload: dict[str, typing.Any]) -> str:
+    error_type = payload.get('__type') or payload.get('Code') or ''
+    if isinstance(error_type, str) and '#' in error_type:
+        error_type = error_type.split('#', 1)[1]
+    return str(error_type or '')
+
+
+def _map_status_error(
+    status_code: int,
+    error_code: str,
+    message: str,
+    error_map: dict[str, type[Exception]],
+) -> Exception:
+    """Translate an AWS error into the appropriate Imbi exception.
+
+    Auth-failure error codes take precedence over the caller-supplied
+    ``error_map`` so the host retry layer can refresh the actor's IAM
+    IC connection — historical maps surface these as
+    :class:`PluginCredentialsMissing`, which means "configure new
+    static creds" and would terminate the request before the refresh
+    has a chance to fire.
+    """
+    if error_code in _AUTH_ERROR_CODES:
+        return PluginAuthenticationFailed(message or error_code)
+    if error_code in error_map:
+        return error_map[error_code](message or error_code)
+    if status_code >= 500:
+        return PluginUnavailableError(message or f'AWS {status_code}')
+    return PluginUnavailableError(
+        message or f'AWS error {error_code or status_code}'
+    )
+
+
+async def call_aws_json(
+    *,
+    service: ServiceName,
+    action: str,
+    body: dict[str, typing.Any],
+    credentials: AwsCredentials,
+    error_map: dict[str, type[Exception]],
+    timeout: float = 15.0,
+    client: httpx.AsyncClient | None = None,
+) -> dict[str, typing.Any]:
+    """Invoke an AWS JSON-1.1 action and return the decoded response.
+
+    ``error_map`` maps AWS error codes (e.g. ``ParameterNotFound``) to
+    the Imbi exception class to raise.  Errors not in the map fall back
+    to :class:`PluginUnavailableError` for 5xx and a default of
+    :class:`PluginUnavailableError` otherwise.
+    """
+    payload = json.dumps(body).encode()
+    target = f'{_TARGET_PREFIX[service]}.{action}'
+    base_headers: dict[str, str] = {
+        'Content-Type': 'application/x-amz-json-1.1',
+        'X-Amz-Target': target,
+    }
+    url = _service_endpoint(service, credentials.region)
+    signed_headers = sign_request(
+        method='POST',
+        url=url,
+        headers=base_headers,
+        body=payload,
+        service=service,
+        credentials=credentials,
+    )
+
+    async def _do(http: httpx.AsyncClient) -> httpx.Response:
+        return await http.post(url, content=payload, headers=signed_headers)
+
+    try:
+        if client is None:
+            async with httpx.AsyncClient(timeout=timeout) as http:
+                response = await _do(http)
+        else:
+            response = await _do(client)
+    except (httpx.ConnectTimeout, httpx.ReadTimeout) as exc:
+        raise PluginTimeoutError(str(exc)) from exc
+
+    if response.status_code == 200:
+        if not response.content:
+            return {}
+        return typing.cast(dict[str, typing.Any], response.json())
+
+    try:
+        err_payload = typing.cast(dict[str, typing.Any], response.json())
+    except ValueError:
+        err_payload = {}
+    error_code = _extract_error_code(err_payload)
+    error_message = str(
+        err_payload.get('message')
+        or err_payload.get('Message')
+        or response.text
+        or ''
+    )
+    raise _map_status_error(
+        response.status_code, error_code, error_message, error_map
+    )