Zabezpieczanie dostępu agentów do narzędzi

from enum import Enum
from dataclasses import dataclass

class Permission(Enum):
    READ_CUSTOMER_PROFILE = 'read_customer_profile'
    SEND_NOTIFICATION = 'send_notification'
    READ_ORDER_HISTORY = 'read_order_history'
    WRITE_CUSTOMER_PROFILE = 'write_customer_profile'  # higher privilege
    PROCESS_REFUND = 'process_refund'                  # highest risk

@dataclass
class AgentIdentity:
    agent_id: str
    allowed_permissions: set[Permission]

# Support chat agent: read-only + send notification only
SUPPORT_AGENT = AgentIdentity(
    agent_id='support-agent-v1',
    allowed_permissions={
        Permission.READ_CUSTOMER_PROFILE,
        Permission.READ_ORDER_HISTORY,
        Permission.SEND_NOTIFICATION  # can notify, but NOT write or refund
    }
)

# NOT: give every agent all permissions for convenience

class ToolGateway:
    def __init__(self, agent: AgentIdentity):
        self.agent = agent
        self.audit_log = []

    def execute_tool(self, tool_name: str, required_permission: Permission, tool_fn, **kwargs) -> dict:
        # Permission check - enforced in code, not in the LLM prompt
        if required_permission not in self.agent.allowed_permissions:
            self.log_denied(tool_name, required_permission, kwargs)
            raise PermissionError(
                f'Agent {self.agent.agent_id} does not have permission: {required_permission.value}'
            )
        
        # Execute tool
        result = tool_fn(**kwargs)
        self.log_allowed(tool_name, kwargs, result)
        return result

    def log_denied(self, tool: str, permission: Permission, args: dict):
        entry = {'type': 'DENIED', 'agent': self.agent.agent_id, 'tool': tool, 'permission': permission.value, 'args': args}
        self.audit_log.append(entry)
        print(f'SECURITY: Permission denied - {entry}')

    def log_allowed(self, tool: str, args: dict, result):
        entry = {'type': 'ALLOWED', 'agent': self.agent.agent_id, 'tool': tool, 'args': args}
        self.audit_log.append(entry)

gateway = ToolGateway(SUPPORT_AGENT)

# This will succeed (agent has READ permission)
gateway.execute_tool('read_profile', Permission.READ_CUSTOMER_PROFILE, read_customer_profile, user_id='123')

# This will raise PermissionError (agent lacks PROCESS_REFUND permission)
gateway.execute_tool('process_refund', Permission.PROCESS_REFUND, process_refund, order_id='456', amount=50.00)

HIGH_RISK_ACTIONS = {
    'delete_customer_record',
    'send_bulk_email',
    'process_refund_over_100',
    'deploy_code',
    'revoke_user_access'
}

def execute_with_confirmation(tool_name: str, tool_args: dict, agent_reasoning: str) -> dict:
    if tool_name in HIGH_RISK_ACTIONS:
        # Pause and request human approval
        approval_request = {
            'action': tool_name,
            'arguments': tool_args,
            'agent_reasoning': agent_reasoning,
            'risk_level': 'HIGH'
        }
        approval = request_human_approval(approval_request)  # blocks until human responds
        
        if not approval.approved:
            return {'status': 'rejected', 'reason': approval.rejection_reason}
        
        # Log the approval for audit trail
        log_approval(tool_name, tool_args, approved_by=approval.approver_id)
    
    # Execute only after confirmation
    return execute_tool(tool_name, **tool_args)

class ScopedCustomerTool:
    def __init__(self, current_user_id: str, session_token: str):
        self.user_id = current_user_id  # agent can only access THIS user's data
        self.token = session_token       # expires at end of session

    def get_customer_profile(self) -> dict:
        # Hardcoded to current user - agent cannot change this via prompt
        return db.query(
            'SELECT * FROM customers WHERE id = %s',
            (self.user_id,)  # parameterized, always this user's ID only
        )

    def get_order_history(self, limit: int = 10) -> list:
        # Even if the agent says 'get orders for user 999', it gets self.user_id
        return db.query(
            'SELECT * FROM orders WHERE customer_id = %s ORDER BY date DESC LIMIT %s',
            (self.user_id, min(limit, 50))  # cap limit too
        )

