pentest-r1-autonomous-penetration-testing

# Pentest-R1 Autonomous Penetration Testing > Skill by [ara.so](https://ara.so) — Security Skills collection Pentest-R1 is a two-stage reinforcement learning framework that trains Large Language Models for autonomous penetration testing. It combines offline RL on expert walkthroughs with online RL in interactive CTF environments to develop robust attack reasoning capabilities. ## Installation ### Prerequisites - Python 3.11.11 - Docker (for Stage 2 and reproducible environments) - NVIDIA Container Toolkit (for GPU support) - CUDA 12.4 runtime (optional, for GPU acceleration) ### Basic Setup ```bash git clone https://github.com/KHenryAegis/Pentest-R1.git cd Pentest-R1 pip install -r requirements.txt ``` ### Docker Environment (Recommended for Reproducibility) ```bash # Build the reproducible research environment source setup-docker.sh # Run container with optimized cache mounting docker run --rm -it \ --name pentest-r1 \ -v "$(pwd)":/root/Pentest-R1 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ -v ~/.cache/triton:/root/.cache/triton \ -v ~/.cache/torch_extensions:/root/.cache/torch_extensions \ -w /root/Pentest-R1 \ --gpus all \ --net=host \ pentest-r1:ubuntu22.04 ``` ## Core Training Pipeline ### Stage 1: Offline Reinforcement Learning Stage 1 trains the base LLM on a curated dataset of 500+ real-world expert penetration testing walkthroughs. ```bash python grpo_stage1.py ``` **Key configuration in `grpo_stage1.py`:** ```python from transformers import AutoTokenizer, AutoModelForCausalLM from unsloth import FastLanguageModel import torch # Model configuration model_name = "unsloth/Meta-Llama-3.1-8B-Instruct" max_seq_length = 4096 dtype = None # Auto-detect load_in_4bit = True # Use 4-bit quantization # Load model with unsloth optimization model, tokenizer = FastLanguageModel.from_pretrained( model_name=model_name, max_seq_length=max_seq_length, dtype=dtype, load_in_4bit=load_in_4bit, ) # PEFT configuration for efficient training from peft import LoraConfig, get_peft_model peft_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM" ) model = get_peft_model(model, peft_config) ``` ### Stage 2: Online Reinforcement Learning Stage 2 fine-tunes the model in interactive CTF environments with real-time feedback. **Setup CTF environment:** ```bash cd train_ctf_env docker build -t intercode-ctf . cd .. ``` **Run Stage 2 training:** ```bash python grpo_multi_turn_stage2.py ``` ## Key Components ### Data Loading and Preprocessing ```python from datasets import load_dataset # Load expert walkthrough dataset dataset = load_dataset("json", data_files="path/to/expert_walkthroughs.jsonl") # Example dataset format # { # "challenge": "SQL Injection in login form", # "steps": [ # {"action": "reconnaissance", "command": "sqlmap -u http://target/login", "reasoning": "..."}, # {"action": "exploit", "command": "sqlmap --dump", "reasoning": "..."} # ], # "flag": "CTF{...}" # } def preprocess_function(examples): """Format data for training""" prompts = [] responses = [] for challenge, steps in zip(examples["challenge"], examples["steps"]): prompt = f"Challenge: {challenge}\nWhat are the steps to solve this?" response = "\n".join([ f"Step {i+1}: {step['reasoning']}\nCommand: {step['command']}" for i, step in enumerate(steps) ]) prompts.append(prompt) responses.append(response) return {"prompt": prompts, "response": responses} tokenized_dataset = dataset.map(preprocess_function, batched=True) ``` ### Reward Model Configuration ```python class PentestRewardModel: """Reward model for evaluating penetration testing actions""" def __init__(self): self.success_reward = 1.0 self.partial_reward = 0.5 self.failure_penalty = -0.1 def calculate_reward(self, action, environment_feedback): """Calculate reward based on action outcome""" if "flag" in environment_feedback.lower(): return self.success_reward elif "error" in environment_feedback.lower(): return self.failure_penalty elif "progress" in environment_feedback.lower(): return self.partial_reward return 0.0 reward_model = PentestRewardModel() ``` ### Interacting with CTF Environment ```python import docker class CTFEnvironment: """Wrapper for InterCode-CTF Docker environment""" def __init__(self, image_name="intercode-ctf"): self.client = docker.from_env() self.image_name = image_name self.container = None def start(self, challenge_id): """Start a CTF challenge container""" self.container = self.client.containers.run( self.image_name, detach=True, environment={"CHALLENGE_ID": challenge_id}, network_mode="host", remove=True ) return self.container def execute_command(self, command): """Execute a command in the container""" if not self.container: raise RuntimeError("Container not started") exec_result = self.container.exec_run(command) return { "stdout": exec_result.output.decode('utf-8'), "exit_code": exec_result.exit_code } def cleanup(self): """Stop and remove container""" if self.container: self.container.stop() self.container = None # Usage example env = CTFEnvironment() env.start(challenge_id="sql_injection_001") result = env.execute_command("sqlmap -u http://localhost/login --batch") print(result["stdout"]) env.cleanup() ``` ## Training Configuration ### GRPO (Group Relative Policy Optimization) Settings ```python from transformers import TrainingArguments training_args = TrainingArguments( output_dir="./pentest-r1-checkpoints", num_train_epochs=3, per_device_train_batch_size=4, gradient_accumulation_steps=4, learning_rate=5e-5, warmup_steps=100, logging_steps=10, save_steps=500, save_total_limit=3, fp16=True, # Mixed precision training report_to="wandb", # Optional: integration with Weights & Biases remove_unused_columns=False, ) ``` ### Environment Variables ```python import os # Model and training configuration os.environ["HF_TOKEN"] = os.getenv("HUGGINGFACE_TOKEN") # For model downloads os.environ["WANDB_API_KEY"] = os.getenv("WANDB_API_KEY") # For logging os.environ["CUDA_VISIBLE_DEVICES"] = "0" # GPU selection # Cache directories (mounted in Docker) os.environ["HF_HOME"] = "/root/.cache/huggingface" os.environ["TRITON_CACHE_DIR"] = "/root/.cache/triton" ``` ## Common Patterns ### Multi-Turn Reasoning ```python class MultiTurnAgent: """Agent for multi-turn penetration testing reasoning""" def __init__(self, model, tokenizer): self.model = model self.tokenizer = tokenizer self.conversation_history = [] def generate_action(self, observation): """Generate next action based on current observation""" # Build prompt with conversation history prompt = self._build_prompt(observation) inputs = self.tokenizer(prompt, return_tensors="pt") outputs = self.model.generate( **inputs, max_new_tokens=512, temperature=0.7, do_sample=True, top_p=0.95 ) action = self.tokenizer.decode(outputs[0], skip_special_tokens=True) self.conversation_history.append({ "observation": observation, "action": action }) return action def _build_prompt(self, observation): """Build prompt with conversation history""" prompt = "You are a penetration testing expert. Analyze and exploit:\n\n" for turn in self.conversation_history[-3:]: # Last 3 turns prompt += f"Observation: {turn['observation']}\n" prompt += f"Action: {turn['action']}\n\n" prompt += f"Current Observation: {observation}\n" prompt += "Next Action:" return prompt ``` ### Evaluation Loop ```python def evaluate_on_ctf_challenges(model, tokenizer, challenge_set): """Evaluate model on a set of CTF challenges""" agent = MultiTurnAgent(model, tokenizer) results = [] for challenge in challenge_set: env = CTFEnvironment() env.start(challenge["id"]) solved = False max_turns = 20 for turn in range(max_turns): # Get current state observation = env.execute_command("cat /challenge/description.txt") # Generate action action = agent.generate_action(observation["stdout"]) # Execute action result = env.execute_command(action) # Check for success if "CTF{" in result["stdout"]: solved = True break results.append({ "challenge_id": challenge["id"], "solved": solved, "turns": turn + 1 }) env.cleanup() return results ``` ## Troubleshooting ### CUDA Out of Memory ```python # Use gradient checkpointing model.gradient_checkpointing_enable() # Reduce batch size training_args.per_device_train_batch_size = 1 training_args.gradient_accumulation_steps = 8 # Use 8-bit or 4-bit quantization load_in_4bit = True ``` ### Docker Container Issues ```bash # Check container logs docker logs intercode-ctf # Verify network connectivity docker run --rm --net=host intercode-ctf ping -c 4 localhost # Rebuild with no cache if issues persist docker build --no-cache -t intercode-ctf train_ctf_env/ ``` ### Unsloth Version Compatibility The research used `unsloth==2025.5.10`. If unavailable, the closest match is specified in `requirements.txt`: ```txt unsloth_zoo==2025.5.11 unsloth @ git+https://github.com/unslothai/unsloth.git@45f26cda996ec0b9a2e28cb18a03251095aa29e8 ``` ### Model Loading Errors ```python # If model fails to load, try without quantization first model, tokenizer = FastLanguageModel.from_pretrained( model_name=model_name, max_seq_length=max_seq_length, dtype=torch.float16, load_in_4bit=False, ) # Clear cache if needed import shutil shutil.rmtree(os.path.expanduser("~/.cache/huggingface"), ignore_errors=True) ``` ## Performance Optimization ### Cache Mounting for Faster Training Always mount HuggingFace, Triton, and PyTorch caches when using Docker: ```bash -v ~/.cache/huggingface:/root/.cache/huggingface \ -v ~/.cache/triton:/root/.cache/triton \ -v ~/.cache/torch_extensions:/root/.cache/torch_extensions ``` ### Distributed Training ```python # For multi-GPU training training_args.ddp_find_unused_parameters = False training_args.local_rank = int(os.environ.get("LOCAL_RANK", -1)) # Launch with torchrun # torchrun --nproc_per_node=4 grpo_stage1.py ``` ## References - Paper: [arXiv:2508.07382](https://arxiv.org/abs/2508.07382) - Base model: Meta-Llama-3.1-8B-Instruct - Framework: Unsloth for optimized fine-tuning - Environment: InterCode-CTF for interactive training