grantry 0.1.0__py3-none-any.whl

grantry/__init__.py

@@ -0,0 +1,3 @@
+"""grantry: a local credential broker for humans and AI agents."""
+
+__version__ = "0.1.0"

grantry/__main__.py

@@ -0,0 +1,6 @@
+"""Enable `python -m grantry`, which the MCP config points at so it works
+regardless of PATH."""
+
+from grantry.cli import main
+
+raise SystemExit(main())

grantry/admin.py

@@ -0,0 +1,125 @@
+"""Org-wide assignment crawl: who has which permission set in which account.
+
+This is admin-only by construction. It calls the sso-admin, organizations, and
+identitystore APIs, which require signed credentials from an account that holds
+those permissions (typically the management or a delegated-admin account).
+grantry mints those credentials through the normal policy path, so a caller who
+cannot assume such a role simply gets nothing; AWS is the gatekeeper.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import Any
+
+ClientFactory = Callable[[str], Any]
+
+
+@dataclass(frozen=True)
+class Assignment:
+    principal_type: str
+    principal_id: str
+    principal_name: str
+    permission_set_name: str
+    account_id: str
+    account_name: str
+
+
+def _paginate(op: Callable[..., dict[str, Any]], key: str, **kwargs: Any) -> Iterator[Any]:
+    token: str | None = None
+    while True:
+        call = dict(kwargs)
+        if token:
+            call["NextToken"] = token
+        resp = op(**call)
+        yield from resp.get(key, [])
+        token = resp.get("NextToken")
+        if not token:
+            return
+
+
+def _principal_name(
+    idstore: Any,
+    identity_store_id: str,
+    ptype: str,
+    pid: str,
+    cache: dict[tuple[str, str], str],
+) -> str:
+    # A group or user appears on many assignments; resolve each unique principal
+    # ONCE and cache it. At 10k+ assignments this turns tens of thousands of
+    # describe calls into one per distinct principal, which is the difference
+    # between a crawl that finishes and one that gets throttled.
+    key = (ptype, pid)
+    if key in cache:
+        return cache[key]
+    try:
+        if ptype == "GROUP":
+            resp = idstore.describe_group(IdentityStoreId=identity_store_id, GroupId=pid)
+            name = str(resp.get("DisplayName", pid))
+        else:
+            resp = idstore.describe_user(IdentityStoreId=identity_store_id, UserId=pid)
+            name = str(resp.get("UserName", pid))
+    except Exception:
+        name = pid
+    cache[key] = name
+    return name
+
+
+def crawl_assignments(
+    make_client: ClientFactory,
+    on_progress: Callable[[int, int], None] | None = None,
+) -> list[Assignment]:
+    sso = make_client("sso-admin")
+    idstore = make_client("identitystore")
+    orgs = make_client("organizations")
+
+    instances = sso.list_instances().get("Instances", [])
+    if not instances:
+        return []
+    instance_arn = instances[0]["InstanceArn"]
+    identity_store_id = instances[0]["IdentityStoreId"]
+
+    accounts: dict[str, str] = {}
+    for acct in _paginate(orgs.list_accounts, "Accounts"):
+        accounts[acct["Id"]] = acct.get("Name", acct["Id"])
+
+    ps_names: dict[str, str] = {}
+    for ps_arn in _paginate(sso.list_permission_sets, "PermissionSets", InstanceArn=instance_arn):
+        d = sso.describe_permission_set(InstanceArn=instance_arn, PermissionSetArn=ps_arn)
+        ps_names[ps_arn] = d["PermissionSet"]["Name"]
+
+    out: list[Assignment] = []
+    name_cache: dict[tuple[str, str], str] = {}
+    total_accounts = len(accounts)
+    for done, (acct_id, acct_name) in enumerate(accounts.items(), start=1):
+        for ps_arn in _paginate(
+            sso.list_permission_sets_provisioned_to_account,
+            "PermissionSets",
+            InstanceArn=instance_arn,
+            AccountId=acct_id,
+        ):
+            for a in _paginate(
+                sso.list_account_assignments,
+                "AccountAssignments",
+                InstanceArn=instance_arn,
+                AccountId=acct_id,
+                PermissionSetArn=ps_arn,
+            ):
+                ptype = a["PrincipalType"]
+                pid = a["PrincipalId"]
+                out.append(
+                    Assignment(
+                        principal_type=ptype,
+                        principal_id=pid,
+                        principal_name=_principal_name(
+                            idstore, identity_store_id, ptype, pid, name_cache
+                        ),
+                        permission_set_name=ps_names.get(ps_arn, ps_arn),
+                        account_id=acct_id,
+                        account_name=acct_name,
+                    )
+                )
+        if on_progress:
+            on_progress(done, total_accounts)
+    return out