# Inject scoped tool into agent - cannot be overridden by prompt
def create_support_agent(user_id: str, session_token: str):
    scoped_tools = ScopedCustomerTool(user_id, session_token)
    return AgentExecutor(llm=llm, tools=[scoped_tools.get_customer_profile, scoped_tools.get_order_history])

from pydantic import BaseModel, validator, constr, confloat
import re

class SendEmailArgs(BaseModel):
    to: str
    subject: constr(max_length=200)
    body: constr(max_length=10000)

    @validator('to')
    def must_be_company_email(cls, v):
        if not re.match(r'^[^@]+@(?:yourcorp\.com|partner\.com)$', v):
            raise ValueError('Email must be sent to yourcorp.com or partner.com domains only')
        return v

class ReadFileArgs(BaseModel):
    filename: constr(pattern=r'^[a-zA-Z0-9_\-\.]+$')  # alphanumeric only, no path traversal

    @validator('filename')
    def no_parent_directory(cls, v):
        if '..' in v or '/' in v or '\\' in v:
            raise ValueError('Path traversal detected')
        return v

# Validate before execution
def validated_send_email(args_dict: dict) -> dict:
    args = SendEmailArgs(**args_dict)  # raises ValueError on invalid input
    return send_email(args.to, args.subject, args.body)

import hashlib
import json
import time

class ImmutableAuditLog:
    def __init__(self, log_backend):
        self.backend = log_backend  # e.g., CloudWatch, BigQuery, or append-only file
        self.previous_hash = '0' * 64  # genesis hash

    def record(self, agent_id: str, tool_name: str, args: dict, result, reasoning: str):
        entry = {
            'agent_id': agent_id,
            'tool': tool_name,
            'args': args,
            'result_summary': str(result)[:500],  # truncate large results
            'reasoning': reasoning[:1000],
            'timestamp': time.time(),
            'previous_hash': self.previous_hash  # chain entries like a blockchain
        }
        
        entry_json = json.dumps(entry, sort_keys=True)
        entry['hash'] = hashlib.sha256(entry_json.encode()).hexdigest()
        
        self.backend.append(entry)  # append-only, never update
        self.previous_hash = entry['hash']
        
        return entry['hash']

from collections import Counter

class AgentBehaviorMonitor:
    def __init__(self):
        self.tool_call_counts = Counter()  # tracks historical tool usage
        self.arg_length_history = {}       # tracks typical argument lengths

    def record_and_check(self, agent_id: str, tool_name: str, args: dict) -> list[str]:
        key = f'{agent_id}:{tool_name}'
        anomalies = []
        
        self.tool_call_counts[key] += 1
        
        # Flag tools never called before by this agent
        if self.tool_call_counts[key] == 1 and tool_name not in COMMON_TOOLS:
            anomalies.append(f'First time agent {agent_id} called unusual tool: {tool_name}')
        
        # Flag unusually long arguments (potential injection payload)
        total_arg_length = sum(len(str(v)) for v in args.values())
        if tool_name not in self.arg_length_history:
            self.arg_length_history[tool_name] = []
        self.arg_length_history[tool_name].append(total_arg_length)
        
        if len(self.arg_length_history[tool_name]) > 10:
            avg = sum(self.arg_length_history[tool_name]) / len(self.arg_length_history[tool_name])
            if total_arg_length > avg * 5:
                anomalies.append(f'Unusually long arguments for {tool_name}: {total_arg_length} chars (avg: {avg:.0f})')
        
        return anomalies

import secrets
import time

class SessionTokenManager:
    def __init__(self, ttl_seconds=3600):
        self.tokens = {}  # token -> (agent_id, user_id, expires_at)
        self.ttl = ttl_seconds

    def issue_token(self, agent_id: str, user_id: str) -> str:
        token = secrets.token_urlsafe(32)
        expires_at = time.time() + self.ttl
        self.tokens[token] = (agent_id, user_id, expires_at)
        return token

    def validate_token(self, token: str) -> dict | None:
        if token not in self.tokens:
            return None
        agent_id, user_id, expires_at = self.tokens[token]
        if time.time() > expires_at:
            del self.tokens[token]  # clean up expired token
            return None
        return {'agent_id': agent_id, 'user_id': user_id}

    def revoke_token(self, token: str):
        self.tokens.pop(token, None)

token_manager = SessionTokenManager(ttl_seconds=1800)  # 30-minute sessions