actshield 0.9.0__py3-none-any.whl

actshield/tools/__init__.py

@@ -0,0 +1,20 @@
+"""Tools package."""
+
+from actshield.tools.interceptor import protected_tool
+from actshield.tools.tool import (
+    Resource,
+    SensitivityLevel,
+    ToolDefinition,
+    ToolRequest,
+    ToolResult,
+)
+
+__all__ = [
+    "Resource",
+    "SensitivityLevel",
+    "ToolDefinition",
+    "ToolRequest",
+    "ToolResult",
+    "protected_tool",
+]
+

actshield/tools/interceptor.py

@@ -0,0 +1,256 @@
+"""Tool execution interception and runtime policy enforcement wrapper."""
+
+import functools
+import inspect
+import time
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+from actshield.context.context import Context
+from actshield.context.provenance import ContextSource
+from actshield.context.taint import TaintState
+from actshield.decisions.decision import DecisionAction, SecurityDecision
+from actshield.tools.tool import Resource, SensitivityLevel, ToolDefinition, ToolRequest, ToolResult
+from actshield.tracing.correlation import generate_id, get_current_context
+from actshield.tracing.events import EventType
+
+if TYPE_CHECKING:
+    from actshield.client import ActShield
+
+
+def protected_tool(
+    guard: "ActShield",
+    name: Optional[str] = None,
+    description: str = "",
+    sensitivity: SensitivityLevel = SensitivityLevel.MEDIUM,
+    required_capabilities: Optional[List[str]] = None,
+    classification: str = "internal",
+    target_resources: Optional[List[Resource]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Decorator to register and wrap a tool function with ActShield security enforcement."""
+
+    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+        tool_name = name or fn.__name__
+        tool_desc = description or fn.__doc__ or ""
+        req_caps = required_capabilities or []
+        res_list = target_resources or []
+
+        tool_def = ToolDefinition(
+            tool_id=generate_id("tool"),
+            name=tool_name,
+            description=tool_desc,
+            sensitivity=sensitivity,
+            required_capabilities=req_caps,
+            classification=classification,
+            target_resources=res_list,
+            metadata=metadata or {},
+        )
+
+        guard.register_tool(tool_def)
+
+        if inspect.iscoroutinefunction(fn):
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                return await _execute_protected_tool_async(guard, tool_def, fn, *args, **kwargs)
+            async_wrapper.tool_def = tool_def  # type: ignore
+            return async_wrapper
+        else:
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                return _execute_protected_tool_sync(guard, tool_def, fn, *args, **kwargs)
+            sync_wrapper.tool_def = tool_def  # type: ignore
+            return sync_wrapper
+
+    return decorator
+
+
+def _prepare_tool_request(
+    guard: "ActShield",
+    tool_def: ToolDefinition,
+    fn: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> ToolRequest:
+    ctx = get_current_context()
+    agent_id = ctx.agent_id if ctx else None
+    
+    # Map positional args to param names
+    sig = inspect.signature(fn)
+    bound_args = sig.bind_partial(*args, **kwargs)
+    bound_args.apply_defaults()
+    raw_arguments = dict(bound_args.arguments)
+
+    # Sanitize arguments for logging/event storage
+    safe_arguments = guard.config.sanitize(raw_arguments)
+
+    context_ids = []
+    if ctx and ctx.context_id:
+        context_ids.append(ctx.context_id)
+
+    target_res = tool_def.target_resources[0] if tool_def.target_resources else None
+
+    return ToolRequest(
+        request_id=generate_id("req"),
+        tool_id=tool_def.tool_id,
+        tool_name=tool_def.name,
+        agent_id=agent_id,
+        arguments=safe_arguments,
+        context_ids=context_ids,
+        target_resource=target_res,
+    )
+
+
+def _execute_protected_tool_sync(
+    guard: "ActShield",
+    tool_def: ToolDefinition,
+    fn: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    request = _prepare_tool_request(guard, tool_def, fn, *args, **kwargs)
+    ctx = get_current_context()
+
+    # Emit tool.requested
+    guard.emit_event(
+        event_type=EventType.TOOL_REQUESTED,
+        agent_id=request.agent_id,
+        payload={
+            "tool_id": tool_def.tool_id,
+            "tool_name": tool_def.name,
+            "arguments": request.arguments,
+            "sensitivity": tool_def.sensitivity.value,
+        }
+    )
+
+    # Emit resource.access_requested if applicable
+    if request.target_resource:
+        guard.emit_event(
+            event_type=EventType.RESOURCE_ACCESS_REQUESTED,
+            agent_id=request.agent_id,
+            payload={
+                "resource_id": request.target_resource.resource_id,
+                "resource_name": request.target_resource.name,
+                "sensitivity": request.target_resource.sensitivity.value,
+            }
+        )
+
+    # Evaluate security policy
+    decision: SecurityDecision = guard.evaluate_tool_invocation(tool_def, request)
+
+    if decision.is_blocked:
+        raise PermissionError(
+            f"ActShield Policy Enforced: Tool '{tool_def.name}' invocation BLOCKED. "
+            f"Reason: [{decision.reason_code}] {decision.explanation}"
+        )
+
+    start_t = time.perf_counter()
+    err_str = None
+    try:
+        raw_output = fn(*args, **kwargs)
+        elapsed_ms = (time.perf_counter() - start_t) * 1000
+        output_for_event = guard.config.sanitize(raw_output)
+        
+        # Emit tool.invoked
+        guard.emit_event(
+            event_type=EventType.TOOL_INVOKED,
+            agent_id=request.agent_id,
+            payload={
+                "tool_id": tool_def.tool_id,
+                "tool_name": tool_def.name,
+                "success": True,
+                "execution_time_ms": elapsed_ms,
+                "output_summary": str(output_for_event)[:200],
+            }
+        )
+        return raw_output
+    except Exception as e:
+        elapsed_ms = (time.perf_counter() - start_t) * 1000
+        err_str = str(e)
+        guard.emit_event(
+            event_type=EventType.TOOL_INVOKED,
+            agent_id=request.agent_id,
+            payload={
+                "tool_id": tool_def.tool_id,
+                "tool_name": tool_def.name,
+                "success": False,
+                "error": err_str,
+                "execution_time_ms": elapsed_ms,
+            }
+        )
+        raise
+
+
+async def _execute_protected_tool_async(
+    guard: "ActShield",
+    tool_def: ToolDefinition,
+    fn: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    request = _prepare_tool_request(guard, tool_def, fn, *args, **kwargs)
+
+    # Emit tool.requested
+    guard.emit_event(
+        event_type=EventType.TOOL_REQUESTED,
+        agent_id=request.agent_id,
+        payload={
+            "tool_id": tool_def.tool_id,
+            "tool_name": tool_def.name,
+            "arguments": request.arguments,
+            "sensitivity": tool_def.sensitivity.value,
+        }
+    )
+
+    if request.target_resource:
+        guard.emit_event(
+            event_type=EventType.RESOURCE_ACCESS_REQUESTED,
+            agent_id=request.agent_id,
+            payload={
+                "resource_id": request.target_resource.resource_id,
+                "resource_name": request.target_resource.name,
+                "sensitivity": request.target_resource.sensitivity.value,
+            }
+        )
+
+    decision: SecurityDecision = guard.evaluate_tool_invocation(tool_def, request)
+
+    if decision.is_blocked:
+        raise PermissionError(
+            f"ActShield Policy Enforced: Tool '{tool_def.name}' invocation BLOCKED. "
+            f"Reason: [{decision.reason_code}] {decision.explanation}"
+        )
+
+    start_t = time.perf_counter()
+    try:
+        raw_output = await fn(*args, **kwargs)
+        elapsed_ms = (time.perf_counter() - start_t) * 1000
+        output_for_event = guard.config.sanitize(raw_output)
+
+        guard.emit_event(
+            event_type=EventType.TOOL_INVOKED,
+            agent_id=request.agent_id,
+            payload={
+                "tool_id": tool_def.tool_id,
+                "tool_name": tool_def.name,
+                "success": True,
+                "execution_time_ms": elapsed_ms,
+                "output_summary": str(output_for_event)[:200],
+            }
+        )
+        return raw_output
+    except Exception as e:
+        elapsed_ms = (time.perf_counter() - start_t) * 1000
+        guard.emit_event(
+            event_type=EventType.TOOL_INVOKED,
+            agent_id=request.agent_id,
+            payload={
+                "tool_id": tool_def.tool_id,
+                "tool_name": tool_def.name,
+                "success": False,
+                "error": str(e),
+                "execution_time_ms": elapsed_ms,
+            }
+        )
+        raise
+
+

actshield/tools/tool.py

@@ -0,0 +1,71 @@
+"""Tool and resource models for protected execution and access governance."""
+
+from enum import Enum
+from typing import Any, Dict, List, Optional
+from pydantic import BaseModel, Field
+
+from actshield.risk.models import RiskLevel
+from actshield.tracing.correlation import generate_id
+
+
+class SensitivityLevel(str, Enum):
+    """Sensitivity classification of tools and resources."""
+    LOW = "LOW"
+    MEDIUM = "MEDIUM"
+    HIGH = "HIGH"
+    CRITICAL = "CRITICAL"
+
+
+class Resource(BaseModel):
+    """A protected asset, system, database, or API target accessed via tools."""
+    resource_id: str = Field(default_factory=lambda: generate_id("rsc"))
+    name: str = Field(..., description="Resource name (e.g. 'financial_database', 'user_pii_vault')")
+    resource_type: str = Field(default="api", description="Resource type ('database', 'filesystem', 'api', 'model')")
+    uri: Optional[str] = Field(default=None, description="URI or path locator")
+    sensitivity: SensitivityLevel = Field(default=SensitivityLevel.MEDIUM)
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolDefinition(BaseModel):
+    """Metadata and security properties defining an executable tool or MCP endpoint."""
+    
+    tool_id: str = Field(default_factory=lambda: generate_id("tool"))
+    name: str = Field(..., description="Canonical tool name (e.g. 'execute_sql_query', 'fetch_webpage')")
+    description: str = Field(default="", description="Functional description of the tool")
+    sensitivity: SensitivityLevel = Field(default=SensitivityLevel.MEDIUM, description="Sensitivity classification")
+    required_capabilities: List[str] = Field(
+        default_factory=list,
+        description="List of capability names an agent must possess to invoke this tool"
+    )
+    classification: str = Field(
+        default="internal",
+        description="Origin classification: 'internal', 'external', 'mcp', 'sensitive'"
+    )
+    target_resources: List[Resource] = Field(
+        default_factory=list,
+        description="Underlying resources accessed by this tool"
+    )
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolRequest(BaseModel):
+    """A concrete invocation request targeting a tool."""
+    
+    request_id: str = Field(default_factory=lambda: generate_id("req"))
+    tool_id: str = Field(default_factory=lambda: generate_id("tool"), description="Tool ID being invoked")
+    tool_name: str = Field(..., description="Tool name")
+    agent_id: Optional[str] = Field(default=None, description="Requesting agent ID")
+
+    arguments: Dict[str, Any] = Field(default_factory=dict, description="Tool invocation parameters")
+    context_ids: List[str] = Field(default_factory=list, description="IDs of context objects influencing this request")
+    target_resource: Optional[Resource] = Field(default=None, description="Specific target resource if applicable")
+
+
+class ToolResult(BaseModel):
+    """The result or error produced by a tool execution."""
+    
+    success: bool = Field(default=True)
+    output: Any = Field(default=None)
+    error: Optional[str] = Field(default=None)
+    execution_time_ms: float = Field(default=0.0)
+

actshield/tracing/__init__.py

@@ -0,0 +1,38 @@
+"""Tracing package."""
+
+from actshield.tracing.correlation import (
+    CorrelationContext,
+    SpanScope,
+    generate_id,
+    get_current_agent_id,
+    get_current_context,
+    get_current_delegation_id,
+    get_current_span_id,
+    get_current_task_id,
+    get_current_trace_id,
+    reset_current_context,
+    set_current_context,
+)
+from actshield.tracing.events import EventType, SecurityEvent
+from actshield.tracing.tracer import CausalEdge, CausalGraph, CausalNode, TraceManager
+
+__all__ = [
+    "CorrelationContext",
+    "SpanScope",
+    "generate_id",
+    "get_current_context",
+    "set_current_context",
+    "reset_current_context",
+    "get_current_trace_id",
+    "get_current_span_id",
+    "get_current_agent_id",
+    "get_current_task_id",
+    "get_current_delegation_id",
+    "EventType",
+    "SecurityEvent",
+    "CausalNode",
+    "CausalEdge",
+    "CausalGraph",
+    "TraceManager",
+]
+

actshield/tracing/correlation.py

@@ -0,0 +1,173 @@
+"""Correlation context management for distributed and multi-agent tracing."""
+
+import contextvars
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+
+def generate_id(prefix: str = "") -> str:
+    """Generate a unique hex/uuid based ID with an optional prefix."""
+    unique_part = uuid.uuid4().hex[:16]
+    if prefix:
+        return f"{prefix}_{unique_part}"
+    return unique_part
+
+
+@dataclass(frozen=True)
+class CorrelationContext:
+    """Immutable correlation context representing current execution coordinates."""
+    trace_id: str
+    span_id: str
+    parent_span_id: Optional[str] = None
+    task_id: Optional[str] = None
+    agent_id: Optional[str] = None
+    parent_agent_id: Optional[str] = None
+    delegation_id: Optional[str] = None
+    parent_event_id: Optional[str] = None
+    context_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def derive_child_span(
+        self,
+        span_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        task_id: Optional[str] = None,
+        delegation_id: Optional[str] = None,
+        context_id: Optional[str] = None,
+        parent_event_id: Optional[str] = None,
+        extra_metadata: Optional[Dict[str, Any]] = None,
+    ) -> "CorrelationContext":
+        """Derive a child correlation context preserving the trace_id and updating span/agent relationships."""
+        new_span_id = span_id or generate_id("spn")
+        new_metadata = dict(self.metadata)
+        if extra_metadata:
+            new_metadata.update(extra_metadata)
+
+        return CorrelationContext(
+            trace_id=self.trace_id,
+            span_id=new_span_id,
+            parent_span_id=self.span_id,
+            task_id=task_id if task_id is not None else self.task_id,
+            agent_id=agent_id if agent_id is not None else self.agent_id,
+            parent_agent_id=self.agent_id if (agent_id is not None and agent_id != self.agent_id) else self.parent_agent_id,
+            delegation_id=delegation_id if delegation_id is not None else self.delegation_id,
+            parent_event_id=parent_event_id if parent_event_id is not None else self.parent_event_id,
+            context_id=context_id if context_id is not None else self.context_id,
+            metadata=new_metadata,
+        )
+
+
+_current_correlation_ctx: contextvars.ContextVar[Optional[CorrelationContext]] = contextvars.ContextVar(
+    "agentguard_correlation_ctx", default=None
+)
+
+
+def get_current_context() -> Optional[CorrelationContext]:
+    """Retrieve the active correlation context for the current async/sync task."""
+    return _current_correlation_ctx.get()
+
+
+def set_current_context(ctx: Optional[CorrelationContext]) -> contextvars.Token:
+    """Set the active correlation context and return the token to reset later."""
+    return _current_correlation_ctx.set(ctx)
+
+
+def reset_current_context(token: contextvars.Token) -> None:
+    """Reset the correlation context using a saved token."""
+    _current_correlation_ctx.reset(token)
+
+
+def get_current_trace_id() -> Optional[str]:
+    """Get the active trace ID if in an active trace scope."""
+    ctx = get_current_context()
+    return ctx.trace_id if ctx else None
+
+
+def get_current_span_id() -> Optional[str]:
+    """Get the active span ID if in an active span scope."""
+    ctx = get_current_context()
+    return ctx.span_id if ctx else None
+
+
+def get_current_agent_id() -> Optional[str]:
+    """Get the active agent ID."""
+    ctx = get_current_context()
+    return ctx.agent_id if ctx else None
+
+
+def get_current_task_id() -> Optional[str]:
+    """Get the active task ID."""
+    ctx = get_current_context()
+    return ctx.task_id if ctx else None
+
+
+def get_current_delegation_id() -> Optional[str]:
+    """Get the active delegation ID."""
+    ctx = get_current_context()
+    return ctx.delegation_id if ctx else None
+
+
+class SpanScope:
+    """Context manager for establishing or nesting a correlation span."""
+
+    def __init__(
+        self,
+        name: str = "span",
+        trace_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        task_id: Optional[str] = None,
+        delegation_id: Optional[str] = None,
+        context_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.name = name
+        self._requested_trace_id = trace_id
+        self._requested_span_id = span_id
+        self._agent_id = agent_id
+        self._task_id = task_id
+        self._delegation_id = delegation_id
+        self._context_id = context_id
+        self._metadata = metadata or {}
+        self._token: Optional[contextvars.Token] = None
+        self.context: Optional[CorrelationContext] = None
+        self.start_time: float = 0.0
+
+    def __enter__(self) -> CorrelationContext:
+        self.start_time = time.time()
+        parent_ctx = get_current_context()
+        if parent_ctx is not None:
+            self.context = parent_ctx.derive_child_span(
+                span_id=self._requested_span_id,
+                agent_id=self._agent_id,
+                task_id=self._task_id,
+                delegation_id=self._delegation_id,
+                context_id=self._context_id,
+                extra_metadata=self._metadata,
+            )
+        else:
+            trace_id = self._requested_trace_id or generate_id("trc")
+            span_id = self._requested_span_id or generate_id("spn")
+            self.context = CorrelationContext(
+                trace_id=trace_id,
+                span_id=span_id,
+                task_id=self._task_id,
+                agent_id=self._agent_id,
+                delegation_id=self._delegation_id,
+                context_id=self._context_id,
+                metadata=self._metadata,
+            )
+        self._token = set_current_context(self.context)
+        return self.context
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        if self._token is not None:
+            reset_current_context(self._token)
+
+    async def __aenter__(self) -> CorrelationContext:
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.__exit__(exc_type, exc_val, exc_tb)

actshield/tracing/events.py

@@ -0,0 +1,75 @@
+"""Strongly typed security events for multi-agent observability and enforcement."""
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Dict, Optional
+from pydantic import BaseModel, ConfigDict, Field
+
+from actshield.tracing.correlation import generate_id
+
+
+class EventType(str, Enum):
+    """Enumeration of standard security-relevant event types."""
+    AGENT_CREATED = "agent.created"
+    AGENT_INVOKED = "agent.invoked"
+    TASK_STARTED = "task.started"
+    TASK_COMPLETED = "task.completed"
+    AGENT_DELEGATED = "agent.delegated"
+    CONTEXT_RECEIVED = "context.received"
+    CONTEXT_GENERATED = "context.generated"
+    CONTEXT_PROPAGATED = "context.propagated"
+    CONTEXT_SANITIZED = "context.sanitized"
+    TOOL_REQUESTED = "tool.requested"
+    TOOL_INVOKED = "tool.invoked"
+    TOOL_COMPLETED = "tool.completed"
+    RESOURCE_ACCESS_REQUESTED = "resource.access_requested"
+    SECURITY_EVALUATED = "security.evaluated"
+    SECURITY_DECISION = "security.decision"
+    HUMAN_APPROVAL_REQUESTED = "human.approval_requested"
+    HUMAN_APPROVAL_RECEIVED = "human.approval_received"
+    APPROVAL_REQUESTED = "approval.requested"
+    APPROVAL_GRANTED = "approval.granted"
+    APPROVAL_REJECTED = "approval.rejected"
+    INCIDENT_CREATED = "incident.created"
+    AI_ANALYSIS_COMPLETED = "ai.analysis_completed"
+    POSTURE_EVALUATED = "posture.evaluated"
+    INCIDENT_TRANSITIONED = "incident.transitioned"
+    RESPONSE_EXECUTED = "response.executed"
+    DRIFT_DETECTED = "drift.detected"
+    GATE_EVALUATED = "gate.evaluated"
+    CUSTOM_EVENT = "custom.event"
+    SECURITY_ALERT = "security.alert"
+
+
+
+
+class SecurityEvent(BaseModel):
+    """Core security event model for distributed provenance and causal reconstruction."""
+    
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    event_id: str = Field(default_factory=lambda: generate_id("evt"), description="Unique event ID")
+    event_type: EventType = Field(..., description="Categorical type of the security event")
+    timestamp: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        description="UTC timestamp of when the event occurred"
+    )
+    
+    # Correlation coordinates
+    trace_id: str = Field(..., description="Root trace ID for causal correlation")
+    span_id: str = Field(..., description="Span ID representing the current execution boundary")
+    parent_event_id: Optional[str] = Field(default=None, description="Immediate causal parent event ID")
+    
+    # Agent coordinates
+    agent_id: Optional[str] = Field(default=None, description="Acting agent ID")
+    parent_agent_id: Optional[str] = Field(default=None, description="Delegating or parent agent ID")
+    
+    # Execution coordinates
+    task_id: Optional[str] = Field(default=None, description="Active task ID")
+    delegation_id: Optional[str] = Field(default=None, description="Active delegation context ID")
+    context_id: Optional[str] = Field(default=None, description="Associated context or data lineage ID")
+    
+    # Event data
+    payload: Dict[str, Any] = Field(default_factory=dict, description="Structured event payload")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Custom metadata and tags")
+