AI Agent Security in Code: A Practical Guide to Defending Against Target Hijacking, Memory Poisoning, and Cascading Failures

[외부 입력] → ASI01(목표 하이재킹)
                  ↓ 오염된 목표로 메모리 쓰기
              ASI06(메모리 포이즈닝)
                  ↓ 오염된 메모리를 공유 RAG에 저장
              ASI08(계단식 실패)
                  ↓ 공유 메모리를 읽는 모든 에이전트로 전파
              [시스템 전체 장애]

import hashlib
import hmac
import re
from dataclasses import dataclass
from typing import Optional
 
SECRET_KEY = b"your-secret-key-stored-in-env"  # 실제 환경에서는 환경 변수로 관리
 
 
@dataclass
class GoalLock:
    original_goal: str
    signature: str
 
    @staticmethod
    def create(goal: str) -> "GoalLock":
        # hmac.new의 첫 번째 인자는 bytes, digestmod는 키워드 인자로 명시
        sig = hmac.new(
            SECRET_KEY, goal.encode(), digestmod=hashlib.sha256
        ).hexdigest()
        return GoalLock(original_goal=goal, signature=sig)
 
    def verify(self, current_goal: str) -> bool:
        expected = hmac.new(
            SECRET_KEY, self.original_goal.encode(), digestmod=hashlib.sha256
        ).hexdigest()
        # hmac.compare_digest: 타이밍 공격(timing attack) 방어
        if not hmac.compare_digest(self.signature, expected):
            return False  # 서명 자체가 조작됨
        return current_goal.strip() == self.original_goal.strip()
 
 
class SecureAgent:
    def __init__(self, goal: str):
        self.goal_lock = GoalLock.create(goal)
        self.current_goal = goal
 
    def process_external_input(self, external_content: str) -> str:
        """외부 입력을 처리하기 전 목표 무결성 확인"""
        sanitized = self._sanitize_input(external_content)
 
        if not self.goal_lock.verify(self.current_goal):
            raise SecurityError(
                f"Goal integrity violation detected. "
                f"Original: '{self.goal_lock.original_goal}' "
                f"Current: '{self.current_goal}'"
            )
        return sanitized
 
    def _sanitize_input(self, content: str) -> str:
        """
        다계층 입력 sanitization.
 
        완전 차단 대신 '[SUSPICIOUS_CONTENT_DETECTED]' 마킹을 사용하는 이유:
        - 완전 차단 시, LLM은 입력이 잘렸다는 사실을 모른 채 불완전한
          컨텍스트로 계획을 진행해 오히려 예측 불가능한 동작을 유발할 수 있습니다.
        - 마킹 방식은 LLM이 의심 콘텐츠의 존재를 인식하고 적절히
          무시하거나 경고를 포함한 응답을 생성하도록 유도합니다.
        - 단, 이 방식은 LLM이 마킹 자체를 무시하거나 학습 컨텍스트로
          처리할 수 있으므로, 임베딩 거리 기반 의미론적 탐지와 병행하는
          것을 권장합니다.
        """
        injection_patterns = [
            r"(?i)ignore\s+(all\s+)?previous\s+instructions?",
            r"(?i)you\s+are\s+now\s+",
            r"(?i)act\s+as\s+",
            r"(?i)forget\s+your\s+(previous\s+)?instructions?",
            r"(?i)new\s+goal\s*:",
            r"(?i)override\s+(previous\s+)?instructions?",
        ]
        for pattern in injection_patterns:
            if re.search(pattern, content):
                content = f"[SUSPICIOUS_CONTENT_DETECTED] {content}"
                break
        return content
 
 
class SecurityError(Exception):
    pass

// Node.js 14.17.0+ 환경 기준
// randomUUID와 createHash 모두 "crypto" 모듈에서 명시적으로 임포트
import { createHash, randomUUID } from "crypto";
 
interface MemoryEntry {
  id: string;
  content: string;
  // Provenance Tracking: 출처 메타데이터
  provenance: {
    source: string;       // 예: "user_upload" | "web_search" | "agent_internal"
    timestamp: number;
    agentId: string;
    trustLevel: "high" | "medium" | "low" | "untrusted";
  };
  // Temporal Decay: 만료 정보
  ttl: number;            // Unix timestamp (ms), 0이면 영구
  contentHash: string;    // 무결성 검증용 SHA-256 해시
}
 
