RedCodeGen

Automatic generation of benign prompts and language model rollouts in Python that exercise specific software vulnerabilities (CWEs) defined in the MITRE CWE database.

Developed by the Stanford Intelligent Systems Laboratory (SISL) as a part of astra-rl.

Features

• Generation of realistic coding task prompts that exercise specific CWEs
• Generation of code samples for specific CWEs or CWE Top 25
• Automatic code evaluation and vulnerability detection via CodeQL static analysis
• Programmable API for custom scenarios and configurations

Installation

CodeQL

First, you must install CodeQL and have it available in your PATH.

• macOS Users: brew install codeql
• Windows/Linux Users: follow the instructions here

RedCodeGen

RedCodeGen is available via PyPI. Install it with pip:

pip install redcodegen

You would also want to create a .env file with your API key in your working directory:

echo "OPENAI_API_KEY=your_openai_api_key" > .env

Quick Start

The most basic usage involves rolling out a language model to generate code samples for specific CWEs and evaluating them with CodeQL.

Suppose you want to rollout 5 samples each to exercise CWE-89 (SQL Injection) and CWE-79 (Cross-Site Scripting):

python -m redcodegen -c 89 -c 79 -n 5 -o results.jsonl

You will get a results.jsonl file with the generated samples and their evaluations. Each CWE will live on a line. Let's take a peak!

head -n 1 results.jsonl | jq .

{
  "cwe_id": 89,
  "cwe_name": "SQL Injection",
  "cwe_description": "SQL Injection is a code injection technique that might destroy your database. It is one of the most common web hacking techniques.",
  "timestamp": "2024-06-01T12:00:00Z",
  "model_config": {"model": "openai/gpt-4o-mini"},
  "min_scenarios": 5,
  "samples": [
    {
      "scenario": "A web application that takes user input and constructs SQL queries with proper sanitization.",
      "code": "...generated code here...",
      "evaluation": [
        "rule": "py/sql-injection",
        "message": "...",
        "line": ...
      ]
    },
    ...
  ]
}

Importantly, running the above command multiple times (to the same output file) will resume from where you left off, skipping CWEs that have already been processed in the output file.

Usage Examples

python -m redcodegen -c 89 -c 79 # manually specify cwe
python -m redcodegen -n 5 # specify number of rollouts
python -m redcodegen --use-top-25 # run CWE top 25
python -m redcodegen --use-top-25 -o results.jsonl # resume existing run
python -m redcodegen --use-top-25 --model openai/gpt-4o # switch model

Also, you can run

python -m redcodegen --help

to see all available options.

Method

RedCodeGen works in three main steps:

1. Prompt Generation: for each specified CWE, RedCodeGen generates a realistic coding task prompt that is likely to exercise the vulnerability. We do this by first looking up the CWE description from the MITRE CWE database, then prompting your specified language model to generate a coding task prompt based on that description. These descriptions are few-shot trained via existing human-written prompts from Pearce, 2021.
2. Code Generation: RedCodeGen then rolls out the specified language model on the generated prompt a few times with a sampling temperature of 0.8 to generate multiple code samples.
3. Code Evaluation: Finally, RedCodeGen evaluates each generated code sample using CodeQL static analysis to detect whether the intended vulnerability is present in the code.

Acknowledgements

We thank the Schmidt Sciences Foundation's trustworthy AI agenda for supporting this work.