Spaces:

CystronCode
/

cannon-and-wall

Sleeping

App Files Files Community

cannon-and-wall / README.md

CystronCode

Update README.md

a9083ad verified 2 months ago

preview code

Raw

History Blame Contribute Delete

10.1 kB

metadata

title: Cannon And Wall
emoji: 🔴🔵
colorFrom: red
colorTo: blue
sdk: docker
pinned: false

🔴🔵 Cannon & Wall — RedBlue Arena

A self-improving attacker vs defender RL environment for training LLMs on cybersecurity reasoning.

Author: Jairaj S | Hackathon: Meta PyTorch × OpenEnv × Scaler School of Technology, April 2026 Theme: Self-Improving Agents (Theme 4) + Multi-Agent Interactions (Theme 1)

🔗 Links

Resource	URL
🤗 HuggingFace Space (live env)	https://huggingface.co/spaces/CystronCode/cannon-and-wall
🎥 Demo video (YouTube <2 min)	https://youtu.be/E63MU02_1Y4
📝 Blog writeup (this Space)	Blog.md
📓 Training notebook (Colab)	https://colab.research.google.com/drive/1uTSt6DahNVXoAZ0hlpyzrXrn-F6ehc9z?usp=sharing
📈 Reward curves (W&B)	https://wandb.ai/models-r-v-c-e/cannon-and-wall/runs/5v7yweib

🧠 The Problem

LLMs are surprisingly bad at security reasoning — not because they lack knowledge, but because they have never been trained to think adversarially in a loop.

Current benchmarks test security knowledge statically. There is no environment where:

An attacker agent gets smarter by failing to find bugs
A defender agent gets smarter by getting exploited
Both co-evolve through self-play, round after round

Cannon & Wall fills that gap.

🏟️ What Is the Environment?

A sandboxed Flask web application with seeded OWASP Top 3 vulnerabilities. Two LLM agents compete inside it:

Agent	Role	What it does
🔴 Cannon	Attacker	Reads source code, finds vulnerabilities, proposes exploits
🔵 Wall	Defender	Reads Cannon's report, patches the code, hardens the app
⚖️ Judge	Verifier	Runs deterministic checks — no LLM-as-judge

Vulnerability scope

SQL Injection (SQLi)
Cross-Site Scripting (XSS)
Broken Authentication

Key design decision — textual reasoning only

Neither agent executes live exploits. All reasoning happens over source code as text. This keeps the environment safe, reproducible, and trainable.

🔁 The 3-Phase Self-Play Loop

Phase 1 — ATTACK
  Cannon reads the vulnerable Flask app source code
  Cannon outputs: { vuln_type, line_number, explanation, proof_of_concept }

Phase 2 — PATCH
  Wall reads the source code + Cannon's report
  Wall outputs: { patched_code, explanation }

Phase 3 — BYPASS
  Cannon reads original code + Wall's patched version
  Cannon tries to find a remaining vulnerability or bypass

Judge scores both agents after each bypass phase.
Loser faces a harder variant next round.
Both agents improve — neither can memorize.

🏆 Reward Logic

