Files
aios/active-runtime/app/task_plane/policy.py
T

246 lines
9.8 KiB
Python

"""Policy and verification logic (fail-closed)."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Callable, Mapping
from .errors import ForbiddenError, LockedError
from .models import ProbeEvidence, ResultEnvelope, TaskRecord
from .scope import TaskActorContext, actor_context as build_actor_context
RISK_RANK = {"low": 1, "medium": 2, "high": 3}
READONLY_WORKER_ACTIONS = frozenset(
{"claim", "result", "fail", "progress", "heartbeat"}
)
ALL_ACTION_CAPABILITIES = frozenset(
{
"task.creator",
"worker.self_register",
"worker.admin",
"task.claim",
"task.cancel",
"task.delegate",
"task.read",
"task.audit",
}
)
@dataclass
class PolicyDecision:
allowed: bool
reason: str = ""
@dataclass(frozen=True)
class ReadonlyWorkerProfile:
"""A narrowly-scoped identity allowed to process safe Shadow tasks only."""
worker_id: str
identity: str
capabilities: frozenset[str]
task_types: frozenset[str] = frozenset({"code.review", "test.execute"})
def permits_task(self, task: TaskRecord) -> bool:
plan = task.task_plan
if plan.task_type not in self.task_types or plan.risk != "low":
return False
if plan.data_classification not in {"public", "internal"}:
return False
if plan.metadata.get("production_side_effects") is not False:
return False
return set(plan.required_capabilities).issubset(self.capabilities)
class IdentityVerifier:
def __init__(self, trusted_identities: set[str] | None = None):
self.trusted_identities = trusted_identities
def verify(self, identity: str) -> bool:
if self.trusted_identities is None:
return True
return identity in self.trusted_identities
class EvidenceVerifier:
def __init__(self, mandatory: list[str] | None = None):
self.mandatory = set(mandatory or [])
def verify(self, task: TaskRecord, evidences: list[dict]) -> PolicyDecision:
required = set(task.task_plan.evidence_requirements or [])
if not required and not self.mandatory:
return PolicyDecision(True)
observed = {e.get("kind") for e in evidences}
missing = required.union(self.mandatory) - observed
if missing:
return PolicyDecision(False, f"missing evidence: {', '.join(sorted(missing))}")
return PolicyDecision(True)
class PolicyEngine:
"""
Fail-closed policy engine with minimal but explicit decisions.
"""
def __init__(
self,
*,
identity_verifier: IdentityVerifier | None = None,
evidence_verifier: EvidenceVerifier | None = None,
execution_enabled: bool = False,
max_risk: str = "medium",
result_signature_verifier: Callable[[ResultEnvelope], bool] | None = None,
probe_attestation_verifier: Callable[[ProbeEvidence], bool] | None = None,
identity_capabilities: dict[str, set[str]] | None = None,
identity_actor_contexts: Mapping[
str, TaskActorContext | Mapping[str, Any]
]
| None = None,
admin_identities: set[str] | None = None,
readonly_worker_execution_enabled: bool = False,
readonly_worker_profiles: list[ReadonlyWorkerProfile] | None = None,
) -> None:
if max_risk not in RISK_RANK:
raise ValueError("invalid max_risk")
self.identity_verifier = identity_verifier or IdentityVerifier()
self.evidence_verifier = evidence_verifier or EvidenceVerifier()
self.execution_enabled = execution_enabled
self.max_risk = max_risk
self.result_signature_verifier = result_signature_verifier
self.probe_attestation_verifier = probe_attestation_verifier
self.identity_capabilities = {
identity: set(capabilities)
for identity, capabilities in (identity_capabilities or {}).items()
}
self.identity_actor_contexts = dict(identity_actor_contexts or {})
self.admin_identities = set(admin_identities or ())
self.readonly_worker_execution_enabled = readonly_worker_execution_enabled
self.readonly_worker_profiles = {
profile.identity: profile for profile in (readonly_worker_profiles or [])
}
def actor_context(self, identity: str) -> TaskActorContext:
configured = self.identity_actor_contexts.get(
identity, self.identity_actor_contexts.get("*")
)
capabilities = set(
self.identity_capabilities.get(
identity, self.identity_capabilities.get("*", set())
)
)
if identity in self.admin_identities:
capabilities.add("task.admin")
return build_actor_context(identity, configured, capabilities)
def _readonly_worker_allows(
self, action: str, actor: str, task: TaskRecord | None
) -> bool:
profile = self.readonly_worker_profiles.get(actor)
return bool(
self.readonly_worker_execution_enabled
and action in READONLY_WORKER_ACTIONS
and profile is not None
and (task is None or profile.permits_task(task))
)
def assert_readonly_worker_binding(self, actor: str, worker_id: str) -> None:
profile = self.readonly_worker_profiles.get(actor)
if profile is not None and profile.worker_id != worker_id:
raise ForbiddenError("worker identity binding mismatch")
def assert_readonly_worker_registration(self, actor: str, payload: dict) -> None:
profile = self.readonly_worker_profiles.get(actor)
if profile is None:
return
if (
str(payload.get("worker_id") or "") != profile.worker_id
or str(payload.get("identity") or "") != profile.identity
or set(payload.get("capabilities") or []) != set(profile.capabilities)
or str(payload.get("trust_zone") or "") != "isolated-review"
or int(payload.get("max_concurrency") or 0) != 1
or int(payload.get("heartbeat_interval_seconds") or 0) != 15
):
raise ForbiddenError("readonly worker registration mismatch")
def verify_result(
self, result: ResultEnvelope, task: TaskRecord | None = None
) -> PolicyDecision:
computed = result.computed_digest()
if result.result_digest and result.result_digest != computed:
return PolicyDecision(False, "result digest mismatch")
sensitive = bool(
task
and (
task.task_plan.data_classification
in {"confidential", "credential-secret"}
or any(
capability.startswith(
("code.write", "deploy", "secret", "production")
)
for capability in task.task_plan.required_capabilities
)
)
)
if sensitive and not (result.signature or result.attestation_ref):
return PolicyDecision(False, "trusted result proof required")
if result.signature is not None or result.attestation_ref is not None:
if self.result_signature_verifier is None:
return PolicyDecision(False, "result signature verifier unavailable")
if not self.result_signature_verifier(result):
return PolicyDecision(False, "result signature invalid")
result.trust_status = "trusted_result"
else:
result.trust_status = "untrusted_shadow_result"
return PolicyDecision(True)
def authorize(self, action: str, actor: str, task: TaskRecord | None = None, context: dict | None = None) -> PolicyDecision:
context = context or {}
if not self.identity_verifier.verify(actor):
return PolicyDecision(False, "identity untrusted")
action_capabilities = {
"create": "task.creator",
"read": "task.read",
"audit": "task.audit",
"claim": "task.claim",
"result": "task.claim",
"fail": "task.claim",
"progress": "task.claim",
"heartbeat": "task.claim",
"cancel": "task.cancel",
"delegate": "task.delegate",
"worker.register": "worker.self_register",
"worker.admin": "worker.admin",
}
required_capability = action_capabilities.get(action)
if required_capability:
granted = set(
self.identity_capabilities.get(
actor, self.identity_capabilities.get("*", set())
)
)
if required_capability not in granted:
return PolicyDecision(False, f"action {action} not permitted")
if action in {"claim", "result", "fail", "cancel", "progress", "delegate", "heartbeat"} and not self.execution_enabled:
if not self._readonly_worker_allows(action, actor, task):
return PolicyDecision(False, "execution disabled")
if action == "delegate" and task and task.status not in {"COMPLETED", "FAILED_FINAL", "CANCELLED"}:
risk = str(context.get("risk", task.task_plan.risk))
if RISK_RANK.get(risk, 2) > RISK_RANK.get(self.max_risk, 1):
return PolicyDecision(False, f"risk {risk} exceeds policy")
return PolicyDecision(True)
def assert_allowed(self, action: str, actor: str, task: TaskRecord | None = None, context: dict | None = None):
decision = self.authorize(action, actor, task, context=context)
if not decision.allowed:
if decision.reason == "identity untrusted":
raise ForbiddenError(decision.reason)
if decision.reason == "execution disabled":
raise LockedError("execution disabled")
raise ForbiddenError(decision.reason)