Skip to content

README.md

Latest commit

 

History

History
335 lines (250 loc) · 9.23 KB

README.md

File metadata and controls

335 lines (250 loc) · 9.23 KB
Raw

GPU-Accelerated RL for CyberBattle Network Security

A high-performance framework combining CyberBattleSim (Microsoft's cybersecurity environment) with GPU-accelerated DQN for training agents on large-scale network penetration testing.

Key Capabilities:

  • Train agents on networks with 100-50,000 nodes
  • 100-200 environment steps/second with GPU acceleration
  • Deep Q-Network (DQN) with dense reward shaping
  • Handles complex credential systems and network topologies
  • Full Docker containerization for reproducibility

Quick Start (5 minutes)

1. Prerequisites

# Ensure you have:
docker --version          # Docker 20.10+
docker run --rm --gpus all nvidia/cuda:11.8.0-base nvidia-smi  # Verify GPU

2. Build Docker Image

cd Docker
docker-compose build

3. Start Container & Build C++ Code

docker-compose up -d pentest-dev
docker-compose exec pentest-dev bash

# Inside container:
cd /app && mkdir -p build && cd build
cmake .. && make -j$(nproc)

4. Run Training

# Quick test (100 nodes, 100 episodes)
bash /app/quick_start.sh test

# Or full training
bash /app/quick_start.sh small    # 1000 nodes, 5K episodes (~5 min)
bash /app/quick_start.sh medium   # 2000 nodes, 10K episodes (~15 min)

5. Monitor Training

# In another terminal:
nvidia-smi dmon -s puc    # Real-time GPU usage
nvtop                      # Interactive GPU monitor

How It Works

System Architecture

The system has 3 main layers:

┌─────────────────────────────────────────┐
│  Python Training (gpu_rl_engine.py)     │  ← Main training loop with DQN
├─────────────────────────────────────────┤
│  C++ Wrapper (cyberbattle_wrapper.cpp)  │  ← Bridges Python ↔ CyberBattleSim
├─────────────────────────────────────────┤
│  CyberBattleSim + GPU Kernels           │  ← Network simulation + CUDA
└─────────────────────────────────────────┘

Training Flow (One Episode)

  1. Environment Reset: Network initialized with hidden nodes
  2. Agent Action: DQN selects action (which node to attack/scan)
  3. Simulation: CyberBattleSim simulates the attack, updates network state
  4. Observation: Agent receives observation (discovered nodes, owned nodes)
  5. Reward Shaping:
    • +1.0 for discovering a new node
    • +5.0 for compromising a node
    • -0.01 per step (time penalty)
  6. Learning: Store experience in replay buffer, train on GPU mini-batches
  7. Repeat until agent owns 60% of network (success) or max steps reached

Key Components

Component Purpose
env_wrapper.py Converts CyberBattleSim to Gymnasium format + reward shaping
gpu_rl_engine.py DQN training loop: experience collection, mini-batch learning
cyberbattle_wrapper.cpp C++ bridge to CyberBattleSim (low-latency)
gpu_kernels.cu CUDA kernels for Q-value computation
CyberBattleSim Microsoft's network security simulation environment

Example Training Output

[CyberBattle] Initialized with 100 nodes, obs_dim=110, action_dim=50

Episode 1/100
[PY-STEP] step=1, shaped_reward=4.98, owned=1/3, disc=8, success=False
[PY-STEP] step=2, shaped_reward=-0.01, owned=1/3, disc=8, success=False
[PY-STEP] step=3, shaped_reward=5.99, owned=2/3, disc=9, success=False
[PY-STEP] step=4, shaped_reward=5.99, owned=3/3, disc=10, success=True ✓
Episode 1 completed: 4 steps, total_reward=16.95

Episode 2/100
[PY-STEP] step=1, shaped_reward=0.99, owned=0/3, disc=3, success=False
...

Episode 100/100
[PY-STEP] step=1, shaped_reward=5.98, owned=1/5, disc=12, success=False
Episode 100 completed: 22 steps, avg_reward=8.5, success_rate=87%

Interpreting Output:

  • owned=1/3: Agent owns 1 node, needs 3 to win
  • disc=8: 8 nodes discovered
  • success=True: Episode succeeded (owned ≥ target)
  • success_rate=87%: 87 out of 100 episodes succeeded

Results & Performance

Training Results

After training on 1000-node network for 5000 episodes:

Metric Result
Episode Length 20-50 steps (average)
Success Rate 70-90% (owns 60% nodes)
Training Time 5-10 minutes (A100)
Steps/Second 100-200 (GPU accelerated)
Convergence Episode 500-1000

Scalability

Network Size Episodes Time GPU Memory
100 nodes 100 30 sec 2GB
1K nodes 5K 5 min 4GB
5K nodes 10K 20 min 8GB
10K nodes 20K 1 hour 12GB
50K nodes SIR only 15 sec 6GB

Agent Behavior

The trained agent learns to:

  • Scan progressively: Discover nodes methodically rather than randomly
  • Exploit credentials: Use discovered credentials to own nodes faster
  • Prioritize targets: Focus on high-value nodes (more connections)
  • Adapt strategy: Change approach based on network topology

Example learned strategy:

Early game (episodes 1-100): Explore randomly
Mid game (episodes 100-500): Focus on discovering nodes with valid credentials
Late game (episodes 500+): Exploit multi-hop attacks to own isolated subnets

Configuration

Training Parameters

Modify in quick_start.sh or directly in training script:

# Network config
NODES=1000           # Network size (100-50000)
EPISODES=5000        # Total training episodes

# RL config
BATCH_SIZE=128       # GPU batch size
LEARNING_RATE=0.0001
GAMMA=0.99           # Discount factor

# Reward shaping
DISCOVERY_BONUS=1.0  # Per discovered node
COMPROMISE_BONUS=5.0 # Per owned node
STEP_PENALTY=-0.01   # Per step

Hardware Requirements

  • GPU: NVIDIA (A10, A100, RTX series) with 6GB+ VRAM
  • CPU: 8+ cores
  • RAM: 16GB minimum, 32GB recommended
  • Storage: 20GB free

Testing

Run Tests

cd /app/build
ctest -V                    # Run all tests
ctest -R "GPUKernels" -V   # GPU kernel tests
ctest -R "CyberBattle1k" -V # Production-scale test

Test Coverage

  • ✓ GPU kernel correctness (matrix ops, activations)
  • ✓ CUDA environment simulation (SIR propagation)
  • ✓ Integration tests (GPU + environment)
  • ✓ Large-scale tests (50K nodes)
  • ✓ End-to-end training (1K nodes, 100 episodes)

Common Commands

# Build inside container
cd /app/build && cmake .. && make -j$(nproc)

# Run quick training
bash /app/quick_start.sh test

# Run tests
ctest --output-on-failure

# Monitor GPU
nvidia-smi dmon -s puc

# View logs
tail -f /app/logs/training.log

# Check results
cat /app/outputs/test_summary.txt

Project Structure

.
├── Docker/                  # Docker config (Dockerfile, docker-compose.yml)
├── src/
│   ├── cpu/environment/     # CyberBattle C++ wrapper
│   ├── gpu/                 # GPU kernels & policy network
│   └── cuda_env/            # Network simulation
├── python/
│   ├── cyberbattle_env/
│   │   ├── env_wrapper.py          # Gymnasium interface + reward shaping
│   │   └── gpu_rl_engine.py        # DQN training loop
│   └── train_large_scale.py
├── tests/                   # C++ & Python tests
├── CMakeLists.txt           # C++ build config
├── quick_start.sh           # Quick training script
└── README.md                # This file

Key Libraries

Library Version Purpose
CyberBattleSim latest Network security simulation
Gymnasium 0.29.1 RL environment API
CUDA 11.8 GPU computing
pybind11 2.11.1 C++ ↔ Python bindings
CMake 3.18+ Build system
NumPy 1.24.3 Numerical computing

Troubleshooting

Problem: GPU not detected in container

docker run --rm --gpus all nvidia/cuda:11.8.0-base nvidia-smi
# Should show GPU info. If not, reinstall nvidia-docker2

Problem: Out of memory

# Reduce batch size or network size
# In quick_start.sh or training script:
BATCH_SIZE=64    # from 128
NODES=500        # from 1000

Problem: Python module not found

export PYTHONPATH=/app/python:$PYTHONPATH

Problem: CMake not finding Python

# Inside container, ensure Python dev headers are installed
apt-get install -y python3.10-dev pybind11-dev

References

Performance Tips

  1. GPU utilization: Monitor with nvidia-smi dmon -s puc
  2. Batch size: Increase for better GPU utilization (watch memory)
  3. Network size: Start small (100 nodes) for testing, scale up gradually
  4. Training stability: If rewards are erratic, reduce learning rate
  5. Convergence: Success rate should improve after 500-1000 episodes

Version: 1.0.0 Last Updated: December 2024