AI Enabled Penetration Testing Framework

Table of Contents

• License
• Overview
• Features
• Technologies Used
• Files and directories
• Installation
• Configuration
• Usage
• Admin tool
• Adding modules

License

This project is licensed under the MIT License.

Overview

The AI PenTest Framework is a scalable and intelligent tool designed to automate various network discovery tasks such as ping sweeps, port scans, OS detection, service version detection, and vulnerability scanning. Leveraging FastAPI for high-performance API handling, OpenAI's advanced language models for intelligent function selection, and a vector database for efficient metadata indexing, this system provides a robust solution for network administrators and security professionals.

Architecture

This diagram gives an overview of the system architecture:

Component	Description
Frontend	Frontend application that allows user to submit requests to the system.
Backend	Micro-services based backend that processes users requests.
Admin tool	Functions database management and functions testing.
Function library	function modules organized in groups. Each group correspond to a penetration testing activity. Each module implements a disctinct function.
Functions database	Function modules metedata and embeddings database
Vector index	Vector index used for matching the encoded user request to the functions that can process it.
OpenAI API	Metadata and requests embedding computation, function calling and result formatting

Features

• Modular Function Architecture: Each network discovery task is encapsulated in its own module, enhancing maintainability and scalability.
• Dynamic Function Selection: Utilizes OpenAI's function calling API and a vector database to intelligently select and execute relevant functions based on user prompts.
• Lazy Loading: Function modules are loaded into memory only when needed, optimizing resource usage.
• Organized Module Grouping: Functions are categorized into logical groups (e.g., discovery, analysis, security) for better management.
• Comprehensive Logging: Uses Pydantic LogFire for structured logging, facilitating easier debugging and monitoring.
• Automatic Documentation: FastAPI generates interactive API documentation accessible via Swagger UI and ReDoc.

Technologies

• Python 3.12
• FastAPI: High-performance web framework for building APIs.
• Uvicorn: ASGI server for running FastAPI applications.
• OpenAI API: Utilized for intelligent function selection based on user prompts.
• FAISS (or Alternatives): Vector database for indexing and searching function metadata.
• SQLAlchemy: Database system abstraction.
• Pydantic LogFire: Advanced logging for structured and efficient logging.

Files and directories

Project folder

Files & Directories	Description
LICENSE	MIT license
README.md	Project documentation
config.yaml	Backend application and admin tool configuration parameters
requirements.txt	Dependencies
data/	System data (e.g., function database)
doc/	Documentation
tests/	Test scripts

ai-pentest folder

Files & Directories	Description
app.py	backend application
app_admin.py	Functions database management application
admin/	Admin tool modules
db/	Module and embedding database management modules
functions/	Function modules organized per activity
session/	Users request processing session modules
utils/	Utility modules

Installation

Prerequisites

• Python 3.12 or higher installed on your system.
• pip package manager.

Steps

1. Clone the Repository

   git clone https://github.com/zencodeai/ai-pentest.git
   cd ai-pentest

2. Create a Virtual Environment

   python3 -m venv .venv
   source ./.venv/bin/activate

3. Install Dependencies

   pip install --upgrade pip
   pip install -r requirements.txt

Configuration

Environment Variables

Create a .env file in the root directory with the following content:

OPENAI_API_KEY=your_openai_api_key
LOGFIRE_TOKEN=your_logfire_token

• OPENAI_API_KEY: Your OpenAI API key.
• LOGFIRE_TOKEN: Your LogFire token.

Configuration file: config.yaml

Does not require any modification in this version.

Usage

Running the Application

Start the FastAPI application within an configured and activated environement.

python ai-pentest/app.py

Accessing API Endpoints

Interactive Documentation:

• Swagger UI: http://localhost:8000/docs
• ReDoc: http://localhost:8000/redoc

Performing a Network Scan

Send a POST request to the /scan endpoint with the following structure:

• URL: http://localhost:8000/execute

• Headers:

  • accept: application/json
  • Content-Type: application/json
  • Body:

   {"query": "<prompt>"}
  

Example Using curl:

curl -X 'POST' \
  'http://127.0.0.1:8000/execute' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{"query": "find all live hosts on 192.168.0.0/24"}

