ai-ml

Reinforcement Learning para Generación de Exploits

Hace unos meses, durante una final de CTF, nuestro equipo se quedó tres horas atascado en un binario de pwn. Stack canary, NX, PIE parcial, un overflow de exactamente 72 bytes en una función poco evidente. Lo intentamos todo. Nada. A las 03:14 de la madrugada, uno de los nuestros lanzó casi en broma un script de cinco líneas que mutaba el input con AFL-like bitflips y reenviaba sin pensar. Veinte minutos después, el script ganó la flag por pura suerte estadística.

Esa noche nos fuimos a casa con una pregunta dándonos vueltas: ¿y si en lugar de un proceso tonto tuviésemos un proceso que aprendiera de cada intento fallido?

Por qué RL encaja con control flow hijacking

El espacio de búsqueda de un exploit es brutal pero no uniforme. La mayoría de inputs producen el mismo comportamiento, y solo una región estrecha lleva al estado de máquina que quieres: registro de instrucción apuntando a memoria que controlas. Lo que diferencia ese subespacio es gradiente de información: un input que se acerca a corromper RIP exhibe efectos secundarios observables — registros parcialmente sobrescritos, valores de stack que coinciden con bytes del input, saltos a direcciones inválidas pero cerca de bytes controlados.

Eso es exactamente lo que las funciones de reward bien diseñadas pueden capturar.

El estado del arte

Avgerinos y Brumley publicaron en 2011 el primer sistema end-to-end con AEG (NDSS 2011) sobre ejecución simbólica. Mayhem (Cha et al., S&P 2012) escaló el análisis. Ambos son cuasi-deterministas pero pagan el precio clásico de symbolic execution: path explosion y solvers SMT atragantados.

Böttinger formalizó el fuzzing como MDP con Deep Reinforcement Fuzzing (2018) usando DQN. PwnGPT (Shao et al., ACL 2025) reporta subir tasa de exploit generation de 26.3% a 57.9% con o1-preview. Y PPO (Schulman et al., 2017) + RLHF (Ouyang et al., 2022) maduraron lo suficiente como para que entrenar un agente sobre entorno binario sea una conversación seria.

Nuestro setup

Entorno tipo Gym sobre ptrace y un sidecar de eBPF. La observación es un tensor denso con:

  • Estado de los 16 registros generales (x86_64), normalizado.
  • Ventana de stack alrededor de RSP (256 bytes) como bytes crudos.
  • Flags de control: NX, canary_present, RIP_corrupted, register_controlled[].
  • Hash de la última traza de cobertura tomada de eBPF (uprobe sobre cada bloque básico).

El espacio de acciones es discreto sobre bytes del input. La recompensa combina cobertura nueva (positiva, decreciente), explotabilidad heurística y penalización por longitud.

import gym
from gwaihir.tools.rl_env import BinaryPwnEnv

class CTFPwnEnv(BinaryPwnEnv):
    def step(self, action):
        mutated = self.apply_action(self.current_input, action)
        trace = self.run_with_ebpf(mutated)

        new_bbs = len(trace.basic_blocks - self.coverage_seen)
        cov_reward = np.log1p(new_bbs) * 0.3
        self.coverage_seen |= trace.basic_blocks

        expl = 0.0
        if trace.rip_controlled_bytes > 0:
            expl += 5.0 + 0.1 * trace.rip_controlled_bytes
        if trace.canary_leaked: expl += 2.5
        if trace.tainted_syscall_args: expl += 1.5

        reward = cov_reward + expl - 0.001 * len(mutated)
        done = trace.rip_fully_controlled or self.steps > 4096
        return self.observe(trace), reward, done, {"trace": trace}

PPO vs DQN

Empezamos con DQN. En cuanto el espacio de acciones creció, DQN se volvió frágil: el target network divergía cada vez que la recompensa de exploitability disparaba un valor enorme y aislado. PPO, con su clipping del ratio de policy, absorbe esos picos sin colapsar. PPO con GAE (λ=0.95), entropía coef 0.01 y batches de 4096 nos dio entrenamientos estables sobre 24 horas en una RTX 4090.

Detalle que no se cuenta a menudo: la mayor parte del tiempo no se va en gradientes, se va en ejecutar el binario. Cada step requiere fork, ptrace, recolección de eBPF, parse. Paralelizamos 32 envs por GPU con SubprocVecEnv y aún así el bottleneck era I/O.

Resultados sobre 10 binarios CTF

Batería: dos babys, tres con canary leak required, dos con format string, dos con ROP chain mínima y uno con heap (UAF). El agente, entrenado desde cero por binario, resolvió 7 de 10. De los tres que falló: el de heap (esperable), uno de ROP que requería un gadget no presente, y un format string donde la política se quedó en mínimo local leakeando direcciones eternamente.

Comparado contra angr puro: RL ganó en 5, empató en 2, perdió en 3 (angr resolvió antes heap y format-write). No es herramienta superior, es complementaria.

Limitaciones brutales

  • Una política por binario. No transferimos conocimiento entre retos.
  • Tiempos de entrenamiento de horas para retos que un humano sénior resuelve en minutos.
  • Reward shaping frágil. Custom allocator o mitigaciones no estándar rompen las heurísticas.
  • Symbolic execution sigue ganando cuando el path es deterministicamente derivable.

El trade-off real: cobertura amplia barata (RL) vs razonamiento profundo caro (symbolic). Los corremos en paralelo dentro de Gwaihir.

Qué viene después

Tres frentes: (a) pretraining sobre exploit-db y writeups CTF, (b) reward learning estilo RLHF, (c) integrar el agente como subherramienta de un LLM coordinador.

Por ahora seguimos entrenando. Y la próxima vez que un equipo se atasque tres horas a las tres de la mañana, queremos que sea el agente quien encuentre la flag — no la suerte.

Bibliografía

  1. Avgerinos, T. et al. (2011). AEG: Automatic Exploit Generation. NDSS.
  2. Cha, S.K. et al. (2012). Unleashing Mayhem. IEEE S&P.
  3. Böttinger, K. et al. (2018). Deep Reinforcement Fuzzing. IEEE SPW. arXiv:1801.04589.
  4. Schulman, J. et al. (2017). PPO. arXiv:1707.06347.
  5. Ouyang, L. et al. (2022). InstructGPT/RLHF. NeurIPS.
  6. Shao, Y. et al. (2025). PwnGPT. ACL.