class SecureMemoryStore {
  private partitions: Map<string, MemoryEntry[]> = new Map();
 
  // 1. Memory Partitioning — 에이전트별 격리
  private getPartition(key: string): MemoryEntry[] {
    if (!this.partitions.has(key)) {
      this.partitions.set(key, []);
    }
    return this.partitions.get(key)!;
  }
 
  // 2. Provenance Tracking — 출처 메타데이터 포함 저장
  async store(
    agentId: string,
    content: string,
    source: string,
    trustLevel: MemoryEntry["provenance"]["trustLevel"],
    ttlHours: number = 24
  ): Promise<string> {
    const entry: MemoryEntry = {
      id: randomUUID(), // "crypto" 모듈에서 임포트한 randomUUID 사용
      content,
      provenance: {
        source,
        timestamp: Date.now(),
        agentId,
        trustLevel,
      },
      ttl: ttlHours > 0 ? Date.now() + ttlHours * 3_600_000 : 0,
      contentHash: createHash("sha256").update(content).digest("hex"),
    };
 
    // 낮은 신뢰도 콘텐츠는 별도 격리 파티션에 저장
    const partitionKey =
      trustLevel === "untrusted" ? `${agentId}:quarantine` : agentId;
 
    this.getPartition(partitionKey).push(entry);
    return entry.id;
  }
 
  // 3. Context Isolation + 4. Temporal Decay — 쿼리 시 실시간 검증
  async query(
    agentId: string,
    minTrustLevel: MemoryEntry["provenance"]["trustLevel"] = "medium"
  ): Promise<MemoryEntry[]> {
    const trustHierarchy: Record<
      MemoryEntry["provenance"]["trustLevel"],
      number
    > = { high: 3, medium: 2, low: 1, untrusted: 0 };
    const minScore = trustHierarchy[minTrustLevel];
    const now = Date.now();
 
    return this.getPartition(agentId).filter((entry) => {
      // Temporal Decay: 만료된 항목 제외
      if (entry.ttl > 0 && entry.ttl < now) return false;
 
      // Context Isolation: 신뢰 수준 필터링
      if (trustHierarchy[entry.provenance.trustLevel] < minScore) return false;
 
      // 무결성 검증: 저장 후 변조 여부 확인
      const currentHash = createHash("sha256")
        .update(entry.content)
        .digest("hex");
      if (currentHash !== entry.contentHash) {
        console.error(`Memory integrity violation detected: entry ${entry.id}`);
        return false;
      }
      return true;
    });
  }
 
  // 5. Behavioral Monitoring — 이상 패턴 탐지
  async detectAnomalies(agentId: string): Promise<string[]> {
    const warnings: string[] = [];
    const partition = this.getPartition(agentId);
 
    // 단시간 내 대량 저장 시도 탐지
    const recentEntries = partition.filter(
      (e) => Date.now() - e.provenance.timestamp < 60_000
    );
    if (recentEntries.length > 50) {
      warnings.push(
        `Anomaly: ${recentEntries.length} entries written in last 60s`
      );
    }
 
    // 비신뢰 소스 비율 탐지
    const untrustedRatio =
      partition.filter((e) => e.provenance.trustLevel === "untrusted").length /
      Math.max(partition.length, 1);
 
    if (untrustedRatio > 0.3) {
      warnings.push(
        `Anomaly: ${(untrustedRatio * 100).toFixed(1)}% entries from untrusted sources`
      );
    }
 
    return warnings;
  }
}

import asyncio
import time
import json
from enum import Enum
from dataclasses import dataclass, field
from typing import Callable, Any, Optional
 
 
class CircuitState(Enum):
    CLOSED = "closed"       # 정상 동작
    OPEN = "open"           # 차단 (요청 즉시 거부)
    HALF_OPEN = "half_open" # 복구 테스트 중
 
 