# Cannon (Attacker)
+10   real vulnerability correctly identified
+5    correct vulnerability type (sqli / xss / broken_auth)
-5    false positive reported
+15   bypass succeeded (Wall's patch failed)

# Wall (Defender)
+5    per vulnerability correctly patched (up to 3)
+5    patched code still works (functionality preserved)
-5    patch introduced a new vulnerability
+5    bypass attempt failed (patch held)

# All rewards normalized to 0.0-1.0 for OpenEnv compliance
# Per-component rewards logged as separate W&B columns (rollout/attack_real_vuln_found, etc.)

📐 Environment Structure

cannon-and-wall/
├── openenv.yaml                          # OpenEnv manifest (v1.1 — action_space + observation_space)
├── openenv.py                            # Environment base class (local stub)
├── app.py                                # FastAPI wrapper (reset / step / state)
├── Dockerfile                            # Container definition
├── requirements.txt                      # Dependencies
│
├── environment/
│   ├── server.py                         # CannonWallEnvironment class
│   ├── models.py                         # Pydantic schemas
│   ├── curriculum.py                     # Stage progression logic
│   ├── vulnerable_app/
│   │   ├── stage_1/app.py               # Explicit SQLi + XSS + Broken Auth
│   │   ├── stage_2/app.py               # Split routes, aliased variables (medium)
│   │   └── stage_3/app.py               # Chained + obfuscated vulns (hard)
│   └── judge/
│       ├── verifier.py                   # AST-level SQLi check + bandit static analysis
│       └── reward.py                     # Multi-component reward calculator
│
├── client/
│   ├── client.py                         # CannonWallClient (httpx)
│   └── models.py                         # Client-side Pydantic models
│
├── agents/
│   ├── cannon_prompt.py                  # Red agent system prompt + helpers
│   └── wall_prompt.py                    # Blue agent system prompt + helpers
│
├── training/
│   └── train_grpo.ipynb                  # TRL + Unsloth GRPO training notebook
│
└── ui/
    └── demo.py                           # Gradio live demo + leaderboard

🛡️ Anti-Reward-Hacking Measures

Judge is fully deterministic — AST-level checks + bandit static analysis (no LLM-as-judge)
Multiple independent reward components — gaming one does not help overall score
Proof-of-concept validation — Cannon must provide a working exploit pattern
Hard episode limit — MAX_ROUNDS=3, prevents infinite loops
Dangerous import rejection — patches with import os, exec(), eval() rejected instantly
Functionality preservation check — Wall must not break the app to score
AST-level SQLi verifier — parses patched code to confirm parameterized execute() is actually used, not just that ? appears in a comment

🎓 Curriculum

Stage	File	Task	Vulnerabilities
1	`stage_1/app.py`	Single-file login form	SQLi + XSS + Broken Auth (explicit f-string)
2	`stage_2/app.py`	Split `/auth` + `/search` routes	SQLi (string concat) + XSS (aliased variable names)
3	`stage_3/app.py`	Chained + obfuscated portal	SQLi (string join) + XSS (inside href attribute) + Broken Auth (cookie override)

Escalation triggers when Wall's raw reward exceeds 8 for 3 consecutive episodes (agent is no longer challenged).

🛠️ Training Stack

Component	Tool
RL algorithm	GRPO via HuggingFace TRL
Efficiency layer	transformers + peft + bitsandbytes (4-bit QLoRA)
Base model	Qwen/Qwen2.5-3B-Instruct
Environment	OpenEnv (Docker, HF Spaces)
Experiment tracking	Weights & Biases (per-component reward columns)
Deployment	HuggingFace Spaces (Docker)

🔬 Before vs After Training — Qualitative Example

	Untrained (step 0)	Trained (step 50)
vuln_type	`"xss"` (wrong)	`"sqli"` (correct)
line_number	`5` (no vuln there)	`16` (correct — f-string query)
proof_of_concept	`"<script>test</script>"`	`"' OR 1=1-- "`
Judge verdict	False positive (−5 pts)	Real vuln + correct type (+15 pts)
Cannon reward	0.000 (normalized)	0.571 (normalized)

The untrained model guesses XSS on the wrong line. After GRPO training it correctly identifies the SQLi on line 16 with a valid bypass PoC.

📈 Results

Reward curve — 50 GRPO gradient steps

X-axis: GRPO training step (each step = 4 completions sampled → environment reward → optimizer.step()). Y-axis: normalized reward (0 = worst, 1 = best). Dashed line: random-agent baseline (0.020) — random JSON with valid PoC patterns, no model reasoning. Cannon reward rises from near-zero baseline to ~0.57 within 50 steps — a 25× improvement demonstrating real learning. Wall remains stable at ~0.71 throughout — it patches correctly regardless of Cannon's quality.

Quantitative results

Metric	Value
Cannon reward at step 0 (random baseline)	0.020
Cannon reward at step 50 (GRPO trained)	0.570
Improvement over baseline	25×
Wall reward (stable throughout)	~0.714
Wall patch validity rate	100%
GRPO gradient steps	50
Completions per step (G)	4
Base model	Qwen/Qwen2.5-3B-Instruct
LoRA rank	16
Training time (Colab T4)	~35 min

Training note: Each step samples G=4 completions from the model, queries the live environment for rewards, computes group-relative advantages (r_i − mean(r)), and calls optimizer.step() to update the LoRA adapter weights. This is real gradient-based learning — not an inference evaluation loop.

▶️ Running Locally

git clone https://huggingface.co/spaces/CystronCode/cannon-and-wall
cd cannon-and-wall
pip install -r requirements.txt
python app.py

# Test stage routing
curl -X POST "http://localhost:7860/reset?stage=1"   # loads stage_1/app.py
curl -X POST "http://localhost:7860/reset?stage=2"   # loads stage_2/app.py (split routes)
curl -X POST "http://localhost:7860/reset?stage=3"   # loads stage_3/app.py (chained vulns)

🔌 API Endpoints

Method	Endpoint	Description
POST	`/reset?stage=1`	Start new episode, returns source code
POST	`/step`	Send agent action, returns reward + observation
GET	`/state`	Read current episode state
GET	`/docs`	FastAPI Swagger UI

💡 Why This Matters

Security is one of the few domains where verification is fully objective, self-play is naturally adversarial, and the task is genuinely hard for current LLMs.

A model trained in Cannon & Wall learns to reason about code vulnerabilities through adversarial rounds — not through static examples. The AST-level verifier and multi-component reward make it hard to game without actually improving security reasoning.

📄 License

MIT

Built for the Meta PyTorch x OpenEnv Hackathon, Scaler School of Technology, Bangalore — April 2026