FEAT : LLM Sandbox

[junkoku38]

Feature: Sandbox et validation du code LLM avant exécution

Contexte

AttackSim exécute du code généré par un LLM (Ollama) dans triage.py via _run_script_operator (ligne 617) sans aucune validation. Le script est lancé avec subprocess.Popen([sys.executable, "-S", tmp_path]) où -S ne supprime que site-packages. L'env restreint ne bloque ni os.system, ni socket, ni subprocess internes.

Un LLM compromis, jailbreaké, ou un prompt injecté dans les findings peut produire un script destructeur (rmtree, reverse shell, exfiltration réseau) exécuté avec les privilèges de l'opérateur.

De plus, exploit_generator._validate_script (ligne 987) utilise une regex sur le texte source qui est trivialement contournable : getattr(os, "system"), __builtins__["eval"], importlib.import_module("subprocess"), encodage en \x-escapes.

Description

1. Remplacer la validation regex par une analyse AST qui interdit les imports et appels dangereux.
2. Ajouter des resource.setrlimit (NPROC=0, FSIZE=0, AS=256MB) avant exécution.
3. Exécuter dans un subprocess avec start_new_session=True pour permettre le kill du process group.
4. Valider systématiquement le code LLM avant exécution dans tous les chemins (triage + exploit_generator).

Spécifications fonctionnelles

Validation AST

Modules autorisés (allow-list)

ALLOWED_MODULES = {
    # stdlib safe pour du PoC network
    "socket", "http.client", "urllib.request", "urllib.parse",
    "json", "struct", "binascii", "base64", "hashlib",
    "re", "sys", "os.path", "time", "datetime",
    "typing", "dataclasses", "enum",
    "ipaddress", "uuid", "random",
    "ssl",  # pour TLS checks (readonly)
}

Modules interdits (explicit deny)

DENIED_MODULES = {
    "subprocess", "os.system", "ctypes", "pickle", "marshal",
    "multiprocessing", "threading", "asyncio.subprocess",
    "importlib", "builtins",
}

Patterns AST interdits

Nœud AST	Raison	Exemple
Import hors allow-list	Import de module dangereux	import subprocess
ImportFrom hors allow-list	Import from dangereux	from os import system
Call vers os.system	RCE shell	os.system("rm -rf /")
Call vers subprocess.*	RCE subprocess	subprocess.run(["sh"])
Call vers eval/exec	Code injection	eval("os.system(...)")
Call vers compile	Code compilation	compile("...", "exec")
Call vers __import__	Import dynamique	__import__("subprocess")
Call vers getattr(os, ...)	Bypass deny-list	getattr(os, "system")
Subscript sur __builtins__	Bypass deny-list	__builtins__["eval"]
Attribute sur __class__/__subclasses__	MRO bypass	().__class__.__bases__[0].__subclasses__()

Algorithme

import ast

def validate_script_ast(source: str) -> tuple[bool, str]:
    try:
        tree = ast.parse(source)
    except SyntaxError as e:
        return False, f"Syntax error: {e}"

    for node in ast.walk(tree):
        # Check imports
        if isinstance(node, ast.Import):
            for alias in node.names:
                mod = alias.name.split(".")[0]
                if mod not in ALLOWED_MODULES:
                    return False, f"Import denied: {alias.name}"
        if isinstance(node, ast.ImportFrom):
            mod = (node.module or "").split(".")[0]
            if mod not in ALLOWED_MODULES:
                return False, f"Import from denied: {node.module}"
        # Check calls
        if isinstance(node, ast.Call):
            fn = _get_call_name(node)
            if fn in DENIED_CALLS:
                return False, f"Call denied: {fn}"
        # Check __builtins__ access
        if isinstance(node, ast.Subscript):
            if _is_builtins_access(node):
                return False, "Access to __builtins__ denied"
        # Check __class__/__subclasses__ MRO chain
        if isinstance(node, ast.Attribute):
            if node.attr in ("__subclasses__", "__bases__", "__mro__"):
                return False, f"MRO chain access denied: {node.attr}"
    return True, "ok"

Sandbox d'exécution

import resource

def _sandbox_execute(script_path: str, env: dict, timeout: int = 60) -> tuple[int, str]:
    def _preexec():
        # Pas de fork de processus
        resource.setrlimit(resource.RLIMIT_NPROC, (0, 0))
        # Pas d'écriture fichier
        resource.setrlimit(resource.RLIMIT_FSIZE, (0, 0))
        # Limite mémoire 256MB
        resource.setrlimit(resource.RLIMIT_AS, (256 * 1024 * 1024, 256 * 1024 * 1024))

    proc = subprocess.Popen(
        [sys.executable, "-S", script_path],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, env=env,
        start_new_session=True,  # permet os.killpg sur timeout
        preexec_fn=_preexec,
    )
    try:
        stdout, _ = proc.communicate(timeout=timeout)
        return proc.returncode, stdout
    except subprocess.TimeoutExpired:
        os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
        proc.wait()
        return -1, "[timeout — script killed]"