@dataclass
class AgentCircuitBreaker:
    agent_id: str
    failure_threshold: int = 5       # 실패 N회 시 OPEN
    recovery_timeout: float = 30.0   # N초 후 HALF_OPEN 시도
    success_threshold: int = 2       # HALF_OPEN에서 성공 N회 시 CLOSED 복귀
 
    state: CircuitState = CircuitState.CLOSED
    failure_count: int = 0
    success_count: int = 0
    last_failure_time: float = 0.0
    safe_snapshot: Optional[dict] = None
 
    def save_snapshot(self, state: dict) -> None:
        """정상 동작 시 SafeMode 스냅샷 저장"""
        self.safe_snapshot = {
            "timestamp": time.time(),
            "agent_id": self.agent_id,
            "state": json.dumps(state),
        }
 
    def restore_from_snapshot(self) -> Optional[dict]:
        """장애 시 마지막 정상 스냅샷으로 복구"""
        if self.safe_snapshot:
            print(
                f"[RECOVERY] Agent {self.agent_id}: "
                f"restoring snapshot (saved at {self.safe_snapshot['timestamp']:.0f})"
            )
            return json.loads(self.safe_snapshot["state"])
        return None
 
    async def call(self, func: Callable, *args, **kwargs) -> Any:
        """에이전트 액션 실행 전 회로 상태 확인"""
        if self.state == CircuitState.OPEN:
            elapsed = time.time() - self.last_failure_time
            if elapsed >= self.recovery_timeout:
                self.state = CircuitState.HALF_OPEN
                self.success_count = 0
                print(f"[CIRCUIT] Agent {self.agent_id}: HALF_OPEN (recovery test)")
            else:
                restored = self.restore_from_snapshot()
                raise CircuitOpenError(
                    f"Agent {self.agent_id} is OPEN. "
                    f"Retry after {self.recovery_timeout - elapsed:.1f}s",
                    restored_state=restored,
                )
 
        try:
            result = await func(*args, **kwargs)
            self._on_success()
            return result
        except Exception:
            self._on_failure()
            raise
 
    def _on_success(self) -> None:
        self.failure_count = 0
        if self.state == CircuitState.HALF_OPEN:
            self.success_count += 1
            if self.success_count >= self.success_threshold:
                self.state = CircuitState.CLOSED
                print(f"[CIRCUIT] Agent {self.agent_id}: CLOSED (recovered)")
 
    def _on_failure(self) -> None:
        self.failure_count += 1
        self.last_failure_time = time.time()
        if self.state == CircuitState.HALF_OPEN:
            self.state = CircuitState.OPEN
        elif self.failure_count >= self.failure_threshold:
            self.state = CircuitState.OPEN
            print(
                f"[CIRCUIT] Agent {self.agent_id}: OPEN "
                f"(failures: {self.failure_count})"
            )
 
 
class MultiAgentOrchestrator:
    def __init__(self, agent_ids: list[str]):
        self.breakers = {
            aid: AgentCircuitBreaker(agent_id=aid) for aid in agent_ids
        }
 
    async def get_available_agent_decision(
        self,
        query: str,
        agent_handlers: dict[str, Callable],
        availability_threshold: float = 0.66,
    ) -> dict:
        """
        다중 에이전트 가용성 기반 결정.
 
        주의: 이 구현은 '응답한 에이전트 비율'이 임계치 이상일 때
        첫 번째 응답을 반환하는 가용성(availability) 검증입니다.
        응답 내용의 의미론적 일치도를 검사하는 진정한 합의(consensus)가
        아닙니다. 프로덕션 환경에서는 응답 간 임베딩 유사도 비교나
        다수결 로직을 추가로 구현하는 것을 권장합니다.
        """
        responses = []
        for agent_id, handler in agent_handlers.items():
            try:
                breaker = self.breakers[agent_id]
                result = await breaker.call(handler, query)
                responses.append({"agent_id": agent_id, "result": result})
            except CircuitOpenError:
                print(f"[DECISION] Agent {agent_id} skipped (circuit open)")
 
        if not responses:
            raise RuntimeError("All agents unavailable — system in SafeMode")
 
        # 가용한 에이전트 비율 확인 (의미론적 일치도 검증은 별도 구현 필요)
        availability_ratio = len(responses) / len(agent_handlers)
        if availability_ratio < availability_threshold:
            raise RuntimeError(
                f"Availability threshold not met: "
                f"only {availability_ratio:.0%} agents responded"
            )
 
        return responses[0]["result"]
 
 
class CircuitOpenError(Exception):
    def __init__(self, message: str, restored_state: Optional[dict] = None):
        super().__init__(message)
        self.restored_state = restored_state