dh-cli 0.8.6__tar.gz → 0.8.7__tar.gz

{dh_cli-0.8.6 → dh_cli-0.8.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.8.6
+Version: 0.8.7
 Summary: Dayhoff Labs developer CLI
 Author-email: Dayhoff Labs <dev@dayhofflabs.com>
 License: # PolyForm Noncommercial License 1.0.0

{dh_cli-0.8.6 → dh_cli-0.8.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dh-cli"
-version = "0.8.6"
+version = "0.8.7"
 description = "Dayhoff Labs developer CLI"
 requires-python = ">=3.11"
 readme = "README.md"

{dh_cli-0.8.6 → dh_cli-0.8.7}/src/dh_cli/bedrock/commands.py

@@ -22,6 +22,8 @@ from typing import Optional
 
 import click
 
+from dh_cli._identity import HandleResolutionError, resolve_handle_from_session
+
 from . import cost_report as cr
 
 # Hard-coded for the dev account. Kept here (not as user-facing flags)
@@ -58,37 +60,21 @@ def bedrock():
 # --------------------------------------------------------------------------
 
 
-def _resolve_handle_from_sts() -> str:
-    """Pull the developer handle from the current SSO session.
-
-    A DeveloperAccess SSO caller identity looks like::
+def _resolve_handle() -> str:
+    """Resolve the developer handle from the current SSO session.
 
-        arn:aws:sts::074735440724:assumed-role/AWSReservedSSO_DeveloperAccess_<id>/<handle>
-
-    We return `<handle>`. Raises ClickException if STS fails or the ARN
-    shape isn't recognisable.
+    Thin wrapper around `_identity.resolve_handle_from_session` that
+    builds a boto3 Session and turns `HandleResolutionError` into a
+    `ClickException` so Click prints it cleanly. Shared with `dh gh`
+    via `_identity.py` so both commands behave identically when the
+    caller's session can't be classified.
     """
     import boto3
 
     try:
-        sts = boto3.client("sts")
-        arn = sts.get_caller_identity()["Arn"]
-    except Exception as exc:
-        raise click.ClickException(
-            f"Could not call sts:GetCallerIdentity to resolve your handle: {exc}\n"
-            "Run `awslogin dev-devaccess` (or pass --handle explicitly)."
-        )
-
-    principal = cr.classify_arn(arn)
-    # Accept any AWSReservedSSO_* session: DeveloperAccess,
-    # DevAdminAccess, SecurityAdministrator, etc. all carry the
-    # IdC username as the role-session-name, which is what we want
-    # as the developer handle.
-    if principal.principal_type in ("claude-code", "claude-code-other", "cursor"):
-        return principal.principal_name
-    raise click.ClickException(
-        f"Couldn't infer a developer handle from your identity ({arn}). Pass --handle explicitly."
-    )
+        return resolve_handle_from_session(boto3.Session())
+    except HandleResolutionError as exc:
+        raise click.ClickException(str(exc))
 
 
 # --------------------------------------------------------------------------
@@ -135,7 +121,7 @@ def bedrock_key(handle: Optional[str], region: str, mode: str):
     """
     import boto3
 
-    resolved = handle or _resolve_handle_from_sts()
+    resolved = handle or _resolve_handle()
     secret_id = f"{_SECRET_PREFIX}{resolved}"
 
     try:
@@ -316,7 +302,7 @@ def bedrock_cost(
 
     my_handle: Optional[str] = None
     if me:
-        my_handle = _resolve_handle_from_sts()
+        my_handle = _resolve_handle()
 
     try:
         records = cr.walk_logs(

{dh_cli-0.8.6 → dh_cli-0.8.7}/tests/bedrock/test_key_command.py

@@ -32,15 +32,32 @@ def _build_sts_stub(arn: str = _DEV_ARN) -> MagicMock:
     return sts
 
 
+def _build_session_stub(sts: MagicMock) -> MagicMock:
+    """boto3.Session() stub whose .client('sts') returns `sts`.
+
+    Handle resolution flows through `_identity.resolve_handle_from_session`,
+    which calls `session.client('sts')`. The Secrets Manager fetch in
+    `bedrock_key` still goes through `boto3.client('secretsmanager', ...)`
+    directly, which is patched separately.
+    """
+    session = MagicMock()
+    session.client.return_value = sts
+    return session
+
+
 def _patch_boto3(*, sts=None, sm=None):
+    sts_stub = sts or _build_sts_stub()
+    session_stub = _build_session_stub(sts_stub)
+
     def _client(name, **_kwargs):
-        if name == "sts":
-            return sts or _build_sts_stub()
         if name == "secretsmanager":
             return sm or _build_sm_stub()
-        raise AssertionError(f"unexpected boto3 client: {name}")
+        raise AssertionError(f"unexpected boto3.client call: {name}")
 
-    return patch("boto3.client", side_effect=_client)
+    return (
+        patch("boto3.Session", return_value=session_stub),
+        patch("boto3.client", side_effect=_client),
+    )
 
 
 def _invoke(args):
@@ -52,7 +69,8 @@ def _invoke(args):
 
 def test_key_default_prints_json_and_resolves_handle_from_sts():
     sm = _build_sm_stub()
-    with _patch_boto3(sts=_build_sts_stub(), sm=sm):
+    p_sess, p_cli = _patch_boto3(sts=_build_sts_stub(), sm=sm)
+    with p_sess, p_cli:
         result = _invoke([])
     assert result.exit_code == 0, result.output
     sm.get_secret_value.assert_called_once_with(SecretId="cursor-bedrock/dma")
@@ -66,26 +84,28 @@ def test_key_handle_override_skips_sts():
     can still fetch (for e.g. break-glass debugging)."""
     sm = _build_sm_stub()
 
-    called: dict[str, int] = {"sts": 0}
+    called: dict[str, int] = {"session": 0}
+
+    def _session_factory(*_args, **_kwargs):
+        called["session"] += 1
+        return _build_session_stub(_build_sts_stub("arn:aws:iam::074735440724:root"))
 
     def _client(name, **_kwargs):
-        if name == "sts":
-            called["sts"] += 1
-            return _build_sts_stub("arn:aws:iam::074735440724:root")
         if name == "secretsmanager":
             return sm
         raise AssertionError(name)
 
-    with patch("boto3.client", side_effect=_client):
+    with patch("boto3.Session", side_effect=_session_factory), patch("boto3.client", side_effect=_client):
         result = _invoke(["--handle", "jason"])
     assert result.exit_code == 0, result.output
-    assert called["sts"] == 0
+    assert called["session"] == 0
     sm.get_secret_value.assert_called_once_with(SecretId="cursor-bedrock/jason")
 
 
 def test_key_export_mode_prints_shell_exports():
     sm = _build_sm_stub()
-    with _patch_boto3(sm=sm):
+    p_sess, p_cli = _patch_boto3(sm=sm)
+    with p_sess, p_cli:
         result = _invoke(["--export"])
     assert result.exit_code == 0, result.output
     out = result.output
@@ -96,7 +116,8 @@ def test_key_export_mode_prints_shell_exports():
 
 def test_key_secret_fetch_failure_exits_non_zero():
     sm = _build_sm_stub(RuntimeError("AccessDeniedException: secretsmanager:GetSecretValue"))
-    with _patch_boto3(sm=sm):
+    p_sess, p_cli = _patch_boto3(sm=sm)
+    with p_sess, p_cli:
         result = _invoke(["--handle", "nobody"])
     assert result.exit_code != 0
     assert "Could not read secret" in result.output
@@ -104,24 +125,30 @@ def test_key_secret_fetch_failure_exits_non_zero():
 
 def test_key_unrecognised_caller_arn_exits_with_hint():
     """If STS returns something we can't classify, the error message must
-    tell the user to pass --handle explicitly instead of silently guessing."""
+    tell the user to pass --handle explicitly instead of silently guessing.
+
+    Message comes from `_identity.HandleResolutionError` (shared with
+    `dh gh login`), so both commands surface the same hint."""
     sts = _build_sts_stub("arn:aws:iam::074735440724:federated-user/weird")
-    with _patch_boto3(sts=sts):
+    p_sess, p_cli = _patch_boto3(sts=sts)
+    with p_sess, p_cli:
         result = _invoke([])
     assert result.exit_code != 0
     assert "--handle" in result.output
+    assert "awslogin" in result.output
 
 
 def test_key_resolves_handle_from_dev_admin_access_session():
     """DevAdminAccess (and any other AWSReservedSSO_* role) must resolve
     to the session-name as the developer handle. Regression: earlier the
     gate only accepted DeveloperAccess sessions, which broke `dh bedrock
-    key` for anyone running under an admin profile."""
-    sts = _build_sts_stub(
-        "arn:aws:sts::074735440724:assumed-role/AWSReservedSSO_DevAdminAccess_2506d659c73583cc/dma"
-    )
+    key` for anyone running under an admin profile. The shared
+    `_identity` helper uses a regex that matches any `AWSReservedSSO_*`
+    role by construction."""
+    sts = _build_sts_stub("arn:aws:sts::074735440724:assumed-role/AWSReservedSSO_DevAdminAccess_2506d659c73583cc/dma")
     sm = _build_sm_stub()
-    with _patch_boto3(sts=sts, sm=sm):
+    p_sess, p_cli = _patch_boto3(sts=sts, sm=sm)
+    with p_sess, p_cli:
         result = _invoke([])
     assert result.exit_code == 0, result.output
     sm.get_secret_value.assert_called_once_with(SecretId="cursor-bedrock/dma")