LLM을 이용한 서비스 취약점 분석 자동화 #2

@mcp.tool()
def find_references(
    symbol_or_pattern: str,
    dir: Optional[str] = None,
    max_results: int = 50
) -> FindReferencesResponse:
    """심볼이나 패턴의 참조를 ripgrep으로 검색"""
    refs, total = find_references_with_ripgrep(
        symbol_or_pattern, dir, max_results
    )

    return {
        "references": [
            {"file": r.file, "line": r.line, "snippet": r.snippet}
            for r in refs
        ],
        "total": total,
        "truncated": total > len(refs)
    }

@mcp.tool()
def read_definition(
    symbol: str,
    file: Optional[str] = None,
    language: Optional[str] = None,
    include_body: bool = True
) -> DefinitionResponse:
    """심볼의 정의를 ctags 인덱스에서 찾아 반환"""
    entries = ctags_index.find_definitions(
        symbol, file=file, language=language
    )

    results = []
    for e in entries[:10]:
        item = {
            "file": e.file,
            "line": e.line,
            "kind": e.kind,
            "language": e.language,
            "signature": e.signature,
            "scope": e.scope
        }

        if include_body:
            bounds = detect_function_bounds(p, e.line)  # tree-sitter
            if bounds:
                item["body"] = "\\n".join(lines[bounds[0]-1:bounds[1]])

        results.append(item)

    return {"symbol"

@mcp.tool()
def read_source(
    path: str,
    line: int,
    before: int = 40,
    after: int = 40
) -> ReadSourceResponse:
    """지정된 라인 주변의 소스 코드를 읽어옴"""
    p = _find_file_in_repo(path)
    lines = p.read_text(errors="ignore").splitlines()

    start = max(1, line - before)
    end = min(len(lines), line + after)
    text = "\\n".join(lines[start-1:end])

    return {
        "path"

@mcp.tool()
def get_project_structure(
    max_depth: int = 5,
    include_file_sizes: bool = False
) -> ProjectStructureResponse:
    """인덱싱된 프로젝트의 디렉터리 구조 정보를 반환"""
    structure = _build_directory_structure(REPO_ROOT, max_depth)

    return {
        "total_files"

# Semgrep Rule - kotlin-anycall-taint.yaml
rules:
  - id: kt-spring-anycall-taint
    languages: [kotlin]
    severity: INFO
    message: Tainted Spring input flows into a call
    mode: taint

    pattern-sources:
      # @RequestParam / @PathVariable / @RequestHeader
      - patterns:
          - pattern-inside: |
              fun $F(..., @RequestParam $X: $T, ...) { ... }
          - pattern: $X

      - patterns:
          - pattern-inside: |
              fun $F(..., @PathVariable $X: $T, ...) { ... }
          - pattern: $X

      - patterns:
          - pattern-inside: |
              fun $F(..., @RequestHeader($H) $X: $T, ...) { ... }
          - pattern: $X

      # @RequestBody — DTO 필드 읽기를 source로 간주
      - patterns:
          - pattern-inside: |
              fun $F(..., @RequestBody $DTO: $T, ...) { ... }
          - pattern: $DTO.$FIELD

      # @RequestPart / @ModelAttribute / @RequestAttribute
      - patterns:
          - pattern-inside: |
              fun $F(..., @RequestPart $X: $T, ...) { ... }
          - pattern: $X

      - patterns:
          - pattern-inside: |
              fun $F(..., @ModelAttribute $X: $T, ...) { ... }
          - pattern: $X

      - patterns:
          - pattern-inside: |
              fun $F(..., @RequestAttribute $X: $T, ...) { ... }
          - pattern: $X

    pattern-sinks:
      - pattern: $F(...)
      - pattern: $OBJ.$M(...)

$ cd target-kotlin-project
$ semgrep scan --config kotlin-anycall-taint.yaml . \\
    --timeout=60 --sarif --output kotlin-semgrep-output.sarif

// kotlin-semgrep-output.sarif
{
  "version": "2.1.0",
  "runs": [
    {
      "invocations": [...],
      "results": [
        {
          "locations": [
            {
              "physicalLocation": {
                "artifactLocation": {
                  "uri": "src/main/kotlin/com/example/demo/UserController.kt",
                  "uriBaseId": "%SRCROOT%"
                },
                "region": {
                  "endColumn": 40,
                  "endLine": 16,
                  "snippet": {
                    "text": " return userService.findUser(id)"
                  },
                  "startColumn": 16,
                  "startLine": 16
                }
              }
            }
          ],
          "message": {
            "text": "Tainted Spring input flows into a call"
          },
          "properties": {},
          "ruleId": "kt-spring-anycall-taint"
        },
        ...
      ],
      ...
    }
  ],
  "$schema": "<https://docs.oasis-open.org/sarif/sarif/v2.1.0/os/schemas/sarif-schema-2.1.0.json>"
}

[
  {
    "id": 1,
    "file_path": "src/main/kotlin/com/example/demo/BrandStorePublicController.kt",
    "line_num": 38,
    "code_snippet": "val result = fileUploadApplication.createPresignedUploadUrl(\\n "
  },
  {
    "id": 2,
    "file_path": "src/main/kotlin/com/example/demo/ProductBrandAdminController.kt",
    "line_num": 39,
    "code_snippet": "val brand = productBrandService.findById(brandId)\\n"
  },
  {
    "id": 3,
    "file_path": "src/main/kotlin/com/example/demo/ProductBrandInternalController.kt",
    "line_num": 40,
    "code_snippet": "val productBrand = productBrandService.getOrCreateByName(request.brandName)\\n"
  },
  {
    "id": 4,
    "file_path": "src/main/kotlin/com/example/demo/ProductCategoryInternalController.kt",
    "line_num": 85,
    "code_snippet": "val result = productCategoryApplication.findCategoryPathsByIdIn(ids)\\n"
  },
  {
    "id": 5,
    "file_path": "src/main/kotlin/com/example/demo/ProductCategoryInternalController.kt",
    "line_num": 86,
    "code_snippet": "val result = productCategoryApplication.findCategoryPathByGroupId(tacaId)\\n"
  },
  ...
]

[
  {
    "id": 1,
    "file_path": "src/main/kotlin/com/example/demo/BrandStorePublicController.kt",
    "line_num": "38-40",
    "code_snippet": "val result = fileUploadApplication.createPresignedUploadUrl(\\n val brand = productBrandService.findById(brandId)\\n val productBrand = productBrandService.getOrCreateByName(request.brandName)\\n"
  },
  {
    "id": 2,
    "file_path": "src/main/kotlin/com/example/demo/ProductCategoryInternalController.kt",
    "line_num": "85-86",
    "code_snippet": "val result = productCategoryApplication.findCategoryPathsByIdIn(ids)\\n val result = productCategoryApplication.findCategoryPathByGroupId(tacaId)\\n"
  },
  ...
]

# GUIDE_LINE
1. 가능한 SourceCode-Browse MCP를 사용하지 말고, 주어진 코드 스니펫을 보고 판단할 것
2. 취약점 가능성을 너무 엄격하게 평가하지 말고, 어느 정도의 발생 가능성이 있으면 분석 경로에 추가할 것

from pydantic import BaseModel, Field
from typing import Union, List

class PathItem(BaseModel):
    """Individual path item structure"""
    id: int = Field(description="PAYLOAD의 id와 동일한 값 (primary key)")
    file_path: str = Field(description="파일 경로")
    line_num: Union[int, str] = Field(description="줄 번호 (정수 또는 'A-B' 형식의 범위)")
    code_snippet: str = Field(description="코드 조각")
    reason: str = Field(description="분류 이유")

class DiscoveryResult(BaseModel):
    """Discovery result structure"""
    include_path: List[PathItem] = Field(
        default_factory=list,
        description="분석이 필요한 엔드포인트 목록 (취약점 발생 가능성이 높음)"
    )

    exclude_path: List[PathItem] = Field(
        default_factory=list,
        description="분석을 건너뛸 수 있는 엔드포인트 목록 (취약점 발생 가능성이 낮음)"
    )

import instructor

payload = [
    {"id": 1, "file_path": "app/routes/user.py", "line_num": 42, "code": "..."},
    {"id": 2, "file_path": "app/routes/auth.py", "line_num": "10-20", "code": "..."},
]

client = instructor.from_openai(
    OpenAI(base_url="<http://localhost:11434/v1>", api_key="ollama"),
    mode=instructor.Mode.JSON,
)

result: DiscoveryResult = client.chat.completions.create(
    model="qwen3:30b",
    response_model=DiscoveryResult,
    max_retries=3,  # 최대 3번 retry
    messages=[
        {
            "role": "system",
            "content": "당신은 보안 분석가입니다. 주어진 코드 경로를 취약점 가능성에 따라 분류하세요."
        },
        {
            "role": "user",
            "content": f"다음 경로들을 분석하세요:\\n{payload}"
        }
    ]
)