import re
from dataclasses import dataclass
from enum import Enum

class RiskLevel(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"

@dataclass
class ValidationResult:
    is_safe: bool
    risk_level: RiskLevel
    matched_rules: list[str]
    sanitized_input: str

INJECTION_PATTERNS = [
    (r"ignore\s+(all\s+)?(previous|prior|above)\s+(instructions|prompts|rules)", "instruction_override"),
    (r"you\s+are\s+now\s+(a|an|the)\s+", "role_hijack"),
    (r"system\s*:\s*", "system_prompt_inject"),
    (r"<\|?(system|im_start|endoftext)\|?>", "control_token_inject"),
    (r"\[INST\]|\[\/INST\]|<<SYS>>", "template_inject"),
    (r"(?:execute|run|eval)\s*\(", "code_execution_attempt"),
    (r"(?:curl|wget|nc|ncat)\s+", "network_command_inject"),
    (r"(?:rm\s+-rf|mkfs|dd\s+if=|chmod\s+777)", "destructive_command"),
    (r"(?:\/etc\/passwd|\/etc\/shadow|\.env\b|\.ssh\/)", "path_traversal"),
    (r"(?:BEGIN\s+(?:RSA|DSA|EC)\s+PRIVATE\s+KEY)", "secret_exfil_attempt"),
]

def validate_agent_input(user_input: str, max_length: int = 4096) -> ValidationResult:
    """Validate and sanitize input before passing to agent."""
    matched = []
    risk = RiskLevel.LOW

    # Length check
    if len(user_input) > max_length:
        matched.append("input_too_long")
        risk = RiskLevel.MEDIUM

    # Null byte and control character removal
    sanitized = user_input.replace("\x00", "")
    sanitized = re.sub(r"[\x01-\x08\x0b\x0c\x0e-\x1f]", "", sanitized)

    # Pattern matching
    for pattern, rule_name in INJECTION_PATTERNS:
        if re.search(pattern, sanitized, re.IGNORECASE):
            matched.append(rule_name)
            risk = RiskLevel.HIGH

    # Stacked injection detection (multiple suspicious patterns)
    if len(matched) >= 3:
        risk = RiskLevel.CRITICAL

    is_safe = risk in (RiskLevel.LOW, RiskLevel.MEDIUM)

    return ValidationResult(
        is_safe=is_safe,
        risk_level=risk,
        matched_rules=matched,
        sanitized_input=sanitized[:max_length] if is_safe else "",
    )

from functools import wraps
from typing import Callable

def input_guard(validator: Callable = validate_agent_input):
    """Decorator that guards agent entry points against unsafe input."""
    def decorator(func):
        @wraps(func)
        async def wrapper(user_input: str, *args, **kwargs):
            result = validator(user_input)

            if result.risk_level == RiskLevel.CRITICAL:
                await log_security_event(
                    event="input_blocked",
                    risk=result.risk_level.value,
                    rules=result.matched_rules,
                    input_hash=hashlib.sha256(user_input.encode()).hexdigest(),
                )
                raise InputRejectedError(
                    f"Input blocked: matched {result.matched_rules}"
                )

            if result.risk_level == RiskLevel.HIGH:
                await log_security_event(
                    event="input_flagged",
                    risk=result.risk_level.value,
                    rules=result.matched_rules,
                )
                # Allow through but flag for review
                kwargs["_security_flags"] = result.matched_rules

            return await func(result.sanitized_input, *args, **kwargs)
        return wrapper
    return decorator

undefined

undefined

undefined

undefined

undefined

from dataclasses import dataclass, field

@dataclass
class ToolPolicy:
    name: str
    allowed_args: dict[str, type]     # parameter name -> expected type
    max_calls_per_session: int = 10
    requires_approval: bool = False
    allowed_patterns: list[str] = field(default_factory=list)
    blocked_patterns: list[str] = field(default_factory=list)

TOOL_ALLOWLIST: dict[str, ToolPolicy] = {
    "read_file": ToolPolicy(
        name="read_file",
        allowed_args={"path": str},
        max_calls_per_session=20,
        allowed_patterns=[r"^/workspace/", r"^/data/public/"],
        blocked_patterns=[r"\.env$", r"\.key$", r"\.pem$", r"/etc/", r"/proc/"],
    ),
    "run_query": ToolPolicy(
        name="run_query",
        allowed_args={"sql": str, "database": str},
        max_calls_per_session=5,
        allowed_patterns=[r"^SELECT\s", r"^EXPLAIN\s"],
        blocked_patterns=[r"\bDROP\b", r"\bDELETE\b", r"\bUPDATE\b", r"\bINSERT\b", r"\bALTER\b"],
    ),
    "http_request": ToolPolicy(
        name="http_request",
        allowed_args={"url": str, "method": str},
        max_calls_per_session=10,
        requires_approval=True,
        allowed_patterns=[r"^https://api\.internal\."],
        blocked_patterns=[r"^https?://169\.254\.", r"^https?://metadata\.google\."],
    ),
    "execute_code": ToolPolicy(
        name="execute_code",
        allowed_args={"code": str, "language": str},
        max_calls_per_session=3,
        requires_approval=True,
        blocked_patterns=[r"import\s+subprocess", r"import\s+os", r"__import__", r"eval\(", r"exec\("],
    ),
}

class ToolGatekeeper:
    def __init__(self, allowlist: dict[str, ToolPolicy]):
        self.allowlist = allowlist
        self.call_counts: dict[str, int] = {}

    async def authorize(self, tool_name: str, args: dict) -> bool:
        if tool_name not in self.allowlist:
            await log_security_event(
                event="tool_denied_not_in_allowlist",
                tool=tool_name,
            )
            return False

        policy = self.allowlist[tool_name]

        # Check call count
        count = self.call_counts.get(tool_name, 0)
        if count >= policy.max_calls_per_session:
            await log_security_event(
                event="tool_denied_rate_limit",
                tool=tool_name,
                count=count,
            )
            return False

        # Validate argument types
        for arg_name, expected_type in policy.allowed_args.items():
            if arg_name in args and not isinstance(args[arg_name], expected_type):
                return False

        # Check patterns against all string arguments
        for arg_value in args.values():
            if not isinstance(arg_value, str):
                continue
            # Must match at least one allowed pattern (if any defined)
            if policy.allowed_patterns:
                if not any(re.search(p, arg_value, re.IGNORECASE) for p in policy.allowed_patterns):
                    return False
            # Must not match any blocked pattern
            if any(re.search(p, arg_value, re.IGNORECASE) for p in policy.blocked_patterns):
                await log_security_event(
                    event="tool_denied_blocked_pattern",
                    tool=tool_name,
                    arg_value_hash=hashlib.sha256(arg_value.encode()).hexdigest(),
                )
                return False

        self.call_counts[tool_name] = count + 1
        return True

undefined

undefined

import httpx

OPA_URL = "http://localhost:8181/v1/data/agent/tool_access"

async def check_tool_permission(agent_role: str, tool_name: str, context: dict) -> dict:
    """Query OPA for tool access decision."""
    payload = {
        "input": {
            "agent_role": agent_role,
            "tool_name": tool_name,
            "session_id": context.get("session_id"),
            "tenant_id": context.get("tenant_id"),
        }
    }
    async with httpx.AsyncClient(timeout=2.0) as client:
        resp = await client.post(OPA_URL, json=payload)
        resp.raise_for_status()
        result = resp.json().get("result", {})
    return {
        "allowed": result.get("allow", False),
        "requires_approval": result.get("requires_approval", False),
    }

undefined

```bash

undefined

import re
from typing import NamedTuple

class PIIMatch(NamedTuple):
    pii_type: str
    start: int
    end: int

PII_PATTERNS = {
    "ssn": r"\b\d{3}-\d{2}-\d{4}\b",
    "credit_card": r"\b(?:\d{4}[\s-]?){3}\d{4}\b",
    "email": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
    "phone_us": r"\b(?:\+1[\s.-]?)?\(?\d{3}\)?[\s.-]?\d{3}[\s.-]?\d{4}\b",
    "aws_key": r"\bAKIA[0-9A-Z]{16}\b",
    "private_key": r"-----BEGIN (?:RSA |EC |DSA )?PRIVATE KEY-----",
    "jwt": r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b",
    "ipv4_internal": r"\b(?:10\.\d{1,3}\.\d{1,3}\.\d{1,3}|172\.(?:1[6-9]|2\d|3[01])\.\d{1,3}\.\d{1,3}|192\.168\.\d{1,3}\.\d{1,3})\b",
    "connection_string": r"(?:mongodb|postgres|mysql|redis):\/\/[^\s\"']+",
}

def scan_for_pii(text: str) -> list[PIIMatch]:
    """Scan text for PII and secrets."""
    matches = []
    for pii_type, pattern in PII_PATTERNS.items():
        for m in re.finditer(pattern, text, re.IGNORECASE):
            matches.append(PIIMatch(pii_type, m.start(), m.end()))
    return matches

def redact_output(text: str) -> tuple[str, list[PIIMatch]]:
    """Redact PII from agent output. Returns redacted text and match list."""
    matches = scan_for_pii(text)
    if not matches:
        return text, []

    # Sort by position descending so replacements don't shift indices
    sorted_matches = sorted(matches, key=lambda m: m.start, reverse=True)
    redacted = text
    for match in sorted_matches:
        placeholder = f"[REDACTED_{match.pii_type.upper()}]"
        redacted = redacted[:match.start] + placeholder + redacted[match.end:]

    return redacted, matches

@dataclass
class OutputPolicy:
    max_length: int = 16384
    block_on_pii: bool = True
    block_on_secrets: bool = True
    allowed_domains: list[str] = field(default_factory=lambda: [
        "docs.example.com", "api.example.com"
    ])

async def validate_agent_output(
    response: str,
    policy: OutputPolicy,
    session_id: str,
) -> str:
    """Validate and filter agent output before returning to user."""
    # Length check
    if len(response) > policy.max_length:
        response = response[:policy.max_length] + "\n\n[Output truncated]"

    # PII/secret scan
    redacted, matches = redact_output(response)
    if matches:
        secret_types = {m.pii_type for m in matches}
        await log_security_event(
            event="output_pii_detected",
            session_id=session_id,
            pii_types=list(secret_types),
            count=len(matches),
        )
        if policy.block_on_secrets and secret_types & {"aws_key", "private_key", "jwt", "connection_string"}:
            return "[Response blocked: contained credentials. This incident has been logged.]"
        if policy.block_on_pii:
            return redacted

    # URL allowlist check — block responses that contain links to unapproved domains
    urls = re.findall(r"https?://([^/\s\"']+)", response)
    for domain in urls:
        if not any(domain.endswith(allowed) for allowed in policy.allowed_domains):
            response = re.sub(
                rf"https?://{re.escape(domain)}[^\s\"']*",
                "[URL_REMOVED]",
                response,
            )

    return response

import json
import time
import hashlib
from datetime import datetime, timezone

class AgentAuditLogger:
    def __init__(self, service_name: str = "agent-platform"):
        self.service_name = service_name

    def log_event(self, event: dict) -> str:
        """Emit a structured audit log entry. Returns the event ID."""
        event_id = hashlib.sha256(
            f"{time.time_ns()}-{json.dumps(event, sort_keys=True)}".encode()
        ).hexdigest()[:16]

        record = {
            "event_id": event_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "service": self.service_name,
            **event,
        }

        # Emit as structured JSON line (ship to SIEM via Fluent Bit / Vector)
        print(json.dumps(record, default=str), flush=True)
        return event_id

    def log_tool_call(self, session_id: str, tool: str, args: dict,
                      result_status: str, duration_ms: float, agent_role: str):
        return self.log_event({
            "event_type": "tool_call",
            "session_id": session_id,
            "tool": tool,
            "args_hash": hashlib.sha256(json.dumps(args, sort_keys=True).encode()).hexdigest(),
            "result_status": result_status,
            "duration_ms": round(duration_ms, 2),
            "agent_role": agent_role,
        })

    def log_input_validation(self, session_id: str, risk_level: str,
                             matched_rules: list[str]):
        return self.log_event({
            "event_type": "input_validation",
            "session_id": session_id,
            "risk_level": risk_level,
            "matched_rules": matched_rules,
        })

    def log_output_filter(self, session_id: str, pii_types: list[str],
                          action_taken: str):
        return self.log_event({
            "event_type": "output_filter",
            "session_id": session_id,
            "pii_types_detected": pii_types,
            "action": action_taken,
        })

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource

undefined

undefined

undefined

import time
from dataclasses import dataclass, field

@dataclass
class TokenBudget:
    max_input_tokens_per_request: int = 4096
    max_output_tokens_per_request: int = 4096
    max_tokens_per_session: int = 100_000
    max_tokens_per_hour: int = 500_000
    max_tool_calls_per_session: int = 50
    max_cost_per_session_usd: float = 5.00

class BudgetEnforcer:
    def __init__(self, budget: TokenBudget):
        self.budget = budget
        self.sessions: dict[str, dict] = {}

    def _get_session(self, session_id: str) -> dict:
        if session_id not in self.sessions:
            self.sessions[session_id] = {
                "total_tokens": 0,
                "tool_calls": 0,
                "estimated_cost_usd": 0.0,
                "hourly_tokens": 0,
                "hour_start": time.time(),
            }
        return self.sessions[session_id]

    def check_budget(self, session_id: str, input_tokens: int,
                     estimated_output_tokens: int) -> tuple[bool, str]:
        """Returns (allowed, reason)."""
        s = self._get_session(session_id)

        # Reset hourly counter if needed
        if time.time() - s["hour_start"] > 3600:
            s["hourly_tokens"] = 0
            s["hour_start"] = time.time()

        if input_tokens > self.budget.max_input_tokens_per_request:
            return False, f"Input tokens {input_tokens} exceeds limit {self.budget.max_input_tokens_per_request}"

        projected = s["total_tokens"] + input_tokens + estimated_output_tokens
        if projected > self.budget.max_tokens_per_session:
            return False, "Session token budget exhausted"

        if s["hourly_tokens"] + input_tokens > self.budget.max_tokens_per_hour:
            return False, "Hourly token budget exhausted"

        if s["estimated_cost_usd"] > self.budget.max_cost_per_session_usd:
            return False, f"Session cost ${s['estimated_cost_usd']:.2f} exceeds limit"

        return True, "ok"

    def record_usage(self, session_id: str, input_tokens: int,
                     output_tokens: int, cost_usd: float):
        s = self._get_session(session_id)
        s["total_tokens"] += input_tokens + output_tokens
        s["hourly_tokens"] += input_tokens + output_tokens
        s["estimated_cost_usd"] += cost_usd

    def record_tool_call(self, session_id: str) -> tuple[bool, str]:
        s = self._get_session(session_id)
        s["tool_calls"] += 1
        if s["tool_calls"] > self.budget.max_tool_calls_per_session:
            return False, "Tool call limit exceeded"
        return True, "ok"

undefined