정규표현식(Regex)에서 AST까지: AI 에이전트 코드를 위한 오염 추적(Taint Tracking) 구축하기

AgentGuard v0.5.0은 AST 기반의 오염 추적(Taint Tracking) 기능을 탑재했습니다. 이 포스트에서는 이것이 어떻게 작동하는지, 그리고 왜 중요한지에 대해 설명합니다.

정규표현식(Regex)의 한계

정규표현식(Regex)은 명백한 패턴을 잡아냅니다:

prompt = f"You are helpful. {user_input}"

정규표현식 규칙은 {user_input}이 포함된 f"..."를 보고 플래그를 표시합니다. 그것으로 끝입니다.

하지만 정규표현식은 다음을 추적할 수 없습니다:

query = request.json.get("query")
processed = query.strip().upper()
template = "Answer: {q}"
...

오염(Taint)은 다음과 같이 흐릅니다: request.json -> query -> processed -> template.format() -> prompt -> openai 호출. 총 4단계의 홉(hop)이 발생합니다. 정규표현식은 각 라인을 독립적으로 보기 때문에 이들을 연결할 수 없습니다.

AST의 구원

Python의 ast 모듈은 소스 코드를 구문 트리(Syntax Tree)로 파싱합니다. 우리는 이 트리를 순회하며 데이터가 어떻게 흐르는지 추적할 수 있습니다.

1단계: 소스(Sources) 식별

"소스(Source)"는 신뢰할 수 없는 데이터를 생성하는 모든 표현식입니다:

SOURCE_PATTERNS = {
    "user_input", "user_msg", "user_message",
    "request", "req", "query", "message", "msg",
...

여기에 속성 접근(Attribute access) 패턴도 포함됩니다: request.args.get("q"), request.json["key"], input().

AST 관점에서, 우리는 ast.Name 노드를 소스 집합과 대조하여 확인하고, ast.Call 노드를 통해 request.args.get 패턴을 확인합니다.

2단계: 전파(Propagation) 추적

소스가 변수에 할당되면, 해당 변수는 오염(Tainted)됩니다:

user_input = request.args.get("q")  # 이제 user_input은 오염되었습니다

하지만 오염은 다음을 통해서도 전파됩니다:

• 메서드 호출 (Method calls): processed = user_input.strip() -- processed는 여전히 오염된 상태입니다.
• F-strings: prompt = f"Hello {user_input}" -- prompt는 오염됩니다.
• .format(): prompt = template.format(q=query) -- query가 오염되어 있다면 prompt도 오염됩니다.
• 문자열 연결 (String concatenation): prompt = "Hello " + user_input -- prompt는 오염됩니다.
• 리스트/딕셔너리 생성 (List/dict construction): messages = [{"role": "user", "content": user_input}] -- messages는 오염됩니다.

트래커는 할당(assignment)을 순서대로 탐색하며 tainted_vars 딕셔너리를 유지합니다. x = tainted_expr을 발견하면 x를 딕셔너리에 추가합니다. x = safe_expr을 발견하면 x를 제거합니다.

3단계: 싱크(Sinks) 식별

"싱크 (sink)"는 오염된 데이터가 LLM에 도달하는 지점입니다:

• 변수 할당 (Variable assignment): prompt = <tainted> 또는 messages = [<tainted>]
• 함수 호출 (Function call): openai.chat.completions.create(messages=<tainted>)

트래커가 오염된 표현식(tainted expression)이 싱크에 도달하는 것을 발견하면 탐지 결과(finding)를 발생시킵니다.

4단계: 새니타이저 (Sanitizers)

모든 변환이 오염(taint)을 보존하는 것은 아닙니다. 일부는 명시적으로 데이터를 안전하게 만듭니다:

safe = str(user_input)[:100]  # 잘림(truncated), 문자열로 형변환(cast)

트래커는 str(), int(), float(), len() 및 명시적인 이스케이프 함수(escape functions)를 새니타이저로 취급합니다. 데이터가 새니타이저를 통과하면 오염이 제거됩니다.

탐지 가능한 항목 (정규표현식(Regex)이 할 수 없는 것)

# 다중 홉 흐름 (Multi-hop flow) -- 4번의 변수 할당
user_input = request.args.get("message")
processed = user_input.strip()
...

(정상적으로) 플래그를 세우지 않는 항목

# 새니타이징된 입력 (Sanitized input)
user_input = request.args.get("q")
safe_input = str(user_input)[:100]
...

한계점

이것은 v0.5.0 -- 첫 번째 반복 버전입니다. 알려진 공백(gaps)은 다음과 같습니다:

1. Python 전용. JavaScript/TypeScript AST 지원은 로드맵에 있습니다.
2. 파일 내부(Intra-file) 전용. 오염이 파일 경계를 넘지 않습니다 (아직 절차 간 분석(interprocedural analysis)은 지원하지 않음).
3. 제어 흐름(Control flow) 미지원. If/else 분기가 별도로 추적되지 않습니다.
4. 보수적인 새니타이저. str()을 새니타이저로 취급하지만, str(user_input) 하나만으로는 모든 컨텍스트에서 입력을 안전하게 만들 수 없습니다.

아키텍처 (The Architecture)

소스 코드 (Source code)
    |
    v
...

사용해 보기

pip install --upgrade dfx-agentguard
agentguard src/ --format text

오염 추적 규칙(ASI01-TAINT-TRACK)은 기존의 정규표현식(regex) 규칙과 함께 실행됩니다. 두 계층은 함께 작동합니다: 속도를 위한 정규표현식(regex), 정밀도를 위한 AST.

AgentGuard는 MIT 라이선스를 따릅니다. v0.5.0 버전에는 38개의 테스트와 100% 탐지율을 기록한 32개 샘플 벤치마크(benchmark)가 포함되어 있습니다.