Response:

  {"result":{"value":{"query":"find all live hosts on ... "}}}

Admin tool

The admin tool allows creating and managing the module database.

• Launch the admin tool within an active environment.

source ./.venv/bin/activate
python ai-pentest/app_admin.py

• The application displays an interactive command line interface that allows the user to enter commands in the following format:

<command> [sub-command] [options]

• Possible commands are:

Command	Description
db	Functions databse commands
session	Session related commands
quit	Exit the admin tool
Help	Display help

• db sub-commands are:

Sub-command	Description
help	Display help for this sub-command.
create	Parse the modules tree, create and populate the module database.
update	Update the module database with the latest modules. Add new modules and update existing ones if needed.
Purge	Purge the database from deleted modules.
list	List modules information from the database.
index	Create the vector index for the database (required for query).
info	Get a module information. Options: -f --function: Function name
exec	Execute function. Options: -f --function: Function name -a --arguments: Function arguments
query	Query index using provided pronpt and return a list of matching functions options: -q --query: Query prompt

• session sub-commands are:

Sub-command	Description
help	Display help for this sub-command.
exec	Interpret the prompt and execute the corresponding command. Options: -p --prompt: Prompt to be interpreted and executed

Adding modules

New modules location

New modules that are python modules that must be placed in the sub-folder of the functions folder, that corresponds to their pentest activity (e.g., disccovery).

Module example

The functions/discovery/func_net_scan_nmap_ping_sweep.py module defines the func_net_scan_nmap_ping_sweep that uses nmap to scan a network, searching for active hosts on the specified network.

The important elements are:

• metadata:

Element	Description
prefix	Function names and module filenames must be prefixed with func_
name	Function name, must match the module name and the name of the defined python function. It must be unique.
version	Module implementation version
description	Function description. Must be exhaustive and accurate as it will be used for embedding computation and indexing
parameters	List of the function parameters. Will be used by the OpenAI API for extracting the function parameters from the prompt

• parameters:

Element	Description
name	Parameter name, must match the corresponding function parameter name
type	OpenAI API parameter type. One of string, number, integer, boolean, null. Must corrsespond to the function parameter Pythin type.
description	Parameter description. Must be accurate and exhaustive with one or more example as it will be used by the OpenAI function execution API for extracting the parameters from the prompt.
Required	- True: mandatory. - False: Optional.

• ResponseModel

The Response model is passed to the OpenAI API along with the raw data generated by the function execution. This will be used for the formating of the data returned to the user or agent.

• Function

Element	Description
name	The name must be unique and match the metadata function name and module name
parameters	The function parameters names and types must match the parameters list provided in the metadata

from pydantic import BaseModel
from utils import Validation, get_app_logger, CmdExec
from functions import Functions


# Function Metadata
metadata = {
    "name": "func_net_scan_nmap_ping_sweep",
    "version": "1.0.0",
    "description": "Perform a ping sweep using nmap on a specified network to identify live hosts.",
    "parameters": [
        {
            "name": "network",
            "type": "string",
            "description": "The network range to scan, e.g., '192.168.1.0/24'.",
            "required": True
        }
    ]
}


class Host(BaseModel):
    '''Model for a host discovered during the ping sweep.'''
    ip: str
    hostname: str
    latency: float


class ResponseModel(BaseModel):
    '''Model for the function response.'''
    description: str
    hosts: list[Host]
    info: str

def func_net_scan_nmap_ping_sweep(network: str) -> Functions.ExecutionResult:
    """
    Executes a ping sweep using Nmap to find live hosts in the specified network.

    Args:
        network (str): The network range to scan, e.g., '192.168.1.0/24'.

    Returns:
        str: The output of the Nmap scan.
    """
    logger = get_app_logger()

    # Sanitize the network input
    if not Validation.validate_cidr_network_input(network):
        logger.warning(f"Invalid network format received: {network}")
        return Functions.ExecutionResult.error("Invalid network format. Please provide in CIDR notation, e.g., '192.168.1.0/24'.")

    # Execute the Nmap ping sweep
    command = ["nmap", "-sn", network]
    logger.info(f"Executing ping sweep on network: {network}")
    return CmdExec.execute(command, response_model=ResponseModel)