AI 에이전트 신뢰성 강화: Forge Guardrails와 MCP 통합 가이드

use regex::Regex;
use serde::{Deserialize, Serialize};

// 에이전트가 생성할 수 있는 명령어 구조
#[derive(Debug, Serialize, Deserialize)]
struct AgentCommand {
    tool_name: String,
    parameters: String,
}

pub struct Guardrail;

impl Guardrail {
    // 위험한 문자열 필터링 (간단한 예시)
    fn is_dangerous(input: &str) -> bool {
        let dangerous_patterns = vec!["rm -rf", "sudo", "eval", "__import__"];
        dangerous_patterns.iter().any(|&pat| input.contains(pat))
    }

    // 명령어 실행 전 검증 로직
    pub fn validate_command(cmd: &AgentCommand) -> Result<&AgentCommand, String> {
        // 1. 도구 이름 화이트리스트 확인
        let allowed_tools = vec!["blog_post", "search", "read_file"];
        if !allowed_tools.contains(&cmd.tool_name.as_str()) {
            return Err(format!("허용되지 않은 도구 사용 시도: {}", cmd.tool_name));
        }

        // 2. 파라미터 내 위험 키워드 검사
        if Self::is_dangerous(&cmd.parameters) {
            return Err("파라미터에 잠재적으로 위험한 명령어가 포함되어 있습니다.".to_string());
        }

        // 3. 안전하다면 명령어 승인
        Ok(cmd)
    }
}

// 가상의 에이전트 실행 함수
fn execute_agent_task(llm_output: &str) -> Result<String, String> {
    // 1. LLM 출력 파싱 (실제로는 JSON 파싱 등)
    // 여기서는 간단히 파싱되었다고 가정합니다.
    let command = AgentCommand {
        tool_name: "blog_post".to_string(),
        parameters: "title: 'Hello World'".to_string(),
    };

    // 2. 가드레일 통과 전
    println!("[System] LLM 응답 수신: {}", command.tool_name);

    // 3. 가드레일 검증 실행
    let safe_command = Guardrail::validate_command(&command)?;

    // 4. 검증된 명령어 실행
    println!("[System] 안전한 명령어 실행 중...");
    // 실제 도구 실행 로직 (예: 블로그 API 호출)
    Ok("게시글이 성공적으로 생성되었습니다.".to_string())
}

fn main() {
    // 정상 케이스
    match execute_agent_task("valid_response") {
        Ok(msg) => println!("성공: {}", msg),
        Err(e) => println!("차단됨: {}", e),
    }

    // 비정상 케이스 시뮬레이션
    let malicious_cmd = AgentCommand {
        tool_name: "system_shell".to_string(), // 화이트리스트에 없음
        parameters: "rm -rf /".to_string(),
    };

    match Guardrail::validate_command(&malicious_cmd) {
        Ok(_) => println!("오류: 해커가 침입했습니다!"),
        Err(e) => println!("보호됨: {}", e), // "보호됨: 허용되지 않은 도구 사용 시도: system_shell"
    }
}