agentcore-cli 0.31.1__py3-none-any.whl

agentcore/primitives/policy_engine.py

@@ -0,0 +1,358 @@
+"""Policy engine add/remove (port of src/cli/primitives/PolicyEnginePrimitive.ts)."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import typer
+
+from agentcore.cli import interactive
+from agentcore.cli.commands.add._common import first_issue
+from agentcore.cli.errors import get_error_message
+from agentcore.cli.naming import generate_unique_name
+from agentcore.cli.run_command import run_cli_command
+from agentcore.lib import ResourceNotFoundError, find_config_root, serialize_result
+from agentcore.lib.errors.types import to_error
+from agentcore.lib.result import Err, Ok
+from agentcore.lib.schemas.io.config_io import to_jsonable
+from agentcore.primitives._policy_form import POLICY_ENGINE_MODE_OPTIONS
+from agentcore.primitives.base import (
+    SOURCE_CODE_NOTE,
+    BasePrimitive,
+    RemovableResource,
+    RemovalPreview,
+    SchemaChange,
+    emit_json,
+)
+from agentcore.primitives.deployed_gateways import merge_deployed_gateways
+from agentcore.schema import PolicyEngineName
+
+_POLICY_ENGINE_MODES = ("LOG_ONLY", "ENFORCE")
+
+
+@dataclass
+class AddPolicyEngineOptions:
+    name: str
+    description: str | None = None
+    encryption_key_arn: str | None = None
+
+
+def parse_policy_engine_mode(value: str) -> str:
+    """``PolicyEngineModeSchema.parse`` equivalent (Zod enum error message)."""
+    if value not in _POLICY_ENGINE_MODES:
+        expected = "|".join(f'"{m}"' for m in _POLICY_ENGINE_MODES)
+        raise ValueError(f"Invalid option: expected one of {expected}")
+    return value
+
+
+def split_csv(value: str) -> list[str]:
+    return [s.strip() for s in value.split(",") if s.strip()]
+
+
+class PolicyEnginePrimitive(BasePrimitive):
+    kind = "policy-engine"
+    label = "Policy Engine"
+
+    def _project(self) -> dict[str, Any]:
+        return to_jsonable(self.read_project_spec())
+
+    # --- add / remove -------------------------------------------------------------------
+
+    def add(self, options: AddPolicyEngineOptions) -> Ok[dict[str, Any]] | Err[Exception]:
+        try:
+            project = self._project()
+            self.check_duplicate([RemovableResource(e["name"]) for e in project["policyEngines"]], options.name)
+            engine: dict[str, Any] = {"name": options.name}
+            if options.description:
+                engine["description"] = options.description
+            if options.encryption_key_arn:
+                engine["encryptionKeyArn"] = options.encryption_key_arn
+            engine["policies"] = []
+            project["policyEngines"].append(engine)
+            self.write_project_spec(project)
+            return Ok({"engineName": engine["name"]})
+        except Exception as exc:  # noqa: BLE001
+            return Err(to_error(exc))
+
+    def remove(self, name: str) -> Ok[None] | Err[Exception]:
+        try:
+            project = self._project()
+            index = next((i for i, e in enumerate(project["policyEngines"]) if e["name"] == name), -1)
+            if index == -1:
+                return Err(ResourceNotFoundError(f'Policy engine "{name}" not found.'))
+            project["policyEngines"].pop(index)
+            self.write_project_spec(project)
+
+            # Clean up any gateway references to this engine.
+            changed = False
+            for gw in project.get("agentCoreGateways", []):
+                if (gw.get("policyEngineConfiguration") or {}).get("policyEngineName") == name:
+                    del gw["policyEngineConfiguration"]
+                    changed = True
+            if changed:
+                self.write_project_spec(project)
+            return Ok(None)
+        except Exception as exc:  # noqa: BLE001
+            return Err(to_error(exc))
+
+    def preview_remove(self, name: str) -> RemovalPreview:
+        project = self._project()
+        engine = next((e for e in project["policyEngines"] if e["name"] == name), None)
+        if engine is None:
+            raise Exception(f'Policy engine "{name}" not found.')
+        summary = [f"Removing policy engine: {name}"]
+        if engine.get("policies"):
+            summary.append(f"Note: {len(engine['policies'])} policy(ies) within this engine will also be removed")
+        after = {**project, "policyEngines": [e for e in project["policyEngines"] if e["name"] != name]}
+        affected = [
+            gw
+            for gw in project.get("agentCoreGateways", [])
+            if (gw.get("policyEngineConfiguration") or {}).get("policyEngineName") == name
+        ]
+        if affected:
+            summary.append(
+                f"Note: {len(affected)} gateway(s) referencing this engine will have policyEngineConfiguration removed"
+            )
+            summary.append(
+                "Warning: this may grant agents escalated permissions to invoke gateway tools that were previously "
+                "restricted"
+            )
+        return RemovalPreview(summary, [], [SchemaChange("agentcore/agentcore.json", project, after)])
+
+    def get_removable(self) -> list[RemovableResource]:
+        try:
+            return [RemovableResource(e["name"]) for e in self._project()["policyEngines"]]
+        except Exception:  # noqa: BLE001
+            return []
+
+    # --- queries used by the policy wizard ----------------------------------------------
+
+    def get_existing_engines(self) -> list[str]:
+        return [r.name for r in self.get_removable()]
+
+    def get_unprotected_gateways(self) -> list[str]:
+        """HTTP gateways (protocolType "None") without a policy engine attached."""
+        try:
+            return [
+                gw["name"]
+                for gw in self._project().get("agentCoreGateways", [])
+                if gw.get("protocolType") == "None" and not gw.get("policyEngineConfiguration")
+            ]
+        except Exception:  # noqa: BLE001
+            return []
+
+    def attach_to_gateways(self, engine_name: str, gateway_names: list[str], mode: str) -> None:
+        if not gateway_names:
+            return
+        project = self._project()
+        names = set(gateway_names)
+        for gw in project.get("agentCoreGateways", []):
+            if gw["name"] in names:
+                gw["policyEngineConfiguration"] = {"policyEngineName": engine_name, "mode": mode}
+        self.write_project_spec(project)
+
+    def get_deployed_engine_id(self, engine_name: str) -> str | None:
+        try:
+            state = self.config_io.read_deployed_state()
+            for target in (state.targets or {}).values():
+                engines = (target.resources.policy_engines if target.resources else None) or {}
+                if engine_name in engines:
+                    return engines[engine_name].policy_engine_id
+            return None
+        except Exception:  # noqa: BLE001
+            return None
+
+    def get_deployed_gateway_arn(self) -> str | None:
+        gateways = self.get_deployed_gateways()
+        return next(iter(gateways.values()), None)
+
+    def get_deployed_gateways(self) -> dict[str, str]:
+        """Deployed gateway name -> ARN, excluding MCP-protocol gateways."""
+        try:
+            state = self.config_io.read_deployed_state()
+            project = self._project()
+            mcp_names = {gw["name"] for gw in project.get("agentCoreGateways", []) if gw.get("protocolType") != "None"}
+            result: dict[str, str] = {}
+            for target in (state.targets or {}).values():
+                for name, gw in merge_deployed_gateways(target).items():
+                    if gw is not None and gw.gateway_arn and name not in mcp_names:
+                        result[name] = gw.gateway_arn
+            return result
+        except Exception:  # noqa: BLE001
+            return {}
+
+    def get_project_gateways(self) -> list[dict[str, Any]]:
+        """HTTP gateways with their httpRuntime target names: ``[{"name", "httpTargets"}]``."""
+        try:
+            return [
+                {
+                    "name": gw["name"],
+                    "httpTargets": [t["name"] for t in gw.get("targets") or [] if t.get("targetType") == "httpRuntime"],
+                }
+                for gw in self._project().get("agentCoreGateways", [])
+                if gw.get("protocolType") == "None"
+            ]
+        except Exception:  # noqa: BLE001
+            return []
+
+    # --- CLI ----------------------------------------------------------------------------
+
+    def register_commands(self, add_app: typer.Typer, remove_app: typer.Typer) -> None:
+        @add_app.command("policy-engine", help="Add a policy engine to the project")
+        def add_policy_engine(
+            name: str | None = typer.Option(None, "--name", help="Policy engine name [non-interactive]"),
+            description: str | None = typer.Option(
+                None, "--description", help="Policy engine description [non-interactive]"
+            ),
+            encryption_key_arn: str | None = typer.Option(
+                None, "--encryption-key-arn", help="KMS encryption key ARN [non-interactive]"
+            ),
+            attach_to_gateways: str | None = typer.Option(
+                None,
+                "--attach-to-gateways",
+                help="Comma-separated gateway names to attach this engine to [non-interactive]",
+            ),
+            attach_mode: str | None = typer.Option(
+                None,
+                "--attach-mode",
+                help="Enforcement mode for attached gateways: LOG_ONLY or ENFORCE [non-interactive]",
+            ),
+            as_json: bool = typer.Option(False, "--json", help="Output as JSON [non-interactive]"),
+        ) -> None:
+            if not find_config_root():
+                typer.echo("No agentcore project found. Run `agentcore create` first.", err=True)
+                raise typer.Exit(1)
+
+            if not (name or description or encryption_key_arn or as_json):
+                try:
+                    interactive.require_tty()
+                    self.add_interactive()
+                except typer.Exit:
+                    raise
+                except Exception as exc:  # noqa: BLE001
+                    typer.echo(get_error_message(exc), err=True)
+                    raise typer.Exit(1) from None
+                raise typer.Exit(0)
+
+            def run() -> None:
+                if not name:
+                    raise Exception("--name is required")
+                result = self.add(
+                    AddPolicyEngineOptions(name=name, description=description, encryption_key_arn=encryption_key_arn)
+                )
+                if result.success and attach_to_gateways:
+                    mode = parse_policy_engine_mode(attach_mode or "LOG_ONLY")
+                    self.attach_to_gateways(name, split_csv(attach_to_gateways), mode)
+                if not result.success:
+                    raise result.error
+                if as_json:
+                    emit_json(serialize_result(result))
+                else:
+                    typer.echo(f"Added policy engine '{result.value['engineName']}'")
+
+            run_cli_command(as_json, run)
+
+        @remove_app.command("policy-engine", help="Remove a policy engine from the project")
+        def remove_policy_engine(
+            name: str | None = typer.Option(None, "--name", help="Name of resource to remove [non-interactive]"),
+            yes: bool = typer.Option(False, "-y", "--yes", help="Skip confirmation prompt [non-interactive]"),
+            as_json: bool = typer.Option(False, "--json", help="Output as JSON [non-interactive]"),
+        ) -> None:
+            self._run_remove_with_text(name, yes, as_json)
+
+    def _run_remove_with_text(self, name: str | None, yes: bool, as_json: bool, key: str | None = None) -> None:
+        """Remove flow whose non-JSON output is plain text (policy-engine and policy upstream)."""
+        try:
+            if not find_config_root():
+                typer.echo("No agentcore project found. Run `agentcore create` first.", err=True)
+                raise typer.Exit(1)
+            if name or yes or as_json:
+                if not name:
+                    emit_json({"success": False, "error": "--name is required"})
+                    raise typer.Exit(1)
+                result = self.remove(key or name)
+                label = self.label.lower()
+                if as_json:
+                    emit_json(
+                        {
+                            "success": result.success,
+                            "resourceType": self.kind,
+                            "resourceName": name,
+                            "message": f"Removed {label} '{name}'" if result.success else None,
+                            "note": SOURCE_CODE_NOTE if result.success else None,
+                            "error": None if result.success else get_error_message(result.error),
+                        }
+                    )
+                elif result.success:
+                    typer.echo(f"Removed {label} '{name}'")
+                else:
+                    typer.echo(get_error_message(result.error), err=True)
+                raise typer.Exit(0 if result.success else 1)
+            interactive.require_tty()
+            self.remove_interactive()
+        except typer.Exit:
+            raise
+        except Exception as exc:  # noqa: BLE001 - top-level command boundary, as upstream
+            if as_json:
+                emit_json({"success": False, "error": get_error_message(exc)})
+            else:
+                typer.echo(f"Error: {get_error_message(exc)}", err=True)
+            raise typer.Exit(1) from None
+
+    # --- interactive (AddPolicyFlow engine branch) --------------------------------------
+
+    def prompt_engine_name(self, existing: list[str]) -> str:
+        def validate(value: str) -> str | None:
+            issue = first_issue(PolicyEngineName, value)
+            if issue:
+                return issue
+            return "Policy engine name already exists" if value in existing else None
+
+        return interactive.text(
+            "Policy engine name", default=generate_unique_name("MyPolicyEngine", existing), validate=validate
+        )
+
+    def run_engine_wizard(self, existing: list[str]) -> str:
+        """Name -> attach gateways -> attach mode, then commit. Returns the engine name."""
+        engine_name = self.prompt_engine_name(existing)
+        gateways: list[str] = []
+        mode: str | None = None
+        unprotected = self.get_unprotected_gateways()
+        if unprotected:
+            gateways = interactive.checkbox(
+                f'Attach "{engine_name}" to gateways (these gateways have no policy engine; select which ones to '
+                "protect)",
+                [interactive.Choice(g, g) for g in unprotected],
+            )
+            if gateways:
+                mode = interactive.select(
+                    f"Select enforcement mode (applies to {len(gateways)} gateway{'s' if len(gateways) > 1 else ''} "
+                    f'using "{engine_name}")',
+                    [interactive.Choice(o.title, o.id, o.description) for o in POLICY_ENGINE_MODE_OPTIONS],
+                )
+        result = self.add(AddPolicyEngineOptions(name=engine_name))
+        if not result.success:
+            raise result.error
+        if gateways and mode:
+            self.attach_to_gateways(engine_name, gateways, mode)
+        typer.echo(f"Added policy engine: {engine_name}")
+        typer.echo("Policy engine added to project config. Deploy with `agentcore deploy` to create it in AWS.")
+        typer.secho(
+            "Note: Once deployed and attached to a gateway, all tool calls become default deny.\n"
+            "You must add permit policies to allow agent tool access.\n"
+            "Note: Natural language policy generation requires a deployed engine. Run `agentcore deploy` before "
+            "using the Generate option.",
+            fg="yellow",
+        )
+        return engine_name
+
+    def add_interactive(self) -> None:
+        # Upstream opens the generic AddFlow; the policy engine branch of AddPolicyFlow is asked here.
+        from agentcore.primitives.policy import PRIMITIVE as policy_primitive
+
+        existing = self.get_existing_engines()
+        engine_name = self.run_engine_wizard(existing)
+        if interactive.confirm("Add a policy to this engine?", default=False):
+            policy_primitive.run_policy_wizard(engine_name, [*existing, engine_name])
+
+
+PRIMITIVE = PolicyEnginePrimitive()

agentcore/primitives/registry.py

@@ -0,0 +1,56 @@
+"""Primitive registry (port of src/cli/primitives/registry.ts).
+
+Primitives are discovered rather than listed by hand: any module in ``agentcore.primitives``
+that defines a module-level ``PRIMITIVE`` instance is registered. ``ORDER`` fixes the upstream
+display order. Kinds in ``ORDER`` without a ported primitive are registered as CLI stubs.
+"""
+
+import importlib
+import pkgutil
+from functools import cache
+
+import agentcore.primitives as _package
+from agentcore.primitives.base import BasePrimitive
+
+# Upstream ALL_PRIMITIVES display order.
+ORDER: tuple[str, ...] = (
+    "agent",
+    "harness",
+    "memory",
+    "dataset",
+    "credential",
+    "evaluator",
+    "online-eval",
+    "online-insights",
+    "gateway",
+    "gateway-target",
+    "knowledge-base",
+    "policy-engine",
+    "policy",
+    "config-bundle",
+    "runtime-endpoint",
+    "payment-manager",
+    "payment-connector",
+    "capacity-provider",
+)
+
+
+@cache
+def all_primitives() -> tuple[BasePrimitive, ...]:
+    found: dict[str, BasePrimitive] = {}
+    for info in pkgutil.iter_modules(_package.__path__):
+        if info.name.startswith("_"):  # private helpers (wizards, validation) aren't primitives
+            continue
+        module = importlib.import_module(f"{_package.__name__}.{info.name}")
+        primitive = getattr(module, "PRIMITIVE", None)
+        if isinstance(primitive, BasePrimitive):
+            found[primitive.kind] = primitive
+    rank = {kind: i for i, kind in enumerate(ORDER)}
+    return tuple(sorted(found.values(), key=lambda p: rank.get(p.kind, len(ORDER))))
+
+
+def get_primitive(kind: str) -> BasePrimitive:
+    for primitive in all_primitives():
+        if primitive.kind == kind:
+            return primitive
+    raise ValueError(f"Unknown primitive kind: {kind}")