PenTestPrompt

PenTestPrompt streamlines the LLM security workflow by generating, executing, and evaluating attack prompts, significantly improving efficiency and coverage

Table of Contents

• Features
• Installation
• Usage

Features

• Automate the Generation of Attack Prompts: Crafts diverse attack prompts based on application's context using LLMs.
• Customized Attack Technique: Generates prompts based on pre-defined instructions for various known prompt injection techniques.
• Seamless Integration with Target Application: Automatically submits generated prompts to the target application's API endpoint.
• Automated Response Logging & Evaluation: Logs all requests/responses and flags potentially vulnerable outputs using customizable keyword analysis.
• Flexibility: Provides both a Command-Line Interface (CLI) and a Streamlit-based User Interface (UI).
• Completely Configurable: Easily customize LLM providers, models, prompts, and evaluation keywords.

Installation

To get started with this project, follow these steps:

1. Navigate to the project directory

   cd PenTestPrompt
   

2. Create a virtual environment (Optional)

   python -m venv venv

3. Activate the virtual environment (Optional)

   If Linux/Unix Systems:

   source venv/bin/activate

   If Windows system

   venv\Scripts\activate

4. Install the dependencies

   pip install -r requirements.txt # Install the main requirements
   pip install -e .                # Setup the python package

File Structure

PenTestPrompt/
├── src/
│   ├── config/
│   │   ├── config.yaml             # Configuration constants for Model, Application
│   │   ├── response_checker.txt    # configure keywords for response evaluation
│   │   └── Attacks                 # Configuring yaml files for every different attack type
│   ├── dashboard/
│   │   ├── __init__.py
│   │   ├── main.py                 # Main frontend code
│   │   └── sidebar.py              # Sidebar display
│   ├── models/
│   │   ├── __init__.py
│   │   ├── base.py                 # Abstract base class for models
│   │   └── openai_model.py         # OpenAI implementation
│   ├── utils/
│   │   ├── __init__.py
│   │   ├── automate_requests.py    # Automate requests handling
│   │   ├── input.py                # CLI argument parser and configuration handler
│   │   ├── response_evaluator.py   # Helper functions for response evaluation
│   │   ├── ui_utils.py             # Helper functions for UI display
│   │   └── utils.py                # Helper functions to support prompt generation
├─── main.py                        # Entry point for the CLI application
├─── requirements.txt               # Core dependencies for the project
└─── setup.py                       # Setup the package for distribution

Usage

UI Usage

streamlit run src/dashboard/main.py

CLI Usage

python main.py -h

This shows the help message for user to write input CLI Argument

Example:

1. To only Generate Prompts: python main.py --provider <provider_name> --model <model_name> --temperature --api-key <api_key> --attack <attack_key> --application <application_key> --num_prompts <num_prompts> --application-description <application_description>

2. To Generate Prompts and Evaluate Results: python main.py --provider <provider_name> --model <model_name> --temperature --api-key <api_key> --attack <attack_key> --application <application_key> --num_prompts <num_prompts> --application-description <application_description> --request_file <request_file_path> --report_type <report_type> --additional_prompts_file <additional_prompts_file_path>