Spécifications techniques

Fichier : exploits/exploit_generator.py

Remplacer _validate_script (ligne 987-1009)

Remplacer la validation regex par validate_script_ast.

Valider dans generate_llm_poc (ligne 930-980)

def generate_llm_poc(...) -> str | None:
    script = _strip_fencing(raw)
    if not script: return None
    ok, msg = validate_script_ast(script)
    if not ok:
        print(f"[!] LLM script rejected: {msg}", file=sys.stderr)
        return None
    return script

Fichier : triage.py

_run_script_operator (ligne 617-676)

1. Valider le script avec validate_script_ast avant exécution
2. Utiliser _sandbox_execute au lieu de Popen direct
3. Si validation échoue : afficher le reason et proposer à l'opérateur d'éditer

def _run_script_operator(script: str, ...):
    ok, msg = validate_script_ast(script)
    if not ok:
        print(f"[!] script validation failed: {msg}", file=sys.stderr)
        if not _confirm("edit and retry?"):
            return None
        script = _open_editor(script)
        return _run_script_operator(script, ...)  # recurse
    return _sandbox_execute(tmp_path, restricted_env, timeout=60)

_llm_complete_poc (ligne ~801)

Le code généré par le LLM doit être validé avant d'être présenté à l'opérateur :

def _llm_complete_poc(finding, ...):
    raw = _ollama_generate(prompt)
    script = _strip_fencing(raw)
    ok, msg = validate_script_ast(script)
    if not ok:
        print(f"[!] LLM generated invalid script: {msg}", file=sys.stderr)
        return None
    return script

Configuration (config.py)

LLM_SCRIPT_TIMEOUT = _env("LLM_SCRIPT_TIMEOUT", 60)
LLM_SCRIPT_MAX_MEMORY = _env("LLM_SCRIPT_MAX_MEMORY", 256)  # MB
LLM_SCRIPT_ALLOW_NETWORK = _env("LLM_SCRIPT_ALLOW_NETWORK", True)  # PoC network

Tests

tests/test_script_validation.py (nouveau) :

• validate_script_ast("import socket; socket.create_connection(('host', 80))") → True
• validate_script_ast("import subprocess; subprocess.run(['rm', '-rf', '/'])") → False ("Import denied: subprocess")
• validate_script_ast("os.system('whoami')") → False ("Call denied: os.system")
• validate_script_ast("eval('1+1')") → False ("Call denied: eval")
• validate_script_ast("__import__('subprocess')") → False ("Call denied: import")
• validate_script_ast("getattr(os, 'system')('id')") → False ("Call denied: getattr(os, 'system')")
• validate_script_ast("__builtins__['eval']('1')") → False ("Access to builtins denied")
• validate_script_ast("().__class__.__bases__[0].__subclasses__()") → False ("MRO chain access denied")
• validate_script_ast("import pickle; pickle.loads(b'...')") → False ("Import denied: pickle")
• validate_script_ast("print('hello')") → True
• validate_script_ast("import os; os.path.exists('/tmp')") → True (os.path dans allow-list)
• validate_script_ast("not valid python {") → False ("Syntax error")

tests/test_script_sandbox.py (nouveau) :

• Exécuter import os; os.fork() → RLIMIT_NPROC → BlockingIOError ou OSError
• Exécuter open('/tmp/test', 'w') → RLIMIT_FSIZE → erreur
• Exécuter x = [0] * (1024 * 1024 * 512) → RLIMIT_AS → MemoryError
• Exécuter import time; time.sleep(999) → timeout → SIGKILL → return -1
• Exécuter print('ok') → return 0, stdout="ok\n"

Acceptance criteria

• Validation AST remplace la regex dans _validate_script
• 10+ patterns de bypass testés et bloqués
• generate_llm_poc valide avant de retourner
• _run_script_operator valide avant d'exécuter
• _llm_complete_poc valide avant de présenter
• resource.setrlimit appliqué (NPROC=0, FSIZE=0, AS=256MB)
• start_new_session=True permet le kill du process group
• Tests de validation passent (10+ cas)
• Tests de sandbox passent (5+ cas)
• Tests existants non cassés (triage, exploit_generator)

Références

• Issue parent : REVIEW.md V1 + V3
• CWE-94 — Code Injection
• CWE-184 — Incomplete List of Disallowed Inputs
• https://docs.python.org/3/library/ast.html
• https://docs.python.org/3/library/resource.html