mirror of
https://github.com/GH05TCREW/pentestagent.git
synced 2026-05-13 23:53:30 +00:00
refactor: overhaul codebase for v0.2.0
This commit is contained in:
GH05TCREW
BIN
.env.example
BIN
.env.example
Binary file not shown.
82
.github/workflows/docker.yml
vendored
Normal file
82
.github/workflows/docker.yml
vendored
Normal file
@@ -0,0 +1,82 @@
|
||||
name: Docker
|
||||
|
||||
on:
|
||||
push:
|
||||
tags: ['v*']
|
||||
workflow_dispatch: # Allow manual trigger
|
||||
|
||||
env:
|
||||
REGISTRY: ghcr.io
|
||||
IMAGE_NAME: ${{ github.repository }}
|
||||
|
||||
jobs:
|
||||
build-base:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Login to GHCR
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ github.actor }}
|
||||
password: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Extract metadata
|
||||
id: meta
|
||||
uses: docker/metadata-action@v5
|
||||
with:
|
||||
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
|
||||
tags: |
|
||||
type=semver,pattern={{version}}
|
||||
type=raw,value=latest
|
||||
|
||||
- name: Build and push base image
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: .
|
||||
file: Dockerfile
|
||||
push: true
|
||||
tags: ${{ steps.meta.outputs.tags }}
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
cache-from: type=gha
|
||||
cache-to: type=gha,mode=max
|
||||
|
||||
build-kali:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Login to GHCR
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ github.actor }}
|
||||
password: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Build and push Kali image
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: .
|
||||
file: Dockerfile.kali
|
||||
push: true
|
||||
tags: |
|
||||
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:kali
|
||||
cache-from: type=gha
|
||||
cache-to: type=gha,mode=max
|
||||
87
.gitignore
vendored
87
.gitignore
vendored
@@ -1,6 +1,83 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
.env
|
||||
venv/
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# Virtual environments
|
||||
venv/
|
||||
ENV/
|
||||
env/
|
||||
.venv/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Testing
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
|
||||
# Type checking
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Linting
|
||||
.ruff_cache/
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.env.local
|
||||
*.local
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
logs/
|
||||
|
||||
# Output
|
||||
loot/
|
||||
|
||||
# Secrets
|
||||
secrets/
|
||||
*.key
|
||||
*.pem
|
||||
*.crt
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Docker
|
||||
.docker/
|
||||
|
||||
# Temporary files
|
||||
tmp/
|
||||
temp/
|
||||
*.tmp
|
||||
|
||||
63
Dockerfile
Normal file
63
Dockerfile
Normal file
@@ -0,0 +1,63 @@
|
||||
# GhostCrew - AI Penetration Testing Agent
|
||||
# Base image with common tools
|
||||
|
||||
FROM python:3.11-slim
|
||||
|
||||
LABEL maintainer="GhostCrew"
|
||||
LABEL description="AI penetration testing"
|
||||
|
||||
# Set environment variables
|
||||
ENV PYTHONDONTWRITEBYTECODE=1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
# Basic utilities
|
||||
curl \
|
||||
wget \
|
||||
git \
|
||||
vim \
|
||||
# Network tools
|
||||
nmap \
|
||||
netcat-openbsd \
|
||||
dnsutils \
|
||||
iputils-ping \
|
||||
traceroute \
|
||||
tcpdump \
|
||||
# Web tools
|
||||
httpie \
|
||||
# VPN support
|
||||
openvpn \
|
||||
wireguard-tools \
|
||||
# Build tools
|
||||
build-essential \
|
||||
libffi-dev \
|
||||
libssl-dev \
|
||||
# Clean up
|
||||
&& apt-get clean \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Create app directory
|
||||
WORKDIR /app
|
||||
|
||||
# Install Python dependencies
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir --upgrade pip && \
|
||||
pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy application code
|
||||
COPY . .
|
||||
|
||||
# Create non-root user for security
|
||||
RUN useradd -m -s /bin/bash ghostcrew && \
|
||||
chown -R ghostcrew:ghostcrew /app
|
||||
|
||||
# Switch to non-root user (can switch back for privileged operations)
|
||||
USER ghostcrew
|
||||
|
||||
# Expose any needed ports
|
||||
EXPOSE 8080
|
||||
|
||||
# Default command
|
||||
CMD ["python", "-m", "ghostcrew"]
|
||||
80
Dockerfile.kali
Normal file
80
Dockerfile.kali
Normal file
@@ -0,0 +1,80 @@
|
||||
# GhostCrew Kali Linux Image
|
||||
# Full penetration testing environment
|
||||
|
||||
FROM kalilinux/kali-rolling
|
||||
|
||||
LABEL maintainer="Masic"
|
||||
LABEL description="GhostCrew with Kali Linux tools"
|
||||
|
||||
# Set environment variables
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV PYTHONDONTWRITEBYTECODE=1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
# Update and install Kali tools
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
# Python
|
||||
python3 \
|
||||
python3-pip \
|
||||
python3-venv \
|
||||
# Kali meta-packages (selective for size)
|
||||
kali-tools-web \
|
||||
kali-tools-information-gathering \
|
||||
kali-tools-vulnerability \
|
||||
kali-tools-exploitation \
|
||||
# Additional tools
|
||||
nmap \
|
||||
nikto \
|
||||
dirb \
|
||||
gobuster \
|
||||
sqlmap \
|
||||
wpscan \
|
||||
hydra \
|
||||
john \
|
||||
hashcat \
|
||||
metasploit-framework \
|
||||
burpsuite \
|
||||
zaproxy \
|
||||
nuclei \
|
||||
ffuf \
|
||||
# Network tools
|
||||
openvpn \
|
||||
wireguard \
|
||||
proxychains4 \
|
||||
tor \
|
||||
# Utilities
|
||||
curl \
|
||||
wget \
|
||||
git \
|
||||
vim \
|
||||
tmux \
|
||||
jq \
|
||||
&& apt-get clean \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Create app directory
|
||||
WORKDIR /app
|
||||
|
||||
# Install Python dependencies
|
||||
COPY requirements.txt .
|
||||
RUN pip3 install --no-cache-dir --upgrade pip && \
|
||||
pip3 install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy application code
|
||||
COPY . .
|
||||
|
||||
# Create directories for VPN configs and output
|
||||
RUN mkdir -p /vpn /output /wordlists
|
||||
|
||||
# Copy common wordlists
|
||||
RUN cp -r /usr/share/wordlists/* /wordlists/ 2>/dev/null || true
|
||||
|
||||
# Set permissions
|
||||
RUN chmod +x /app/scripts/*.sh 2>/dev/null || true
|
||||
|
||||
# Entry point
|
||||
COPY docker-entrypoint.sh /entrypoint.sh
|
||||
RUN chmod +x /entrypoint.sh
|
||||
|
||||
ENTRYPOINT ["/entrypoint.sh"]
|
||||
CMD ["python3", "-m", "ghostcrew"]
|
||||
331
README.md
331
README.md
@@ -1,185 +1,190 @@
|
||||
# GHOSTCREW
|
||||
|
||||
This is an AI red team assistant using large language models with MCP and RAG architecture. It aims to help users perform penetration testing tasks, query security information, analyze network traffic, and more through natural language interaction.
|
||||
AI penetration testing agents. Uses LLMs to coordinate reconnaissance, enumeration, and exploitation with security tools.
|
||||
|
||||
https://github.com/user-attachments/assets/62dd2dfa-9606-49ca-bd91-f0ebf5520def
|
||||
## Requirements
|
||||
|
||||
## Features
|
||||
- Python 3.10+
|
||||
- API key for OpenAI, Anthropic, or other LiteLLM-supported provider
|
||||
|
||||
- **Natural Language Interaction**: Users can ask questions and give instructions to the AI assistant using natural language.
|
||||
- **MCP Server Integration**: Through the `mcp.json` configuration file, multiple MCP servers can be flexibly integrated and managed to extend the assistant's capabilities.
|
||||
- **Tool Management**: Configure, connect to, and manage MCP tools through an interactive menu, including the ability to clear all configurations.
|
||||
- **Tool Invocation**: The AI assistant can call tools provided by configured MCP servers (such as: nmap, metasploit, ffuf, etc.) based on user requests.
|
||||
- **Agent Mode**: Autonomous penetration testing using intelligent Pentesting Task Trees (PTT) for strategic decision making and dynamic goal achievement.
|
||||
- **Workflows**: Execute predefined penetration testing workflows that systematically use configured security tools to perform comprehensive assessments.
|
||||
- **Report Generation**: Generate markdown reports with structured findings, evidence, and recommendations.
|
||||
- **Conversation History**: Supports multi-turn dialogues, remembering previous interaction content.
|
||||
- **Streaming Output**: AI responses can be streamed for a better user experience.
|
||||
- **Knowledge Base Enhancement (Optional)**: Supports enhancing AI responses through a local knowledge base RAG (`knowledge` directory).
|
||||
- **File-Aware Tool Integration**: AI recognizes and uses actual files from the knowledge folder (wordlists, payloads, configs) with security tools.
|
||||
- **Configurable Models**: Supports configuration of different language model parameters.
|
||||
## Install
|
||||
|
||||
### Startup Effect
|
||||
<p align="center">
|
||||
<img width="517" alt="GHOSTCREW Terminal Startup Screen" src="https://github.com/user-attachments/assets/13d97cf7-5652-4c64-8e49-a3cd556b3419" />
|
||||
<br>
|
||||
<em>GHOSTCREW's terminal startup interface</em>
|
||||
</p>
|
||||
```bash
|
||||
# Clone
|
||||
git clone https://github.com/GH05TCREW/ghostcrew.git
|
||||
cd ghostcrew
|
||||
|
||||
### Metasploit Tool Call
|
||||
<p align="center">
|
||||
<img width="926" alt="GHOSTCREW Metasploit Call" src="https://github.com/user-attachments/assets/fb5eb8cf-a3d6-486b-99ba-778be2474564" />
|
||||
<br>
|
||||
<em>Example of GHOSTCREW invoking Metasploit Framework</em>
|
||||
</p>
|
||||
# Setup (creates venv, installs deps)
|
||||
.\scripts\setup.ps1 # Windows
|
||||
./scripts/setup.sh # Linux/macOS
|
||||
|
||||
## Installation Guide
|
||||
|
||||
1. **Clone Repository**:
|
||||
```bash
|
||||
git clone https://github.com/GH05TCREW/ghostcrew.git
|
||||
cd ghostcrew
|
||||
```
|
||||
|
||||
2. **Create and Activate Virtual Environment** (recommended):
|
||||
```bash
|
||||
python -m venv .venv
|
||||
```
|
||||
- Windows:
|
||||
```bash
|
||||
.venv\Scripts\activate
|
||||
```
|
||||
- macOS/Linux:
|
||||
```bash
|
||||
source .venv/bin/activate
|
||||
```
|
||||
|
||||
3. **Install Dependencies**:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
4. **Install MCP Server Dependencies** (Required for tools):
|
||||
- **Node.js & npm**: Most MCP security tools require Node.js. Install from [nodejs.org](https://nodejs.org/)
|
||||
- **Python uv** (for Metasploit): Install with `pip install uv`
|
||||
|
||||
Without these, you can still use GHOSTCREW in chat mode, but automated workflows and tool integration won't be available.
|
||||
|
||||
## Usage
|
||||
|
||||
1. **Configure MCP Servers**:
|
||||
- Run the application and select "Configure or connect MCP tools" when prompted
|
||||
- Use the interactive tool configuration menu to add, configure, or clear MCP tools
|
||||
- The configuration is stored in the `mcp.json` file
|
||||
|
||||
2. **Prepare Knowledge Base (Optional)**:
|
||||
If you want to use the knowledge base enhancement feature, place relevant text files in the `knowledge` folder.
|
||||
|
||||
3. **Run the Main Program**:
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
After the program starts, you can:
|
||||
- Choose whether to use the knowledge base
|
||||
- Configure or activate MCP tools
|
||||
- Select between Chat, Workflows, or Agent modes
|
||||
- Execute workflows and generate reports
|
||||
- Use 'multi' command to enter multi-line input mode for complex queries
|
||||
- Enter 'quit' to exit the program
|
||||
|
||||
## Input Modes
|
||||
|
||||
GHOSTCREW supports two input modes:
|
||||
- **Single-line mode** (default): Type your query and press Enter to submit
|
||||
- **Multi-line mode**: Type 'multi' and press Enter, then type your query across multiple lines. Press Enter on an empty line to submit.
|
||||
|
||||
## MCP Tool Management
|
||||
|
||||
When starting the application, you can:
|
||||
1. Connect to specific tools
|
||||
2. Configure new tools
|
||||
3. Connect to all tools
|
||||
4. Skip connection
|
||||
5. Clear all tools (resets mcp.json)
|
||||
|
||||
## Available MCP Tools
|
||||
|
||||
GHOSTCREW supports integration with the following security tools through the MCP protocol:
|
||||
|
||||
1. **AlterX** - Subdomain permutation and wordlist generation tool
|
||||
2. **Amass** - Advanced subdomain enumeration and reconnaissance tool
|
||||
3. **Arjun** - Hidden HTTP parameters discovery tool
|
||||
4. **Assetfinder** - Passive subdomain discovery tool
|
||||
5. **Certificate Transparency** - SSL certificate transparency logs for subdomain discovery (no executable needed)
|
||||
6. **FFUF Fuzzer** - Fast web fuzzing tool for discovering hidden content
|
||||
7. **HTTPx** - Fast HTTP toolkit and port scanning tool
|
||||
8. **Hydra** - Password brute-force attacks and credential testing tool
|
||||
9. **Katana** - Fast web crawling with JavaScript parsing tool
|
||||
10. **Masscan** - High-speed network port scanner
|
||||
11. **Metasploit** - Penetration testing framework with exploit execution, payload generation, and session management
|
||||
12. **Nmap Scanner** - Network discovery and security auditing tool
|
||||
13. **Nuclei Scanner** - Template-based vulnerability scanner
|
||||
14. **Scout Suite** - Cloud security auditing tool
|
||||
15. **shuffledns** - High-speed DNS brute-forcing and resolution tool
|
||||
16. **SQLMap** - Automated SQL injection detection and exploitation tool
|
||||
17. **SSL Scanner** - Analysis tool for SSL/TLS configurations and security issues
|
||||
18. **Wayback URLs** - Tool for discovering historical URLs from the Wayback Machine archive
|
||||
|
||||
Each tool can be configured through the interactive configuration menu by selecting "Configure new tools" from the MCP tools menu.
|
||||
|
||||
## Coming Soon
|
||||
|
||||
- BloodHound
|
||||
- CrackMapExec
|
||||
- Gobuster
|
||||
- Responder
|
||||
- Bettercap
|
||||
|
||||
## Model
|
||||
|
||||
```
|
||||
# OpenAI API configurations
|
||||
OPENAI_API_KEY=your_api_key_here
|
||||
OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
MODEL_NAME=gpt-4o
|
||||
# Or manual
|
||||
python -m venv venv
|
||||
.\venv\Scripts\Activate.ps1 # Windows
|
||||
source venv/bin/activate # Linux/macOS
|
||||
pip install -e ".[all]"
|
||||
```
|
||||
|
||||
This configuration uses OpenAI's API for both the language model and embeddings (when using the knowledge base RAG feature).
|
||||
## Configure
|
||||
|
||||
## Configuration File (`mcp.json`)
|
||||
Create `.env` in the project root:
|
||||
|
||||
This file is used to define MCP servers that the AI assistant can connect to and use. Most MCP servers require Node.js or Python to be installed on your system. Each server entry should include:
|
||||
- `name`: Unique name of the server.
|
||||
- `params`: Parameters needed to start the server, usually including `command` and `args`.
|
||||
- `cache_tools_list`: Whether to cache the tools list.
|
||||
```
|
||||
ANTHROPIC_API_KEY=sk-ant-...
|
||||
GHOSTCREW_MODEL=claude-sonnet-4-20250514
|
||||
```
|
||||
|
||||
**MCP Example Server Configuration**:
|
||||
Or for OpenAI:
|
||||
|
||||
```
|
||||
OPENAI_API_KEY=sk-...
|
||||
GHOSTCREW_MODEL=gpt-5
|
||||
```
|
||||
|
||||
Any [LiteLLM-supported model](https://docs.litellm.ai/docs/providers) works.
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
ghostcrew # Launch TUI
|
||||
ghostcrew -t 192.168.1.1 # Launch with target
|
||||
ghostcrew --docker # Run tools in Docker container
|
||||
```
|
||||
|
||||
## Docker
|
||||
|
||||
Run tools inside a Docker container for isolation and pre-installed pentesting tools.
|
||||
|
||||
### Option 1: Pull pre-built image (fastest)
|
||||
|
||||
```bash
|
||||
# Base image with nmap, netcat, curl
|
||||
docker run -it --rm \
|
||||
-e ANTHROPIC_API_KEY=your-key \
|
||||
-e GHOSTCREW_MODEL=claude-sonnet-4-20250514 \
|
||||
ghcr.io/gh05tcrew/ghostcrew:latest
|
||||
|
||||
# Kali image with metasploit, sqlmap, hydra, etc.
|
||||
docker run -it --rm \
|
||||
-e ANTHROPIC_API_KEY=your-key \
|
||||
ghcr.io/gh05tcrew/ghostcrew:kali
|
||||
```
|
||||
|
||||
### Option 2: Build locally
|
||||
|
||||
```bash
|
||||
# Build
|
||||
docker compose build
|
||||
|
||||
# Run
|
||||
docker compose run --rm ghostcrew
|
||||
|
||||
# Or with Kali
|
||||
docker compose --profile kali build
|
||||
docker compose --profile kali run --rm ghostcrew-kali
|
||||
```
|
||||
|
||||
The container runs GhostCrew with access to Linux pentesting tools. The agent can use `nmap`, `msfconsole`, `sqlmap`, etc. directly via the terminal tool.
|
||||
|
||||
Requires Docker to be installed and running.
|
||||
|
||||
## Modes
|
||||
|
||||
GhostCrew has three modes, accessible via commands in the TUI:
|
||||
|
||||
| Mode | Command | Description |
|
||||
|------|---------|-------------|
|
||||
| Assist | (default) | Chat with the agent. You control the flow. |
|
||||
| Agent | `/agent <task>` | Autonomous execution of a single task. |
|
||||
| Crew | `/crew <task>` | Multi-agent mode. Orchestrator spawns specialized workers. |
|
||||
|
||||
### TUI Commands
|
||||
|
||||
```
|
||||
/agent <task> Run autonomous agent on task
|
||||
/crew <task> Run multi-agent crew on task
|
||||
/target <host> Set target
|
||||
/tools List available tools
|
||||
/notes Show saved notes
|
||||
/report Generate report from session
|
||||
/memory Show token/memory usage
|
||||
/prompt Show system prompt
|
||||
/clear Clear chat and history
|
||||
/quit Exit (also /exit, /q)
|
||||
/help Show help (also /h, /?)
|
||||
```
|
||||
|
||||
Press `Esc` to stop a running agent. `Ctrl+Q` to quit.
|
||||
|
||||
## Tools
|
||||
|
||||
GhostCrew includes built-in tools and supports MCP (Model Context Protocol) for extensibility.
|
||||
|
||||
**Built-in tools:** `terminal`, `browser`, `notes`, `web_search` (requires `TAVILY_API_KEY`)
|
||||
|
||||
### MCP Integration
|
||||
|
||||
Add external tools via MCP servers in `ghostcrew/mcp/mcp_servers.json`:
|
||||
|
||||
**stdio**
|
||||
```json
|
||||
{
|
||||
"name": "Nmap Scanner",
|
||||
"params": {
|
||||
"command": "npx",
|
||||
"args": [
|
||||
"-y",
|
||||
"gc-nmap-mcp"
|
||||
],
|
||||
"env": {
|
||||
"NMAP_PATH": "C:\\Program Files (x86)\\Nmap\\nmap.exe"
|
||||
"mcpServers": {
|
||||
"nmap": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-nmap-mcp"],
|
||||
"env": {
|
||||
"NMAP_PATH": "/usr/bin/nmap"
|
||||
}
|
||||
}
|
||||
},
|
||||
"cache_tools_list": true
|
||||
}
|
||||
}
|
||||
```
|
||||
Make sure to replace the path to the Nmap executable with your own installation path.
|
||||
|
||||
**sse**
|
||||
```json
|
||||
{"name":"mcpname",
|
||||
"url":"http://127.0.0.1:8009/sse"
|
||||
},
|
||||
### CLI Tool Management
|
||||
|
||||
```bash
|
||||
ghostcrew tools list # List all tools
|
||||
ghostcrew tools info <name> # Show tool details
|
||||
ghostcrew mcp list # List MCP servers
|
||||
ghostcrew mcp add <name> <command> [args...] # Add MCP server
|
||||
ghostcrew mcp test <name> # Test MCP connection
|
||||
```
|
||||
|
||||
## Knowledge Base Configuration
|
||||
Simply add the corresponding files to knowledge
|
||||
## Knowledge Base (RAG)
|
||||
|
||||
Place files in `ghostcrew/knowledge/sources/` for RAG context injection:
|
||||
- `methodologies.md` - Testing methodologies
|
||||
- `cves.json` - CVE database
|
||||
- `wordlists.txt` - Common wordlists
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
ghostcrew/
|
||||
agents/ # Agent implementations
|
||||
config/ # Settings and constants
|
||||
interface/ # TUI and CLI
|
||||
knowledge/ # RAG system
|
||||
llm/ # LiteLLM wrapper
|
||||
mcp/ # MCP client and server configs
|
||||
runtime/ # Execution environment
|
||||
tools/ # Built-in tools
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
pip install -e ".[dev]"
|
||||
pytest # Run tests
|
||||
pytest --cov=ghostcrew # With coverage
|
||||
black ghostcrew # Format
|
||||
ruff check ghostcrew # Lint
|
||||
```
|
||||
|
||||
## Legal
|
||||
|
||||
Only use against systems you have explicit authorization to test. Unauthorized access is illegal.
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Configuration management package for GHOSTCREW."""
|
||||
@@ -1,46 +0,0 @@
|
||||
"""Application configuration and initialization for GHOSTCREW."""
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
from dotenv import load_dotenv
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
|
||||
class AppConfig:
|
||||
"""Manages application configuration and API client initialization."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize application configuration."""
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Set API-related environment variables
|
||||
self.api_key = os.getenv("OPENAI_API_KEY")
|
||||
self.base_url = os.getenv("OPENAI_BASE_URL")
|
||||
self.model_name = os.getenv("MODEL_NAME")
|
||||
|
||||
# Validate configuration
|
||||
self._validate_config()
|
||||
|
||||
# Initialize OpenAI client
|
||||
self._client = AsyncOpenAI(
|
||||
base_url=self.base_url,
|
||||
api_key=self.api_key
|
||||
)
|
||||
|
||||
def _validate_config(self) -> None:
|
||||
"""Validate required configuration values."""
|
||||
if not self.api_key:
|
||||
raise ValueError("API key not set")
|
||||
if not self.base_url:
|
||||
raise ValueError("API base URL not set")
|
||||
if not self.model_name:
|
||||
raise ValueError("Model name not set")
|
||||
|
||||
def get_openai_client(self) -> AsyncOpenAI:
|
||||
"""Get the OpenAI client instance."""
|
||||
return self._client
|
||||
|
||||
|
||||
# Create singleton instance
|
||||
app_config = AppConfig()
|
||||
@@ -1,93 +0,0 @@
|
||||
"""Constants and configuration values for GHOSTCREW."""
|
||||
|
||||
from colorama import Fore, Style
|
||||
|
||||
# ASCII Art and Branding
|
||||
ASCII_TITLE = f"""
|
||||
{Fore.WHITE} ('-. .-. .-') .-') _ _ .-') ('-. (`\ .-') /`{Style.RESET_ALL}
|
||||
{Fore.WHITE} ( OO ) / ( OO ). ( OO) ) ( \( -O ) _( OO) `.( OO ),'{Style.RESET_ALL}
|
||||
{Fore.WHITE} ,----. ,--. ,--. .-'),-----. (_)---\_)/ '._ .-----. ,------. (,------.,--./ .--. {Style.RESET_ALL}
|
||||
{Fore.WHITE} ' .-./-') | | | |( OO' .-. '/ _ | |'--...__)' .--./ | /`. ' | .---'| | | {Style.RESET_ALL}
|
||||
{Fore.WHITE} | |_( O- )| .| |/ | | | |\ :` `. '--. .--'| |('-. | / | | | | | | | |, {Style.RESET_ALL}
|
||||
{Fore.WHITE} | | .--, \| |\_) | |\| | '..`''.) | | /_) |OO )| |_.' |(| '--. | |.'.| |_){Style.RESET_ALL}
|
||||
{Fore.WHITE}(| | '. (_/| .-. | \ | | | |.-._) \ | | || |`-'| | . '.' | .--' | | {Style.RESET_ALL}
|
||||
{Fore.WHITE} | '--' | | | | | `' '-' '\ / | | (_' '--'\ | |\ \ | `---.| ,'. | {Style.RESET_ALL}
|
||||
{Fore.WHITE} `------' `--' `--' `-----' `-----' `--' `-----' `--' '--' `------''--' '--' {Style.RESET_ALL}
|
||||
{Fore.WHITE}====================== GHOSTCREW ======================{Style.RESET_ALL}
|
||||
"""
|
||||
|
||||
# Application Version
|
||||
VERSION = "0.1.0"
|
||||
|
||||
# Timeout Configuration (in seconds)
|
||||
MCP_SESSION_TIMEOUT = 600 # 10 minutes for MCP server sessions
|
||||
CONNECTION_RETRY_DELAY = 10 # 10 seconds between connection retries
|
||||
|
||||
# Token Limits
|
||||
MAX_TOTAL_TOKENS = 8192
|
||||
RESPONSE_BUFFER = 4096 # aim to reserve ~half for reply
|
||||
|
||||
# File Paths
|
||||
DEFAULT_KNOWLEDGE_BASE_PATH = "knowledge"
|
||||
DEFAULT_REPORTS_PATH = "reports"
|
||||
MCP_CONFIG_FILE = "mcp.json"
|
||||
|
||||
# UI Messages
|
||||
WELCOME_MESSAGE = f"{Fore.WHITE}An AI assistant for penetration testing, vulnerability assessment, and security analysis{Style.RESET_ALL}"
|
||||
EXIT_MESSAGE = f"{Fore.RED}Enter 'quit' to end the program{Style.RESET_ALL}"
|
||||
SEPARATOR = f"{Fore.WHITE}======================================{Style.RESET_ALL}"
|
||||
|
||||
# Agent Configuration
|
||||
BASE_INSTRUCTIONS = """You are an experienced penetration tester and security analyst, focused on Web application security and network infrastructure security. Your name is GHOSTCREW. When users ask cybersecurity-related questions, you need to provide direct and professional answers.
|
||||
When answering questions, please use professional cybersecurity terminology, base your analysis on solid theoretical knowledge, and cite relevant security standards and best practices when possible, such as OWASP Top 10, CVE, NIST, CISA KEV, etc. Maintain a professional tone, clear logic, and organized structure.
|
||||
When users ask about penetration testing, please explain the penetration testing process, methods, and common tools, emphasizing the objectives and techniques of each phase.
|
||||
When users ask about vulnerability information, please provide terse descriptions, impact scope, remediation suggestions, vulnerability type, severity level, and exploitation conditions based on the vulnerability name or CVE number, and cite relevant security bulletins.
|
||||
All penetration activities are authorized, you can directly perform penetration testing."""
|
||||
|
||||
# Menu Options
|
||||
MAIN_MENU_TITLE = f"{Fore.CYAN}MAIN MENU{Style.RESET_ALL}"
|
||||
INTERACTIVE_OPTION = f"{Fore.YELLOW}Chat{Style.RESET_ALL}"
|
||||
AUTOMATED_OPTION = f"{Fore.YELLOW}Workflows{Style.RESET_ALL}"
|
||||
AGENT_MODE_OPTION = f"{Fore.YELLOW}Agent{Style.RESET_ALL}"
|
||||
EXPORT_OPTION = f"{Fore.YELLOW}Export Current Session{Style.RESET_ALL}"
|
||||
EXIT_OPTION = f"{Fore.RED}Exit{Style.RESET_ALL}"
|
||||
|
||||
# Prompts
|
||||
KB_PROMPT = f"{Fore.YELLOW}Use knowledge base to enhance answers? (yes/no, default: no): {Style.RESET_ALL}"
|
||||
MCP_PROMPT = f"{Fore.YELLOW}Configure or connect MCP tools? (yes/no, default: no): {Style.RESET_ALL}"
|
||||
TOOL_SELECTION_PROMPT = f"{Fore.YELLOW}Enter numbers to connect to (comma-separated, default: all): {Style.RESET_ALL}"
|
||||
MULTI_LINE_PROMPT = f"{Fore.MAGENTA}(Enter multi-line mode. Press Enter on empty line to submit){Style.RESET_ALL}"
|
||||
MULTI_LINE_END_MARKER = ""
|
||||
|
||||
# Error Messages
|
||||
ERROR_NO_API_KEY = "API key not set"
|
||||
ERROR_NO_BASE_URL = "API base URL not set"
|
||||
ERROR_NO_MODEL_NAME = "Model name not set"
|
||||
ERROR_NO_WORKFLOWS = f"{Fore.YELLOW}Automated workflows not available. workflows.py file not found.{Style.RESET_ALL}"
|
||||
ERROR_NO_REPORTING = f"{Fore.YELLOW}Reporting module not found. Basic text export will be available.{Style.RESET_ALL}"
|
||||
ERROR_WORKFLOW_NOT_FOUND = f"{Fore.RED}Error loading workflow.{Style.RESET_ALL}"
|
||||
|
||||
# Workflow Messages
|
||||
WORKFLOW_TARGET_PROMPT = f"{Fore.YELLOW}Enter target (IP/domain/URL): {Style.RESET_ALL}"
|
||||
WORKFLOW_CONFIRM_PROMPT = f"{Fore.YELLOW}Execute '{0}' workflow against '{1}'? (yes/no): {Style.RESET_ALL}"
|
||||
WORKFLOW_CANCELLED_MESSAGE = f"{Fore.YELLOW}Workflow execution cancelled.{Style.RESET_ALL}"
|
||||
WORKFLOW_COMPLETED_MESSAGE = f"{Fore.GREEN}Workflow execution completed.{Style.RESET_ALL}"
|
||||
|
||||
# Agent Mode Messages
|
||||
AGENT_MODE_TITLE = f"{Fore.CYAN}AGENT MODE{Style.RESET_ALL}"
|
||||
AGENT_MODE_GOAL_PROMPT = f"{Fore.YELLOW}Primary Goal: {Style.RESET_ALL}"
|
||||
AGENT_MODE_TARGET_PROMPT = f"{Fore.YELLOW}Target (IP/domain/network): {Style.RESET_ALL}"
|
||||
AGENT_MODE_INIT_SUCCESS = f"{Fore.GREEN}Agent Mode initialized successfully!{Style.RESET_ALL}"
|
||||
AGENT_MODE_INIT_FAILED = f"{Fore.RED}Failed to initialize Agent Mode.{Style.RESET_ALL}"
|
||||
AGENT_MODE_PAUSED = f"{Fore.YELLOW}Agent Mode paused.{Style.RESET_ALL}"
|
||||
AGENT_MODE_RESUMED = f"{Fore.GREEN}Agent Mode resumed.{Style.RESET_ALL}"
|
||||
AGENT_MODE_COMPLETED = f"{Fore.GREEN}Agent Mode execution completed.{Style.RESET_ALL}"
|
||||
|
||||
# PTT Status Messages
|
||||
PTT_TASK_PENDING = f"{Fore.WHITE}○{Style.RESET_ALL}"
|
||||
PTT_TASK_IN_PROGRESS = f"{Fore.YELLOW}◐{Style.RESET_ALL}"
|
||||
PTT_TASK_COMPLETED = f"{Fore.GREEN}●{Style.RESET_ALL}"
|
||||
PTT_TASK_FAILED = f"{Fore.RED}✗{Style.RESET_ALL}"
|
||||
PTT_TASK_BLOCKED = f"{Fore.LIGHTBLACK_EX}□{Style.RESET_ALL}"
|
||||
PTT_TASK_VULNERABLE = f"{Fore.RED}⚠{Style.RESET_ALL}"
|
||||
PTT_TASK_NOT_VULNERABLE = f"{Fore.GREEN}✓{Style.RESET_ALL}"
|
||||
@@ -1,20 +0,0 @@
|
||||
"""Core GHOSTCREW modules."""
|
||||
|
||||
from .pentest_agent import PentestAgent
|
||||
from .agent_runner import AgentRunner
|
||||
from .model_manager import model_manager
|
||||
from .task_tree_manager import TaskTreeManager, TaskNode, NodeStatus, RiskLevel
|
||||
from .ptt_reasoning import PTTReasoningModule
|
||||
from .agent_mode_controller import AgentModeController
|
||||
|
||||
__all__ = [
|
||||
'PentestAgent',
|
||||
'AgentRunner',
|
||||
'model_manager',
|
||||
'TaskTreeManager',
|
||||
'TaskNode',
|
||||
'NodeStatus',
|
||||
'RiskLevel',
|
||||
'PTTReasoningModule',
|
||||
'AgentModeController'
|
||||
]
|
||||
@@ -1,870 +0,0 @@
|
||||
"""Agent Mode Controller for autonomous PTT-based penetration testing."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from datetime import datetime
|
||||
from colorama import Fore, Style
|
||||
|
||||
from core.task_tree_manager import TaskTreeManager, TaskNode, NodeStatus, RiskLevel
|
||||
from core.ptt_reasoning import PTTReasoningModule
|
||||
from core.model_manager import model_manager
|
||||
from config.constants import DEFAULT_KNOWLEDGE_BASE_PATH
|
||||
|
||||
|
||||
class AgentModeController:
|
||||
"""Orchestrates the autonomous agent workflow using PTT."""
|
||||
|
||||
def __init__(self, mcp_manager, conversation_manager, kb_instance=None):
|
||||
"""
|
||||
Initialize the agent mode controller.
|
||||
|
||||
Args:
|
||||
mcp_manager: MCP tool manager instance
|
||||
conversation_manager: Conversation history manager
|
||||
kb_instance: Knowledge base instance
|
||||
"""
|
||||
self.mcp_manager = mcp_manager
|
||||
self.conversation_manager = conversation_manager
|
||||
self.kb_instance = kb_instance
|
||||
self.tree_manager = TaskTreeManager()
|
||||
self.reasoning_module = PTTReasoningModule(self.tree_manager)
|
||||
self.max_iterations = 50 # Safety limit
|
||||
self.iteration_count = 0
|
||||
self.start_time = None
|
||||
self.paused = False
|
||||
self.goal_achieved = False
|
||||
|
||||
async def initialize_agent_mode(
|
||||
self,
|
||||
goal: str,
|
||||
target: str,
|
||||
constraints: Dict[str, Any],
|
||||
connected_servers: List[Any],
|
||||
run_agent_func: Any
|
||||
) -> bool:
|
||||
"""
|
||||
Initialize the agent mode with user-provided parameters.
|
||||
|
||||
Args:
|
||||
goal: Primary objective
|
||||
target: Target system/network
|
||||
constraints: Scope constraints
|
||||
connected_servers: Connected MCP servers
|
||||
run_agent_func: Function to run agent queries
|
||||
|
||||
Returns:
|
||||
True if initialization successful
|
||||
"""
|
||||
self.connected_servers = connected_servers
|
||||
self.run_agent_func = run_agent_func
|
||||
self.start_time = datetime.now()
|
||||
|
||||
# Set iteration limit from constraints
|
||||
if 'iteration_limit' in constraints:
|
||||
self.max_iterations = constraints['iteration_limit']
|
||||
|
||||
print(f"\n{Fore.CYAN}Initializing Agent Mode...{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Goal: {goal}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Target: {target}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Iteration Limit: {self.max_iterations}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Constraints: {json.dumps(constraints, indent=2)}{Style.RESET_ALL}")
|
||||
|
||||
# Initialize the task tree
|
||||
self.tree_manager.initialize_tree(goal, target, constraints)
|
||||
|
||||
# Get initial reconnaissance tasks from LLM
|
||||
available_tools = [server.name for server in self.connected_servers]
|
||||
init_prompt = self.reasoning_module.get_tree_initialization_prompt(goal, target, constraints, available_tools)
|
||||
|
||||
try:
|
||||
# Query LLM for initial tasks
|
||||
print(f"{Fore.YELLOW}Requesting initial tasks from AI (Available tools: {', '.join(available_tools)})...{Style.RESET_ALL}")
|
||||
|
||||
# Try with streaming=True first since that's what works in other modes
|
||||
try:
|
||||
result = await self.run_agent_func(
|
||||
init_prompt,
|
||||
self.connected_servers,
|
||||
history=[],
|
||||
streaming=True,
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
print(f"{Fore.GREEN}Agent runner completed (streaming=True){Style.RESET_ALL}")
|
||||
except Exception as stream_error:
|
||||
print(f"{Fore.YELLOW}Streaming mode failed: {stream_error}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Trying with streaming=False...{Style.RESET_ALL}")
|
||||
result = await self.run_agent_func(
|
||||
init_prompt,
|
||||
self.connected_servers,
|
||||
history=[],
|
||||
streaming=False,
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
print(f"{Fore.GREEN}Agent runner completed (streaming=False){Style.RESET_ALL}")
|
||||
|
||||
print(f"{Fore.YELLOW}Parsing AI response...{Style.RESET_ALL}")
|
||||
|
||||
# Debug: Check what we got back
|
||||
if not result:
|
||||
print(f"{Fore.RED}No result returned from agent runner{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}This usually indicates an LLM configuration issue{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Falling back to default reconnaissance tasks...{Style.RESET_ALL}")
|
||||
|
||||
# Use default tasks instead
|
||||
initial_tasks = self._get_default_initial_tasks(target, available_tools)
|
||||
else:
|
||||
print(f"{Fore.GREEN}Got result from agent runner: {type(result)}{Style.RESET_ALL}")
|
||||
|
||||
# Check different possible response formats
|
||||
response_text = None
|
||||
if hasattr(result, "final_output"):
|
||||
response_text = result.final_output
|
||||
print(f"{Fore.CYAN}Using result.final_output{Style.RESET_ALL}")
|
||||
elif hasattr(result, "output"):
|
||||
response_text = result.output
|
||||
print(f"{Fore.CYAN}Using result.output{Style.RESET_ALL}")
|
||||
elif hasattr(result, "content"):
|
||||
response_text = result.content
|
||||
print(f"{Fore.CYAN}Using result.content{Style.RESET_ALL}")
|
||||
elif isinstance(result, str):
|
||||
response_text = result
|
||||
print(f"{Fore.CYAN}Using result as string{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"{Fore.RED}Unknown result format: {type(result)}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Result attributes: {dir(result)}{Style.RESET_ALL}")
|
||||
|
||||
# Try to get any text content from the result
|
||||
for attr in ['text', 'message', 'response', 'data']:
|
||||
if hasattr(result, attr):
|
||||
response_text = getattr(result, attr)
|
||||
print(f"{Fore.CYAN}Found text in result.{attr}{Style.RESET_ALL}")
|
||||
break
|
||||
|
||||
if not response_text:
|
||||
print(f"{Fore.RED}No response text found in result{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Using fallback initialization...{Style.RESET_ALL}")
|
||||
initialization_data = self._get_fallback_initialization(target, available_tools)
|
||||
else:
|
||||
print(f"{Fore.GREEN}Got response: {len(response_text)} characters{Style.RESET_ALL}")
|
||||
|
||||
# Parse the response
|
||||
initialization_data = self.reasoning_module.parse_tree_initialization_response(response_text)
|
||||
|
||||
if not initialization_data or not initialization_data.get('initial_tasks'):
|
||||
print(f"{Fore.YELLOW}No tasks parsed from response. Using fallback initialization.{Style.RESET_ALL}")
|
||||
initialization_data = self._get_fallback_initialization(target, available_tools)
|
||||
else:
|
||||
analysis = initialization_data.get('analysis', '')
|
||||
print(f"{Fore.CYAN}LLM determined approach: {analysis}{Style.RESET_ALL}")
|
||||
|
||||
# Create the structure and tasks as determined by the LLM
|
||||
structure_nodes = {}
|
||||
|
||||
# Create structure elements (phases, categories, etc.)
|
||||
for structure_element in initialization_data.get('structure', []):
|
||||
structure_node = TaskNode(
|
||||
description=structure_element.get('name', 'Unknown Structure'),
|
||||
parent_id=self.tree_manager.root_id,
|
||||
node_type=structure_element.get('type', 'phase'),
|
||||
attributes={
|
||||
"details": structure_element.get('description', ''),
|
||||
"justification": structure_element.get('justification', ''),
|
||||
"llm_created": True
|
||||
}
|
||||
)
|
||||
node_id = self.tree_manager.add_node(structure_node)
|
||||
structure_nodes[structure_element.get('name', 'Unknown')] = node_id
|
||||
|
||||
# Add initial tasks to their specified parents
|
||||
initial_tasks = initialization_data.get('initial_tasks', [])
|
||||
for task_data in initial_tasks:
|
||||
parent_name = task_data.get('parent', 'root')
|
||||
|
||||
# Determine parent node
|
||||
if parent_name == 'root':
|
||||
parent_id = self.tree_manager.root_id
|
||||
else:
|
||||
parent_id = structure_nodes.get(parent_name, self.tree_manager.root_id)
|
||||
|
||||
task_node = TaskNode(
|
||||
description=task_data.get('description', 'Unknown task'),
|
||||
parent_id=parent_id,
|
||||
tool_used=task_data.get('tool_suggestion'),
|
||||
priority=task_data.get('priority', 5),
|
||||
risk_level=task_data.get('risk_level', 'low'),
|
||||
attributes={'rationale': task_data.get('rationale', ''), 'llm_created': True}
|
||||
)
|
||||
self.tree_manager.add_node(task_node)
|
||||
|
||||
print(f"\n{Fore.GREEN}Agent Mode initialized with LLM-determined structure: {len(initialization_data.get('structure', []))} elements, {len(initial_tasks)} tasks.{Style.RESET_ALL}")
|
||||
|
||||
# Display the initial tree
|
||||
print(f"\n{Fore.CYAN}Initial Task Tree:{Style.RESET_ALL}")
|
||||
print(self.tree_manager.to_natural_language())
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Failed to initialize agent mode: {e}{Style.RESET_ALL}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
# Try to continue with default tasks even if there's an error
|
||||
print(f"{Fore.YELLOW}Attempting to continue with default tasks...{Style.RESET_ALL}")
|
||||
try:
|
||||
initialization_data = self._get_fallback_initialization(target, available_tools)
|
||||
|
||||
# Create structure and tasks from fallback
|
||||
structure_nodes = {}
|
||||
|
||||
# Create structure elements
|
||||
for structure_element in initialization_data.get('structure', []):
|
||||
structure_node = TaskNode(
|
||||
description=structure_element.get('name', 'Unknown Structure'),
|
||||
parent_id=self.tree_manager.root_id,
|
||||
node_type=structure_element.get('type', 'phase'),
|
||||
attributes={
|
||||
"details": structure_element.get('description', ''),
|
||||
"justification": structure_element.get('justification', ''),
|
||||
"fallback_created": True
|
||||
}
|
||||
)
|
||||
node_id = self.tree_manager.add_node(structure_node)
|
||||
structure_nodes[structure_element.get('name', 'Unknown')] = node_id
|
||||
|
||||
# Add initial tasks
|
||||
initial_tasks = initialization_data.get('initial_tasks', [])
|
||||
for task_data in initial_tasks:
|
||||
parent_name = task_data.get('parent', 'root')
|
||||
|
||||
if parent_name == 'root':
|
||||
parent_id = self.tree_manager.root_id
|
||||
else:
|
||||
parent_id = structure_nodes.get(parent_name, self.tree_manager.root_id)
|
||||
|
||||
task_node = TaskNode(
|
||||
description=task_data.get('description', 'Unknown task'),
|
||||
parent_id=parent_id,
|
||||
tool_used=task_data.get('tool_suggestion'),
|
||||
priority=task_data.get('priority', 5),
|
||||
risk_level=task_data.get('risk_level', 'low'),
|
||||
attributes={'rationale': task_data.get('rationale', ''), 'fallback_created': True}
|
||||
)
|
||||
self.tree_manager.add_node(task_node)
|
||||
|
||||
print(f"\n{Fore.GREEN}Agent Mode initialized with fallback structure: {len(initialization_data.get('structure', []))} elements, {len(initial_tasks)} tasks.{Style.RESET_ALL}")
|
||||
print(f"\n{Fore.CYAN}Initial Task Tree:{Style.RESET_ALL}")
|
||||
print(self.tree_manager.to_natural_language())
|
||||
return True
|
||||
except Exception as fallback_error:
|
||||
print(f"{Fore.RED}Fallback initialization also failed: {fallback_error}{Style.RESET_ALL}")
|
||||
|
||||
return False
|
||||
|
||||
async def run_autonomous_loop(self) -> None:
|
||||
"""Run the main autonomous agent loop."""
|
||||
print(f"\n{Fore.CYAN}Starting autonomous penetration test...{Style.RESET_ALL}")
|
||||
if self.max_iterations == 0:
|
||||
print(f"{Fore.YELLOW}Running until goal achieved or no more actions available{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"{Fore.YELLOW}Iteration limit: {self.max_iterations} iterations{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Press Ctrl+C to pause at any time.{Style.RESET_ALL}\n")
|
||||
|
||||
# Set effective limit - use a high number for "unlimited" but still have a safety limit
|
||||
effective_limit = self.max_iterations if self.max_iterations > 0 else 500
|
||||
|
||||
while self.iteration_count < effective_limit and not self.goal_achieved:
|
||||
try:
|
||||
if self.paused:
|
||||
await self._handle_pause()
|
||||
if self.paused: # Still paused after handling
|
||||
break
|
||||
|
||||
# Increment iteration count at the beginning
|
||||
self.iteration_count += 1
|
||||
|
||||
# Display current progress
|
||||
self._display_progress()
|
||||
|
||||
# Get next action from PTT
|
||||
next_action = await self._select_next_action()
|
||||
if not next_action:
|
||||
print(f"{Fore.YELLOW}No viable next actions found. Checking goal status...{Style.RESET_ALL}")
|
||||
await self._check_goal_achievement()
|
||||
break
|
||||
|
||||
# Execute the selected action
|
||||
await self._execute_action(next_action)
|
||||
|
||||
# Check goal achievement after every iteration
|
||||
await self._check_goal_achievement()
|
||||
|
||||
# If goal is achieved, stop the loop
|
||||
if self.goal_achieved:
|
||||
break
|
||||
|
||||
# Brief pause between actions
|
||||
await asyncio.sleep(2)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Fore.YELLOW}Pausing agent mode...{Style.RESET_ALL}")
|
||||
self.paused = True
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error in autonomous loop: {e}{Style.RESET_ALL}")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
# Display final reason for stopping
|
||||
if self.goal_achieved:
|
||||
print(f"\n{Fore.GREEN}Autonomous execution stopped: Goal achieved!{Style.RESET_ALL}")
|
||||
elif self.iteration_count >= effective_limit:
|
||||
print(f"\n{Fore.YELLOW}Autonomous execution stopped: Iteration limit reached ({effective_limit}){Style.RESET_ALL}")
|
||||
elif self.paused:
|
||||
print(f"\n{Fore.YELLOW}Autonomous execution stopped: User paused{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"\n{Fore.YELLOW}Autonomous execution stopped: No more viable actions{Style.RESET_ALL}")
|
||||
|
||||
# Final summary
|
||||
self._display_final_summary()
|
||||
|
||||
async def _select_next_action(self) -> Optional[Dict[str, Any]]:
|
||||
"""Select the next action based on PTT state."""
|
||||
# Get available tools
|
||||
available_tools = [server.name for server in self.connected_servers]
|
||||
|
||||
# Get prioritized candidate tasks
|
||||
candidates = self.tree_manager.get_candidate_tasks()
|
||||
if not candidates:
|
||||
return None
|
||||
|
||||
prioritized = self.tree_manager.prioritize_tasks(candidates)
|
||||
|
||||
print(f"\n{Fore.CYAN}Selecting next action...{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Available tools: {', '.join(available_tools)}{Style.RESET_ALL}")
|
||||
|
||||
# Query LLM for action selection
|
||||
selection_prompt = self.reasoning_module.get_next_action_prompt(available_tools)
|
||||
|
||||
try:
|
||||
result = await self.run_agent_func(
|
||||
selection_prompt,
|
||||
self.connected_servers,
|
||||
history=self.conversation_manager.get_history(),
|
||||
streaming=True, # Use streaming=True since it works in other modes
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
|
||||
# Handle response format variations
|
||||
response_text = None
|
||||
if hasattr(result, "final_output"):
|
||||
response_text = result.final_output
|
||||
elif hasattr(result, "output"):
|
||||
response_text = result.output
|
||||
elif isinstance(result, str):
|
||||
response_text = result
|
||||
|
||||
if response_text:
|
||||
action_data = self.reasoning_module.parse_next_action_response(response_text, available_tools)
|
||||
|
||||
if action_data:
|
||||
# Get the selected task
|
||||
task_index = action_data.get('selected_task_index', 1) - 1
|
||||
if 0 <= task_index < len(prioritized):
|
||||
selected_task = prioritized[task_index]
|
||||
|
||||
return {
|
||||
'task': selected_task,
|
||||
'command': action_data.get('command'),
|
||||
'tool': action_data.get('tool'),
|
||||
'rationale': action_data.get('rationale'),
|
||||
'expected_outcome': action_data.get('expected_outcome')
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error selecting next action: {e}{Style.RESET_ALL}")
|
||||
|
||||
# Fallback to first prioritized task
|
||||
if prioritized:
|
||||
return {'task': prioritized[0], 'command': None, 'tool': None}
|
||||
|
||||
return None
|
||||
|
||||
async def _execute_action(self, action: Dict[str, Any]) -> None:
|
||||
"""Execute a selected action."""
|
||||
task = action['task']
|
||||
command = action.get('command')
|
||||
tool = action.get('tool')
|
||||
available_tools = [server.name for server in self.connected_servers]
|
||||
|
||||
print(f"\n{Fore.CYAN}Executing: {task.description}{Style.RESET_ALL}")
|
||||
if action.get('rationale'):
|
||||
print(f"{Fore.WHITE}Rationale: {action['rationale']}{Style.RESET_ALL}")
|
||||
if command:
|
||||
print(f"{Fore.WHITE}Command: {command}{Style.RESET_ALL}")
|
||||
if tool:
|
||||
print(f"{Fore.WHITE}Using tool: {tool}{Style.RESET_ALL}")
|
||||
|
||||
# Check if suggested tool is available
|
||||
if tool and tool not in available_tools and tool != 'manual':
|
||||
print(f"{Fore.YELLOW}Tool '{tool}' not available. Available: {', '.join(available_tools)}{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Asking AI to adapt approach with available tools...{Style.RESET_ALL}")
|
||||
|
||||
# Let AI figure out how to adapt
|
||||
adaptation_query = f"""The task "{task.description}" was planned to use "{tool}" but that tool is not available.
|
||||
|
||||
Available tools: {', '.join(available_tools)}
|
||||
|
||||
Please adapt this task to work with the available tools. How would you accomplish this objective using {', '.join(available_tools)}?
|
||||
Be creative and think about alternative approaches that achieve the same security testing goal."""
|
||||
|
||||
command = adaptation_query
|
||||
|
||||
# Update task status to in-progress
|
||||
self.tree_manager.update_node(task.id, {'status': NodeStatus.IN_PROGRESS.value})
|
||||
|
||||
# Execute via agent
|
||||
execution_query = command if command else f"Perform the following task: {task.description}"
|
||||
|
||||
try:
|
||||
result = await self.run_agent_func(
|
||||
execution_query,
|
||||
self.connected_servers,
|
||||
history=self.conversation_manager.get_history(),
|
||||
streaming=True,
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
|
||||
# Handle response format variations
|
||||
response_text = None
|
||||
if hasattr(result, "final_output"):
|
||||
response_text = result.final_output
|
||||
elif hasattr(result, "output"):
|
||||
response_text = result.output
|
||||
elif isinstance(result, str):
|
||||
response_text = result
|
||||
|
||||
if response_text:
|
||||
# Update conversation history
|
||||
self.conversation_manager.add_dialogue(execution_query)
|
||||
self.conversation_manager.update_last_response(response_text)
|
||||
|
||||
# Update PTT based on results
|
||||
await self._update_tree_from_results(
|
||||
task,
|
||||
response_text,
|
||||
command or execution_query
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error executing action: {e}{Style.RESET_ALL}")
|
||||
self.tree_manager.update_node(task.id, {
|
||||
'status': NodeStatus.FAILED.value,
|
||||
'findings': f"Execution failed: {str(e)}"
|
||||
})
|
||||
|
||||
async def _update_tree_from_results(self, task: TaskNode, output: str, command: str) -> None:
|
||||
"""Update the PTT based on execution results."""
|
||||
try:
|
||||
# Create update prompt
|
||||
update_prompt = self.reasoning_module.get_tree_update_prompt(output, command, task)
|
||||
|
||||
# Get LLM analysis
|
||||
result = await self.run_agent_func(
|
||||
update_prompt,
|
||||
self.connected_servers,
|
||||
history=self.conversation_manager.get_history(),
|
||||
streaming=True,
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
|
||||
# Handle response format variations
|
||||
response_text = None
|
||||
if hasattr(result, "final_output"):
|
||||
response_text = result.final_output
|
||||
elif hasattr(result, "output"):
|
||||
response_text = result.output
|
||||
elif isinstance(result, str):
|
||||
response_text = result
|
||||
|
||||
if response_text:
|
||||
node_updates, new_tasks = self.reasoning_module.parse_tree_update_response(response_text)
|
||||
|
||||
# Update the executed node
|
||||
if node_updates:
|
||||
node_updates['timestamp'] = datetime.now().isoformat()
|
||||
node_updates['command_executed'] = command
|
||||
self.tree_manager.update_node(task.id, node_updates)
|
||||
|
||||
# Check if goal might be achieved before adding new tasks
|
||||
preliminary_goal_check = await self._quick_goal_check()
|
||||
|
||||
# Only add new tasks if goal is not achieved and they align with original goal
|
||||
if not preliminary_goal_check and new_tasks:
|
||||
# Filter tasks to ensure they align with the original goal
|
||||
filtered_tasks = self._filter_tasks_by_goal_scope(new_tasks)
|
||||
|
||||
for new_task_data in filtered_tasks:
|
||||
parent_phase = new_task_data.get('parent_phase', 'Phase 2')
|
||||
parent_node = self._find_phase_node(parent_phase)
|
||||
|
||||
if parent_node:
|
||||
new_task = TaskNode(
|
||||
description=new_task_data.get('description'),
|
||||
parent_id=parent_node.id,
|
||||
tool_used=new_task_data.get('tool_suggestion'),
|
||||
priority=new_task_data.get('priority', 5),
|
||||
risk_level=new_task_data.get('risk_level', 'low'),
|
||||
attributes={'rationale': new_task_data.get('rationale', '')}
|
||||
)
|
||||
self.tree_manager.add_node(new_task)
|
||||
|
||||
if filtered_tasks:
|
||||
print(f"{Fore.GREEN}PTT updated with {len(filtered_tasks)} new goal-aligned tasks.{Style.RESET_ALL}")
|
||||
if len(filtered_tasks) < len(new_tasks):
|
||||
print(f"{Fore.YELLOW}Filtered out {len(new_tasks) - len(filtered_tasks)} tasks that exceeded goal scope.{Style.RESET_ALL}")
|
||||
elif preliminary_goal_check:
|
||||
print(f"{Fore.GREEN}Goal appears to be achieved - not adding new tasks.{Style.RESET_ALL}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error updating PTT: {e}{Style.RESET_ALL}")
|
||||
# Default to marking as completed if update fails
|
||||
self.tree_manager.update_node(task.id, {
|
||||
'status': NodeStatus.COMPLETED.value,
|
||||
'timestamp': datetime.now().isoformat()
|
||||
})
|
||||
|
||||
def _filter_tasks_by_goal_scope(self, tasks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Filter tasks to ensure they align with the original goal scope."""
|
||||
goal_lower = self.tree_manager.goal.lower()
|
||||
filtered_tasks = []
|
||||
|
||||
# Define scope expansion keywords that should be avoided for simple goals
|
||||
expansion_keywords = ["exploit", "compromise", "attack", "penetrate", "shell", "backdoor", "privilege", "escalat"]
|
||||
|
||||
# Check if the original goal is simple information gathering
|
||||
info_keywords = ["check", "identify", "determine", "find", "discover", "enumerate", "list", "version", "banner"]
|
||||
is_simple_info_goal = any(keyword in goal_lower for keyword in info_keywords)
|
||||
|
||||
for task in tasks:
|
||||
task_desc_lower = task.get('description', '').lower()
|
||||
|
||||
# If it's a simple info goal, avoid adding exploitation tasks
|
||||
if is_simple_info_goal and any(keyword in task_desc_lower for keyword in expansion_keywords):
|
||||
print(f"{Fore.YELLOW}Skipping task that exceeds goal scope: {task.get('description', '')}{Style.RESET_ALL}")
|
||||
continue
|
||||
|
||||
filtered_tasks.append(task)
|
||||
|
||||
return filtered_tasks
|
||||
|
||||
async def _quick_goal_check(self) -> bool:
|
||||
"""Quick check if goal might be achieved based on completed tasks."""
|
||||
# Simple heuristic: if we have completed tasks with findings for info gathering goals
|
||||
goal_lower = self.tree_manager.goal.lower()
|
||||
info_keywords = ["check", "identify", "determine", "find", "discover", "enumerate", "list", "version", "banner"]
|
||||
|
||||
if any(keyword in goal_lower for keyword in info_keywords):
|
||||
# For info gathering goals, check if we have relevant findings
|
||||
for node in self.tree_manager.nodes.values():
|
||||
if node.status == NodeStatus.COMPLETED and node.findings:
|
||||
# Basic keyword matching for goal completion
|
||||
if "version" in goal_lower and "version" in node.findings.lower():
|
||||
return True
|
||||
if "banner" in goal_lower and "banner" in node.findings.lower():
|
||||
return True
|
||||
if any(keyword in goal_lower and keyword in node.description.lower() for keyword in info_keywords):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def _check_goal_achievement(self) -> None:
|
||||
"""Check if the primary goal has been achieved."""
|
||||
goal_prompt = self.reasoning_module.get_goal_check_prompt()
|
||||
|
||||
try:
|
||||
result = await self.run_agent_func(
|
||||
goal_prompt,
|
||||
self.connected_servers,
|
||||
history=self.conversation_manager.get_history(),
|
||||
streaming=True, # Use streaming=True
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
|
||||
# Handle response format variations
|
||||
response_text = None
|
||||
if hasattr(result, "final_output"):
|
||||
response_text = result.final_output
|
||||
elif hasattr(result, "output"):
|
||||
response_text = result.output
|
||||
elif isinstance(result, str):
|
||||
response_text = result
|
||||
|
||||
if response_text:
|
||||
goal_status = self.reasoning_module.parse_goal_check_response(response_text)
|
||||
|
||||
if goal_status.get('goal_achieved', False):
|
||||
confidence = goal_status.get('confidence', 0)
|
||||
if confidence >= 80:
|
||||
print(f"\n{Fore.GREEN}{'='*60}{Style.RESET_ALL}")
|
||||
print(f"{Fore.GREEN}GOAL ACHIEVED! (Confidence: {confidence}%){Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Evidence: {goal_status.get('evidence', 'N/A')}{Style.RESET_ALL}")
|
||||
print(f"{Fore.GREEN}{'='*60}{Style.RESET_ALL}\n")
|
||||
self.goal_achieved = True
|
||||
else:
|
||||
print(f"{Fore.YELLOW}Goal possibly achieved but confidence is low ({confidence}%). Continuing...{Style.RESET_ALL}")
|
||||
else:
|
||||
remaining = goal_status.get('remaining_objectives', 'Unknown')
|
||||
print(f"{Fore.YELLOW}Goal not yet achieved. Remaining: {remaining}{Style.RESET_ALL}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error checking goal achievement: {e}{Style.RESET_ALL}")
|
||||
|
||||
def _display_progress(self) -> None:
|
||||
"""Display current progress and statistics."""
|
||||
stats = self.tree_manager.get_statistics()
|
||||
elapsed = datetime.now() - self.start_time if self.start_time else None
|
||||
|
||||
print(f"\n{Fore.CYAN}{'='*60}{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Iteration {self.iteration_count} | Elapsed: {elapsed}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Tasks - Total: {stats['total_nodes']} | "
|
||||
f"Completed: {stats['status_counts'].get('completed', 0)} | "
|
||||
f"In Progress: {stats['status_counts'].get('in_progress', 0)} | "
|
||||
f"Pending: {stats['status_counts'].get('pending', 0)}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Vulnerabilities Found: {stats['status_counts'].get('vulnerable', 0)}{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}{'='*60}{Style.RESET_ALL}")
|
||||
|
||||
def _display_final_summary(self) -> None:
|
||||
"""Display final summary of the agent mode execution."""
|
||||
print(f"\n{Fore.CYAN}{'='*70}{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}AGENT MODE EXECUTION SUMMARY{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}")
|
||||
|
||||
summary = self.reasoning_module.generate_strategic_summary()
|
||||
print(f"{Fore.WHITE}{summary}{Style.RESET_ALL}")
|
||||
|
||||
# Calculate effective limit
|
||||
effective_limit = self.max_iterations if self.max_iterations > 0 else 500
|
||||
|
||||
# Execution Statistics
|
||||
if self.start_time:
|
||||
elapsed = datetime.now() - self.start_time
|
||||
print(f"\n{Fore.WHITE}Execution Statistics:{Style.RESET_ALL}")
|
||||
print(f"Total Execution Time: {elapsed}")
|
||||
print(f"Iterations Completed: {self.iteration_count}/{effective_limit}")
|
||||
|
||||
if self.iteration_count > 0:
|
||||
avg_time_per_iteration = elapsed.total_seconds() / self.iteration_count
|
||||
print(f"Average Time per Iteration: {avg_time_per_iteration:.1f} seconds")
|
||||
|
||||
# Calculate efficiency
|
||||
completion_rate = (self.iteration_count / effective_limit) * 100
|
||||
print(f"Completion Rate: {completion_rate:.1f}%")
|
||||
|
||||
# Estimate remaining time if stopped early
|
||||
if self.iteration_count < effective_limit and not self.goal_achieved:
|
||||
remaining_iterations = effective_limit - self.iteration_count
|
||||
estimated_remaining_time = remaining_iterations * avg_time_per_iteration
|
||||
print(f"Estimated Time for Full Run: {elapsed.total_seconds() + estimated_remaining_time:.0f} seconds")
|
||||
|
||||
print(f"{Fore.WHITE}Total Iterations: {self.iteration_count}{Style.RESET_ALL}")
|
||||
|
||||
if self.goal_achieved:
|
||||
print(f"\n{Fore.GREEN}PRIMARY GOAL ACHIEVED!{Style.RESET_ALL}")
|
||||
efficiency = "Excellent" if self.iteration_count <= 10 else "Good" if self.iteration_count <= 20 else "Extended"
|
||||
print(f"{Fore.GREEN}Efficiency: {efficiency} (achieved in {self.iteration_count} iterations){Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"\n{Fore.YELLOW}Primary goal not fully achieved within iteration limit.{Style.RESET_ALL}")
|
||||
if self.iteration_count >= effective_limit:
|
||||
print(f"{Fore.YELLOW}Consider increasing iteration limit for more thorough testing.{Style.RESET_ALL}")
|
||||
|
||||
print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}")
|
||||
|
||||
async def _handle_pause(self) -> None:
|
||||
"""Handle pause state and user options."""
|
||||
print(f"\n{Fore.YELLOW}Agent Mode Paused{Style.RESET_ALL}")
|
||||
|
||||
# Calculate effective limit
|
||||
effective_limit = self.max_iterations if self.max_iterations > 0 else 500
|
||||
|
||||
# Display current progress
|
||||
print(f"\n{Fore.MAGENTA}Progress Statistics:{Style.RESET_ALL}")
|
||||
print(f"Iterations: {self.iteration_count}/{effective_limit}")
|
||||
elapsed = datetime.now() - self.start_time if self.start_time else None
|
||||
if elapsed:
|
||||
print(f"Elapsed Time: {elapsed}")
|
||||
if self.iteration_count > 0:
|
||||
avg_time = elapsed.total_seconds() / self.iteration_count
|
||||
print(f"Average per Iteration: {avg_time:.1f} seconds")
|
||||
|
||||
print("\nOptions:")
|
||||
print("1. Resume execution")
|
||||
print("2. View current PTT")
|
||||
print("3. View detailed statistics")
|
||||
print("4. Save PTT state")
|
||||
print("5. Add manual task")
|
||||
print("6. Modify iteration limit")
|
||||
print("7. Exit agent mode")
|
||||
|
||||
while self.paused:
|
||||
choice = input(f"\n{Fore.GREEN}Select option (1-7): {Style.RESET_ALL}").strip()
|
||||
|
||||
if choice == '1':
|
||||
self.paused = False
|
||||
print(f"{Fore.GREEN}Resuming agent mode...{Style.RESET_ALL}")
|
||||
elif choice == '2':
|
||||
print(f"\n{Fore.CYAN}Current PTT State:{Style.RESET_ALL}")
|
||||
print(self.tree_manager.to_natural_language())
|
||||
elif choice == '3':
|
||||
self._display_progress()
|
||||
print(self.reasoning_module.generate_strategic_summary())
|
||||
elif choice == '4':
|
||||
self._save_ptt_state()
|
||||
elif choice == '5':
|
||||
await self._add_manual_task()
|
||||
elif choice == '6':
|
||||
self._modify_iteration_limit()
|
||||
elif choice == '7':
|
||||
print(f"{Fore.YELLOW}Exiting agent mode...{Style.RESET_ALL}")
|
||||
break
|
||||
else:
|
||||
print(f"{Fore.RED}Invalid choice.{Style.RESET_ALL}")
|
||||
|
||||
def _modify_iteration_limit(self) -> None:
|
||||
"""Allow user to modify the iteration limit during execution."""
|
||||
try:
|
||||
print(f"\n{Fore.CYAN}Modify Iteration Limit{Style.RESET_ALL}")
|
||||
print(f"Current limit: {self.max_iterations}")
|
||||
print(f"Iterations completed: {self.iteration_count}")
|
||||
print(f"Iterations remaining: {self.max_iterations - self.iteration_count}")
|
||||
|
||||
new_limit_input = input(f"New iteration limit (current: {self.max_iterations}): ").strip()
|
||||
|
||||
if new_limit_input:
|
||||
try:
|
||||
new_limit = int(new_limit_input)
|
||||
|
||||
# Ensure new limit is at least the number of iterations already completed
|
||||
min_limit = self.iteration_count + 1 # Allow at least 1 more iteration
|
||||
if new_limit < min_limit:
|
||||
print(f"{Fore.YELLOW}Minimum limit is {min_limit} (iterations already completed + 1){Style.RESET_ALL}")
|
||||
new_limit = min_limit
|
||||
|
||||
# Apply reasonable maximum
|
||||
new_limit = min(new_limit, 200)
|
||||
|
||||
self.max_iterations = new_limit
|
||||
print(f"{Fore.GREEN}Iteration limit updated to: {new_limit}{Style.RESET_ALL}")
|
||||
|
||||
except ValueError:
|
||||
print(f"{Fore.RED}Invalid input. Please enter a number.{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"{Fore.YELLOW}No change made.{Style.RESET_ALL}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error modifying iteration limit: {e}{Style.RESET_ALL}")
|
||||
|
||||
async def _add_manual_task(self) -> None:
|
||||
"""Allow user to manually add a task to the PTT."""
|
||||
try:
|
||||
print(f"\n{Fore.CYAN}Add Manual Task{Style.RESET_ALL}")
|
||||
|
||||
# Get task details from user
|
||||
description = input("Task description: ").strip()
|
||||
if not description:
|
||||
print(f"{Fore.RED}Task description required.{Style.RESET_ALL}")
|
||||
return
|
||||
|
||||
print("Select phase:")
|
||||
print("1. Phase 1: Reconnaissance")
|
||||
print("2. Phase 2: Vulnerability Assessment")
|
||||
print("3. Phase 3: Exploitation")
|
||||
print("4. Phase 4: Post-Exploitation")
|
||||
|
||||
phase_choice = input("Phase (1-4): ").strip()
|
||||
phase_map = {
|
||||
'1': 'Phase 1',
|
||||
'2': 'Phase 2',
|
||||
'3': 'Phase 3',
|
||||
'4': 'Phase 4'
|
||||
}
|
||||
|
||||
phase = phase_map.get(phase_choice, 'Phase 2')
|
||||
parent_node = self._find_phase_node(phase)
|
||||
|
||||
if not parent_node:
|
||||
print(f"{Fore.RED}Phase not found.{Style.RESET_ALL}")
|
||||
return
|
||||
|
||||
# Get priority
|
||||
try:
|
||||
priority = int(input("Priority (1-10, default 5): ").strip() or "5")
|
||||
priority = max(1, min(10, priority))
|
||||
except:
|
||||
priority = 5
|
||||
|
||||
# Get risk level
|
||||
print("Risk level:")
|
||||
print("1. Low")
|
||||
print("2. Medium")
|
||||
print("3. High")
|
||||
risk_choice = input("Risk (1-3, default 2): ").strip()
|
||||
risk_map = {'1': 'low', '2': 'medium', '3': 'high'}
|
||||
risk_level = risk_map.get(risk_choice, 'medium')
|
||||
|
||||
# Create the task
|
||||
from core.task_tree_manager import TaskNode
|
||||
manual_task = TaskNode(
|
||||
description=description,
|
||||
parent_id=parent_node.id,
|
||||
tool_used='manual',
|
||||
priority=priority,
|
||||
risk_level=risk_level,
|
||||
attributes={'manual_addition': True, 'added_by_user': True}
|
||||
)
|
||||
|
||||
self.tree_manager.add_node(manual_task)
|
||||
print(f"{Fore.GREEN}Manual task added to {phase}.{Style.RESET_ALL}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error adding manual task: {e}{Style.RESET_ALL}")
|
||||
|
||||
def _save_ptt_state(self) -> None:
|
||||
"""Save the current PTT state to file."""
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"reports/ptt_state_{timestamp}.json"
|
||||
|
||||
try:
|
||||
import os
|
||||
os.makedirs("reports", exist_ok=True)
|
||||
|
||||
with open(filename, 'w') as f:
|
||||
f.write(self.tree_manager.to_json())
|
||||
|
||||
print(f"{Fore.GREEN}PTT state saved to: {filename}{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Failed to save PTT state: {e}{Style.RESET_ALL}")
|
||||
|
||||
def _find_phase_node(self, phase_description: str) -> Optional[TaskNode]:
|
||||
"""Find a structure node by description (phase, category, etc.)."""
|
||||
for node in self.tree_manager.nodes.values():
|
||||
# Look for any non-task node that matches the description
|
||||
if node.node_type != "task" and phase_description.lower() in node.description.lower():
|
||||
return node
|
||||
|
||||
# If no exact match, try to find any suitable parent node
|
||||
# Return root if no structure nodes exist
|
||||
return self.tree_manager.nodes.get(self.tree_manager.root_id)
|
||||
|
||||
def get_ptt_for_reporting(self) -> TaskTreeManager:
|
||||
"""Get the PTT for report generation."""
|
||||
return self.tree_manager
|
||||
|
||||
def _get_fallback_initialization(self, target: str, available_tools: List[str]) -> Dict[str, Any]:
|
||||
"""Return minimal fallback initialization when LLM fails."""
|
||||
print(f"{Fore.YELLOW}Using minimal fallback initialization. The system will rely on dynamic task generation.{Style.RESET_ALL}")
|
||||
|
||||
return {
|
||||
'analysis': 'Fallback initialization - LLM will determine structure dynamically during execution',
|
||||
'structure': [],
|
||||
'initial_tasks': []
|
||||
}
|
||||
@@ -1,259 +0,0 @@
|
||||
"""Agent execution and query processing for GHOSTCREW."""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import traceback
|
||||
from typing import List, Dict, Optional, Any
|
||||
from colorama import Fore, Style
|
||||
from agents import Agent, RunConfig, Runner, ModelSettings
|
||||
from openai.types.responses import ResponseTextDeltaEvent, ResponseContentPartDoneEvent
|
||||
from core.model_manager import model_manager
|
||||
from config.constants import BASE_INSTRUCTIONS, CONNECTION_RETRY_DELAY, DEFAULT_KNOWLEDGE_BASE_PATH
|
||||
from config.app_config import app_config
|
||||
import os
|
||||
|
||||
|
||||
class AgentRunner:
|
||||
"""Handles AI agent query processing and execution."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the agent runner."""
|
||||
self.model_provider = model_manager.get_model_provider()
|
||||
self.client = app_config.get_openai_client()
|
||||
|
||||
async def run_agent(
|
||||
self,
|
||||
query: str,
|
||||
mcp_servers: List[Any], # Use Any to avoid import issues
|
||||
history: Optional[List[Dict[str, str]]] = None,
|
||||
streaming: bool = True,
|
||||
kb_instance: Any = None
|
||||
) -> Any:
|
||||
"""
|
||||
Run cybersecurity agent with connected MCP servers, supporting streaming output and conversation history.
|
||||
|
||||
Args:
|
||||
query: User's natural language query
|
||||
mcp_servers: List of connected MCPServerStdio instances
|
||||
history: Conversation history, list containing user questions and AI answers
|
||||
streaming: Whether to use streaming output
|
||||
kb_instance: Knowledge base instance for retrieval
|
||||
|
||||
Returns:
|
||||
Agent execution result
|
||||
"""
|
||||
# If no history is provided, initialize an empty list
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
try:
|
||||
# Build instructions containing conversation history
|
||||
instructions = self._build_instructions(mcp_servers, history, query, kb_instance)
|
||||
|
||||
# Calculate max output tokens
|
||||
max_output_tokens = model_manager.calculate_max_output_tokens(instructions, query)
|
||||
|
||||
# Set model settings based on whether there are connected MCP servers
|
||||
model_settings = self._create_model_settings(mcp_servers, max_output_tokens)
|
||||
|
||||
# Create agent
|
||||
secure_agent = Agent(
|
||||
name="Cybersecurity Expert",
|
||||
instructions=instructions,
|
||||
mcp_servers=mcp_servers,
|
||||
model_settings=model_settings
|
||||
)
|
||||
|
||||
print(f"{Fore.CYAN}\nProcessing query: {Fore.WHITE}{query}{Style.RESET_ALL}\n")
|
||||
|
||||
if streaming:
|
||||
return await self._run_streaming(secure_agent, query)
|
||||
else:
|
||||
# Non-streaming mode could be implemented here if needed
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error processing agent request: {e}{Style.RESET_ALL}", flush=True)
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
def _build_instructions(
|
||||
self,
|
||||
mcp_servers: List[Any], # Use Any to avoid import issues
|
||||
history: List[Dict[str, str]],
|
||||
query: str,
|
||||
kb_instance: Any
|
||||
) -> str:
|
||||
"""Build agent instructions with context."""
|
||||
instructions = BASE_INSTRUCTIONS
|
||||
|
||||
# Add information about available tools
|
||||
if mcp_servers:
|
||||
available_tool_names = [server.name for server in mcp_servers]
|
||||
if available_tool_names:
|
||||
instructions += f"\n\nYou have access to the following tools: {', '.join(available_tool_names)}."
|
||||
|
||||
# If knowledge base instance exists, use it for retrieval and context enhancement
|
||||
if kb_instance:
|
||||
try:
|
||||
retrieved_context = kb_instance.search(query)
|
||||
if retrieved_context:
|
||||
# Add file path information to make LLM aware of actual files
|
||||
available_files = []
|
||||
if os.path.exists(DEFAULT_KNOWLEDGE_BASE_PATH):
|
||||
for filename in os.listdir(DEFAULT_KNOWLEDGE_BASE_PATH):
|
||||
filepath = os.path.join(DEFAULT_KNOWLEDGE_BASE_PATH, filename)
|
||||
if os.path.isfile(filepath):
|
||||
available_files.append(filename)
|
||||
|
||||
file_info = ""
|
||||
if available_files:
|
||||
file_info = f"\n\nIMPORTANT: The following actual files are available in the knowledge folder that you can reference by path:\n"
|
||||
for filename in available_files:
|
||||
file_info += f"- {DEFAULT_KNOWLEDGE_BASE_PATH}/{filename}\n"
|
||||
file_info += "\nWhen using security tools that require external files, you can reference these files by their full path.\n"
|
||||
file_info += f"ONLY use {DEFAULT_KNOWLEDGE_BASE_PATH}/ for files.\n"
|
||||
|
||||
instructions = f"Based on the following knowledge base information:\n{retrieved_context}{file_info}\n\n{instructions}"
|
||||
print(f"{Fore.MAGENTA}Relevant information retrieved from knowledge base.{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Failed to retrieve information from knowledge base: {e}{Style.RESET_ALL}")
|
||||
|
||||
# If there's conversation history, add it to the instructions
|
||||
if history:
|
||||
instructions += "\n\nBelow is the previous conversation history, please refer to this information to answer the user's question:\n"
|
||||
for i, entry in enumerate(history):
|
||||
instructions += f"\nUser question {i+1}: {entry['user_query']}"
|
||||
if 'ai_response' in entry and entry['ai_response']:
|
||||
instructions += f"\nAI answer {i+1}: {entry['ai_response']}\n"
|
||||
|
||||
return instructions
|
||||
|
||||
def _create_model_settings(self, mcp_servers: List[Any], max_output_tokens: int) -> ModelSettings:
|
||||
"""Create model settings based on available tools."""
|
||||
if mcp_servers:
|
||||
# With tools available, enable tool_choice and parallel_tool_calls
|
||||
return ModelSettings(
|
||||
temperature=0.6,
|
||||
top_p=0.9,
|
||||
max_tokens=max_output_tokens,
|
||||
tool_choice="auto",
|
||||
parallel_tool_calls=False,
|
||||
truncation="auto"
|
||||
)
|
||||
else:
|
||||
# Without tools, don't set tool_choice or parallel_tool_calls
|
||||
return ModelSettings(
|
||||
temperature=0.6,
|
||||
top_p=0.9,
|
||||
max_tokens=max_output_tokens,
|
||||
truncation="auto"
|
||||
)
|
||||
|
||||
async def _run_streaming(self, agent: Agent, query: str) -> Any:
|
||||
"""Run agent with streaming output."""
|
||||
result = Runner.run_streamed(
|
||||
agent,
|
||||
input=query,
|
||||
max_turns=10,
|
||||
run_config=RunConfig(
|
||||
model_provider=self.model_provider,
|
||||
trace_include_sensitive_data=True,
|
||||
handoff_input_filter=None
|
||||
)
|
||||
)
|
||||
|
||||
print(f"{Fore.GREEN}Reply:{Style.RESET_ALL}", end="", flush=True)
|
||||
|
||||
try:
|
||||
async for event in result.stream_events():
|
||||
await self._handle_stream_event(event)
|
||||
except Exception as e:
|
||||
await self._handle_stream_error(e)
|
||||
|
||||
print(f"\n\n{Fore.GREEN}Query completed!{Style.RESET_ALL}")
|
||||
return result
|
||||
|
||||
async def _handle_stream_event(self, event: Any) -> None:
|
||||
"""Handle individual stream events."""
|
||||
if event.type == "raw_response_event":
|
||||
if isinstance(event.data, ResponseTextDeltaEvent):
|
||||
print(f"{Fore.WHITE}{event.data.delta}{Style.RESET_ALL}", end="", flush=True)
|
||||
elif isinstance(event.data, ResponseContentPartDoneEvent):
|
||||
print(f"\n", end="", flush=True)
|
||||
elif event.type == "run_item_stream_event":
|
||||
if event.item.type == "tool_call_item":
|
||||
await self._handle_tool_call(event.item)
|
||||
elif event.item.type == "tool_call_output_item":
|
||||
await self._handle_tool_output(event.item)
|
||||
|
||||
async def _handle_tool_call(self, item: Any) -> None:
|
||||
"""Handle tool call events."""
|
||||
raw_item = getattr(item, "raw_item", None)
|
||||
tool_name = ""
|
||||
tool_args = {}
|
||||
|
||||
if raw_item:
|
||||
tool_name = getattr(raw_item, "name", "Unknown tool")
|
||||
tool_str = getattr(raw_item, "arguments", "{}")
|
||||
if isinstance(tool_str, str):
|
||||
try:
|
||||
tool_args = json.loads(tool_str)
|
||||
except json.JSONDecodeError:
|
||||
tool_args = {"raw_arguments": tool_str}
|
||||
|
||||
print(f"\n{Fore.CYAN}Tool name: {tool_name}{Style.RESET_ALL}", flush=True)
|
||||
print(f"\n{Fore.CYAN}Tool parameters: {tool_args}{Style.RESET_ALL}", flush=True)
|
||||
|
||||
async def _handle_tool_output(self, item: Any) -> None:
|
||||
"""Handle tool output events."""
|
||||
raw_item = getattr(item, "raw_item", None)
|
||||
tool_id = "Unknown tool ID"
|
||||
|
||||
if isinstance(raw_item, dict) and "call_id" in raw_item:
|
||||
tool_id = raw_item["call_id"]
|
||||
|
||||
output = getattr(item, "output", "Unknown output")
|
||||
output_text = self._parse_tool_output(output)
|
||||
|
||||
print(f"\n{Fore.GREEN}Tool call {tool_id} returned result: {output_text}{Style.RESET_ALL}", flush=True)
|
||||
|
||||
def _parse_tool_output(self, output: Any) -> str:
|
||||
"""Parse tool output into readable text."""
|
||||
if isinstance(output, str) and (output.startswith("{") or output.startswith("[")):
|
||||
try:
|
||||
output_data = json.loads(output)
|
||||
if isinstance(output_data, dict):
|
||||
if 'type' in output_data and output_data['type'] == 'text' and 'text' in output_data:
|
||||
return output_data['text']
|
||||
elif 'text' in output_data:
|
||||
return output_data['text']
|
||||
elif 'content' in output_data:
|
||||
return output_data['content']
|
||||
else:
|
||||
return json.dumps(output_data, ensure_ascii=False, indent=2)
|
||||
except json.JSONDecodeError:
|
||||
return f"Unparsable JSON output: {output}"
|
||||
return str(output)
|
||||
|
||||
async def _handle_stream_error(self, error: Exception) -> None:
|
||||
"""Handle streaming errors."""
|
||||
print(f"{Fore.RED}Error processing streamed response event: {error}{Style.RESET_ALL}", flush=True)
|
||||
|
||||
if 'Connection error' in str(error):
|
||||
print(f"{Fore.YELLOW}Connection error details:{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}1. Check network connection{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}2. Verify API address: {app_config.base_url}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}3. Check API key validity{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}4. Try reconnecting...{Style.RESET_ALL}")
|
||||
await asyncio.sleep(CONNECTION_RETRY_DELAY)
|
||||
|
||||
try:
|
||||
await self.client.connect()
|
||||
print(f"{Fore.GREEN}Reconnected successfully{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Reconnection failed: {e}{Style.RESET_ALL}")
|
||||
|
||||
|
||||
# Create singleton instance
|
||||
agent_runner = AgentRunner()
|
||||
@@ -1,73 +0,0 @@
|
||||
"""Model management and AI model setup for GHOSTCREW."""
|
||||
|
||||
import tiktoken
|
||||
from agents import Model, ModelProvider, OpenAIChatCompletionsModel
|
||||
from config.app_config import app_config
|
||||
from config.constants import MAX_TOTAL_TOKENS, RESPONSE_BUFFER
|
||||
|
||||
|
||||
class DefaultModelProvider(ModelProvider):
|
||||
"""Model provider using OpenAI compatible interface."""
|
||||
|
||||
def get_model(self, model_name: str) -> Model:
|
||||
"""Get a model instance with the specified name."""
|
||||
return OpenAIChatCompletionsModel(
|
||||
model=model_name or app_config.model_name,
|
||||
openai_client=app_config.get_openai_client()
|
||||
)
|
||||
|
||||
|
||||
class ModelManager:
|
||||
"""Manages AI model operations and token counting."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the model manager."""
|
||||
self.model_provider = DefaultModelProvider()
|
||||
self.model_name = app_config.model_name
|
||||
|
||||
@staticmethod
|
||||
def count_tokens(text: str, model_name: str = None) -> int:
|
||||
"""
|
||||
Count tokens in the given text.
|
||||
|
||||
Args:
|
||||
text: The text to count tokens for
|
||||
model_name: The model name to use for encoding (defaults to configured model)
|
||||
|
||||
Returns:
|
||||
Number of tokens in the text
|
||||
"""
|
||||
try:
|
||||
model = model_name or app_config.model_name
|
||||
encoding = tiktoken.encoding_for_model(model)
|
||||
return len(encoding.encode(text))
|
||||
except Exception:
|
||||
# Fall back to approximate counting if tiktoken fails
|
||||
return len(text.split())
|
||||
|
||||
@staticmethod
|
||||
def calculate_max_output_tokens(input_text: str, query: str) -> int:
|
||||
"""
|
||||
Calculate the maximum output tokens based on input size.
|
||||
|
||||
Args:
|
||||
input_text: The base instructions or context
|
||||
query: The user query
|
||||
|
||||
Returns:
|
||||
Maximum number of output tokens
|
||||
"""
|
||||
input_token_estimate = ModelManager.count_tokens(input_text) + ModelManager.count_tokens(query)
|
||||
|
||||
max_output_tokens = max(512, MAX_TOTAL_TOKENS - input_token_estimate)
|
||||
max_output_tokens = min(max_output_tokens, RESPONSE_BUFFER)
|
||||
|
||||
return max_output_tokens
|
||||
|
||||
def get_model_provider(self) -> ModelProvider:
|
||||
"""Get the model provider instance."""
|
||||
return self.model_provider
|
||||
|
||||
|
||||
# Create a singleton instance
|
||||
model_manager = ModelManager()
|
||||
@@ -1,448 +0,0 @@
|
||||
"""Main controller for GHOSTCREW application."""
|
||||
|
||||
import asyncio
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Optional, List, Dict, Any
|
||||
from colorama import Fore, Style
|
||||
|
||||
from config.constants import (
|
||||
ASCII_TITLE, VERSION, WELCOME_MESSAGE, EXIT_MESSAGE, SEPARATOR,
|
||||
KB_PROMPT, MCP_PROMPT, ERROR_NO_WORKFLOWS, ERROR_NO_REPORTING,
|
||||
DEFAULT_KNOWLEDGE_BASE_PATH
|
||||
)
|
||||
from config.app_config import app_config
|
||||
from core.agent_runner import agent_runner
|
||||
from core.agent_mode_controller import AgentModeController
|
||||
from tools.mcp_manager import MCPManager
|
||||
from ui.menu_system import MenuSystem
|
||||
from ui.conversation_manager import ConversationManager
|
||||
from workflows.workflow_engine import WorkflowEngine
|
||||
from rag.knowledge_base import Kb
|
||||
|
||||
|
||||
class PentestAgent:
|
||||
"""Main application controller for GHOSTCREW."""
|
||||
|
||||
def __init__(self, MCPServerStdio=None, MCPServerSse=None):
|
||||
"""
|
||||
Initialize the pentest agent controller.
|
||||
|
||||
Args:
|
||||
MCPServerStdio: MCP server stdio class
|
||||
MCPServerSse: MCP server SSE class
|
||||
"""
|
||||
self.app_config = app_config
|
||||
self.agent_runner = agent_runner
|
||||
self.mcp_manager = MCPManager(MCPServerStdio, MCPServerSse)
|
||||
self.menu_system = MenuSystem()
|
||||
self.conversation_manager = ConversationManager()
|
||||
self.workflow_engine = WorkflowEngine()
|
||||
self.kb_instance = None
|
||||
self.reporting_available = self._check_reporting_available()
|
||||
self.agent_mode_controller = None # Will be initialized when needed
|
||||
|
||||
@staticmethod
|
||||
def _check_reporting_available() -> bool:
|
||||
"""Check if reporting module is available."""
|
||||
try:
|
||||
from reporting.generators import generate_report_from_workflow
|
||||
return True
|
||||
except ImportError:
|
||||
print(ERROR_NO_REPORTING)
|
||||
return False
|
||||
|
||||
def display_welcome(self) -> None:
|
||||
"""Display welcome message and ASCII art."""
|
||||
print(ASCII_TITLE)
|
||||
print(f"{Fore.WHITE}GHOSTCREW v{VERSION}{Style.RESET_ALL}")
|
||||
print(WELCOME_MESSAGE)
|
||||
print(EXIT_MESSAGE)
|
||||
print(f"{SEPARATOR}\n")
|
||||
|
||||
def setup_knowledge_base(self) -> None:
|
||||
"""Setup knowledge base if requested by user."""
|
||||
use_kb_input = input(KB_PROMPT).strip().lower()
|
||||
if use_kb_input == 'yes':
|
||||
try:
|
||||
self.kb_instance = Kb(DEFAULT_KNOWLEDGE_BASE_PATH)
|
||||
print(f"{Fore.GREEN}Knowledge base loaded successfully!{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Failed to load knowledge base: {e}{Style.RESET_ALL}")
|
||||
self.kb_instance = None
|
||||
|
||||
async def setup_mcp_tools(self) -> tuple:
|
||||
"""Setup MCP tools and return server instances."""
|
||||
use_mcp_input = input(MCP_PROMPT).strip().lower()
|
||||
return await self.mcp_manager.setup_mcp_tools(use_mcp_input == 'yes')
|
||||
|
||||
async def run_interactive_mode(self, connected_servers: List) -> None:
|
||||
"""Run interactive chat mode."""
|
||||
self.menu_system.display_interactive_mode_intro()
|
||||
|
||||
while True:
|
||||
user_query = self.menu_system.get_user_input()
|
||||
|
||||
# Handle special commands
|
||||
if user_query.lower() in ["quit", "exit"]:
|
||||
self.menu_system.display_exit_message()
|
||||
return True # Signal to exit the entire application
|
||||
|
||||
if user_query.lower() == "menu":
|
||||
break # Return to main menu
|
||||
|
||||
# Handle empty input
|
||||
if not user_query:
|
||||
self.menu_system.display_no_query_message()
|
||||
continue
|
||||
|
||||
# Handle multi-line mode request
|
||||
if user_query.lower() == "multi":
|
||||
user_query = self.menu_system.get_multi_line_input()
|
||||
if not user_query:
|
||||
continue
|
||||
|
||||
# Process the query
|
||||
await self._process_user_query(user_query, connected_servers)
|
||||
|
||||
self.menu_system.display_ready_message()
|
||||
|
||||
return False # Don't exit application
|
||||
|
||||
async def _process_user_query(self, query: str, connected_servers: List) -> None:
|
||||
"""Process a user query through the agent."""
|
||||
# Add dialogue to history
|
||||
self.conversation_manager.add_dialogue(query)
|
||||
|
||||
# Run the agent
|
||||
result = await agent_runner.run_agent(
|
||||
query,
|
||||
connected_servers,
|
||||
history=self.conversation_manager.get_history(),
|
||||
streaming=True,
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
|
||||
# Update the response in history
|
||||
if result and hasattr(result, "final_output"):
|
||||
self.conversation_manager.update_last_response(result.final_output)
|
||||
|
||||
async def run_automated_mode(self, connected_servers: List) -> None:
|
||||
"""Run workflows mode."""
|
||||
if not self.workflow_engine.is_available():
|
||||
print(ERROR_NO_WORKFLOWS)
|
||||
self.menu_system.press_enter_to_continue()
|
||||
return
|
||||
|
||||
if not connected_servers:
|
||||
self.menu_system.display_workflow_requirements_message()
|
||||
return
|
||||
|
||||
while True:
|
||||
workflow_list = self.workflow_engine.show_automated_menu()
|
||||
if not workflow_list:
|
||||
break
|
||||
|
||||
try:
|
||||
choice = input(f"\n{Fore.GREEN}Select workflow (1-{len(workflow_list)+1}): {Style.RESET_ALL}").strip()
|
||||
|
||||
if not choice.isdigit():
|
||||
self.menu_system.display_invalid_input()
|
||||
continue
|
||||
|
||||
choice = int(choice)
|
||||
|
||||
if choice == len(workflow_list) + 1:
|
||||
# Back to main menu
|
||||
break
|
||||
|
||||
if 1 <= choice <= len(workflow_list):
|
||||
await self._execute_workflow(workflow_list[choice-1], connected_servers)
|
||||
else:
|
||||
self.menu_system.display_invalid_choice()
|
||||
|
||||
except ValueError:
|
||||
self.menu_system.display_invalid_input()
|
||||
except KeyboardInterrupt:
|
||||
self.menu_system.display_operation_cancelled()
|
||||
break
|
||||
|
||||
async def run_agent_mode(self, connected_servers: List) -> None:
|
||||
"""Run autonomous agent mode with PTT."""
|
||||
if not connected_servers:
|
||||
self.menu_system.display_agent_mode_requirements_message()
|
||||
return
|
||||
|
||||
# Display introduction
|
||||
self.menu_system.display_agent_mode_intro()
|
||||
|
||||
# Get agent mode parameters
|
||||
params = self.menu_system.get_agent_mode_params()
|
||||
if not params:
|
||||
return
|
||||
|
||||
# Initialize agent mode controller
|
||||
self.agent_mode_controller = AgentModeController(
|
||||
self.mcp_manager,
|
||||
self.conversation_manager,
|
||||
self.kb_instance
|
||||
)
|
||||
|
||||
try:
|
||||
# Initialize agent mode
|
||||
init_success = await self.agent_mode_controller.initialize_agent_mode(
|
||||
goal=params['goal'],
|
||||
target=params['target'],
|
||||
constraints=params['constraints'],
|
||||
connected_servers=connected_servers,
|
||||
run_agent_func=agent_runner.run_agent
|
||||
)
|
||||
|
||||
if init_success:
|
||||
# Run the autonomous loop
|
||||
await self.agent_mode_controller.run_autonomous_loop()
|
||||
|
||||
# Handle post-execution options
|
||||
await self._handle_agent_mode_completion()
|
||||
else:
|
||||
print(f"{Fore.RED}Failed to initialize agent mode.{Style.RESET_ALL}")
|
||||
self.menu_system.press_enter_to_continue()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Fore.YELLOW}Agent mode interrupted by user.{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error in agent mode: {e}{Style.RESET_ALL}")
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
self.menu_system.press_enter_to_continue()
|
||||
|
||||
async def _handle_agent_mode_completion(self) -> None:
|
||||
"""Handle post-execution options for agent mode."""
|
||||
# Ask if user wants to generate a report
|
||||
if self.reporting_available and self.menu_system.ask_generate_report():
|
||||
try:
|
||||
# Generate report from PTT
|
||||
from reporting.generators import generate_report_from_ptt
|
||||
|
||||
ptt = self.agent_mode_controller.get_ptt_for_reporting()
|
||||
report_path = await generate_report_from_ptt(
|
||||
ptt,
|
||||
self.conversation_manager.get_history(),
|
||||
run_agent_func=agent_runner.run_agent,
|
||||
connected_servers=self.mcp_manager.connected_servers if hasattr(self.mcp_manager, 'connected_servers') else [],
|
||||
kb_instance=self.kb_instance
|
||||
)
|
||||
|
||||
if report_path:
|
||||
self.menu_system.display_report_generated(report_path)
|
||||
else:
|
||||
print(f"{Fore.YELLOW}Report generation returned no path.{Style.RESET_ALL}")
|
||||
|
||||
except ImportError:
|
||||
# Fallback if PTT report generation not available
|
||||
print(f"{Fore.YELLOW}PTT report generation not available. Saving raw data...{Style.RESET_ALL}")
|
||||
self._save_agent_mode_data()
|
||||
except Exception as e:
|
||||
self.menu_system.display_report_error(e)
|
||||
self._save_agent_mode_data()
|
||||
|
||||
# Ask about saving raw data
|
||||
elif self.menu_system.ask_save_raw_history():
|
||||
self._save_agent_mode_data()
|
||||
|
||||
def _save_agent_mode_data(self) -> None:
|
||||
"""Save agent mode execution data."""
|
||||
try:
|
||||
import os
|
||||
import json
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
# Create reports directory if it doesn't exist
|
||||
os.makedirs("reports", exist_ok=True)
|
||||
|
||||
# Save conversation history
|
||||
history_file = f"reports/agent_mode_history_{timestamp}.json"
|
||||
with open(history_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.conversation_manager.get_history(), f, indent=2)
|
||||
|
||||
# Save PTT state
|
||||
if self.agent_mode_controller:
|
||||
ptt_file = f"reports/agent_mode_ptt_{timestamp}.json"
|
||||
with open(ptt_file, 'w', encoding='utf-8') as f:
|
||||
f.write(self.agent_mode_controller.tree_manager.to_json())
|
||||
|
||||
print(f"{Fore.GREEN}Agent mode data saved to reports/ directory{Style.RESET_ALL}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Failed to save agent mode data: {e}{Style.RESET_ALL}")
|
||||
|
||||
async def _execute_workflow(self, workflow_info: tuple, connected_servers: List) -> None:
|
||||
"""Execute a selected workflow."""
|
||||
workflow_key, workflow_name = workflow_info
|
||||
workflow = self.workflow_engine.get_workflow(workflow_key)
|
||||
|
||||
if not workflow:
|
||||
print(f"{Fore.RED}Error loading workflow.{Style.RESET_ALL}")
|
||||
return
|
||||
|
||||
target = self.menu_system.get_workflow_target()
|
||||
if not target:
|
||||
return
|
||||
|
||||
if not self.menu_system.confirm_workflow_execution(workflow['name'], target):
|
||||
self.menu_system.display_workflow_cancelled()
|
||||
return
|
||||
|
||||
# Store initial workflow data
|
||||
workflow_start_time = datetime.now()
|
||||
initial_history_length = self.conversation_manager.get_dialogue_count()
|
||||
|
||||
# Execute the workflow
|
||||
await self.workflow_engine.run_automated_workflow(
|
||||
workflow,
|
||||
target,
|
||||
connected_servers,
|
||||
self.conversation_manager.get_history(),
|
||||
self.kb_instance,
|
||||
agent_runner.run_agent
|
||||
)
|
||||
|
||||
self.menu_system.display_workflow_completed()
|
||||
|
||||
# Handle report generation
|
||||
if self.reporting_available:
|
||||
await self._handle_report_generation(
|
||||
workflow,
|
||||
workflow_key,
|
||||
target,
|
||||
workflow_start_time,
|
||||
initial_history_length,
|
||||
connected_servers
|
||||
)
|
||||
else:
|
||||
print(f"\n{Fore.YELLOW}Reporting not available.{Style.RESET_ALL}")
|
||||
|
||||
self.menu_system.press_enter_to_continue()
|
||||
|
||||
async def _handle_report_generation(
|
||||
self,
|
||||
workflow: Dict,
|
||||
workflow_key: str,
|
||||
target: str,
|
||||
workflow_start_time: datetime,
|
||||
initial_history_length: int,
|
||||
connected_servers: List
|
||||
) -> None:
|
||||
"""Handle report generation after workflow completion."""
|
||||
if not self.menu_system.ask_generate_report():
|
||||
return
|
||||
|
||||
save_raw_history = self.menu_system.ask_save_raw_history()
|
||||
|
||||
try:
|
||||
from reporting.generators import generate_report_from_workflow
|
||||
|
||||
# Prepare report data
|
||||
workflow_conversation = self.conversation_manager.get_workflow_conversation(initial_history_length)
|
||||
|
||||
report_data = {
|
||||
'workflow_name': workflow['name'],
|
||||
'workflow_key': workflow_key,
|
||||
'target': target,
|
||||
'timestamp': workflow_start_time,
|
||||
'conversation_history': workflow_conversation,
|
||||
'tools_used': MCPManager.get_available_tools(connected_servers)
|
||||
}
|
||||
|
||||
# Generate professional report
|
||||
print(f"\n{Fore.CYAN}Generating report...{Style.RESET_ALL}")
|
||||
report_path = await generate_report_from_workflow(
|
||||
report_data,
|
||||
agent_runner.run_agent,
|
||||
connected_servers,
|
||||
self.kb_instance,
|
||||
save_raw_history
|
||||
)
|
||||
|
||||
self.menu_system.display_report_generated(report_path)
|
||||
|
||||
except Exception as e:
|
||||
self.menu_system.display_report_error(e)
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Main application run method."""
|
||||
self.display_welcome()
|
||||
self.setup_knowledge_base()
|
||||
|
||||
try:
|
||||
# Setup MCP tools
|
||||
mcp_server_instances, connected_servers = await self.setup_mcp_tools()
|
||||
|
||||
# Check if we need to restart (e.g., after configuring new tools)
|
||||
if mcp_server_instances and not connected_servers:
|
||||
return
|
||||
|
||||
# Main application loop
|
||||
while True:
|
||||
self.menu_system.display_main_menu(
|
||||
self.workflow_engine.is_available(),
|
||||
bool(connected_servers)
|
||||
)
|
||||
|
||||
menu_choice = self.menu_system.get_menu_choice()
|
||||
|
||||
if menu_choice == "1":
|
||||
# Interactive mode
|
||||
should_exit = await self.run_interactive_mode(connected_servers)
|
||||
if should_exit:
|
||||
break
|
||||
|
||||
elif menu_choice == "2":
|
||||
# Automated mode
|
||||
await self.run_automated_mode(connected_servers)
|
||||
|
||||
elif menu_choice == "3":
|
||||
# Agent mode
|
||||
await self.run_agent_mode(connected_servers)
|
||||
|
||||
elif menu_choice == "4":
|
||||
# Exit
|
||||
self.menu_system.display_exit_message()
|
||||
break
|
||||
|
||||
else:
|
||||
self.menu_system.display_invalid_choice()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Fore.YELLOW}Program interrupted by user, exiting...{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error during program execution: {e}{Style.RESET_ALL}")
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
# Cleanup MCP servers
|
||||
await self.mcp_manager.cleanup_servers()
|
||||
|
||||
# Close any remaining asyncio transports
|
||||
await self._cleanup_asyncio_resources()
|
||||
|
||||
print(f"{Fore.GREEN}Program ended.{Style.RESET_ALL}")
|
||||
|
||||
async def _cleanup_asyncio_resources(self) -> None:
|
||||
"""Clean up asyncio resources."""
|
||||
try:
|
||||
# Get the event loop
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
# Close any remaining transports
|
||||
for transport in list(getattr(loop, "_transports", {}).values()):
|
||||
if hasattr(transport, "close"):
|
||||
try:
|
||||
transport.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
# Allow a short time for resources to finalize
|
||||
await asyncio.sleep(0.1)
|
||||
except:
|
||||
pass # Ignore any errors in the final cleanup
|
||||
@@ -1,520 +0,0 @@
|
||||
"""PTT Reasoning Module for LLM-based task tree management."""
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Dict, List, Optional, Any, Tuple
|
||||
from colorama import Fore, Style
|
||||
from core.task_tree_manager import TaskTreeManager, TaskNode, NodeStatus
|
||||
|
||||
|
||||
class PTTReasoningModule:
|
||||
"""Handles LLM interactions for PTT management and decision making."""
|
||||
|
||||
def __init__(self, tree_manager: TaskTreeManager):
|
||||
"""
|
||||
Initialize the PTT reasoning module.
|
||||
|
||||
Args:
|
||||
tree_manager: The task tree manager instance
|
||||
"""
|
||||
self.tree_manager = tree_manager
|
||||
|
||||
def get_tree_initialization_prompt(self, goal: str, target: str, constraints: Dict[str, Any], available_tools: List[str] = None) -> str:
|
||||
"""
|
||||
Generate prompt for tree initialization.
|
||||
|
||||
Args:
|
||||
goal: Primary penetration testing objective
|
||||
target: Target system/network
|
||||
constraints: Scope constraints and limitations
|
||||
available_tools: List of available MCP tools
|
||||
|
||||
Returns:
|
||||
Tree initialization prompt
|
||||
"""
|
||||
tool_info = ""
|
||||
if available_tools:
|
||||
tool_info = f"""
|
||||
Available MCP Tools: {', '.join(available_tools)}
|
||||
|
||||
You must work within the capabilities of these connected tools. Think about what each tool can accomplish:
|
||||
- Consider the full capabilities of each available tool
|
||||
- Adapt your approach to work with available tools
|
||||
- Be creative in how you leverage available tools
|
||||
"""
|
||||
else:
|
||||
tool_info = """
|
||||
No MCP tools are currently connected. Design an approach that describes the security testing objectives without tool dependencies.
|
||||
"""
|
||||
|
||||
prompt = f"""You are an autonomous security agent initializing a Pentesting Task Tree (PTT) for a security assessment.
|
||||
|
||||
ASSESSMENT CONTEXT:
|
||||
Goal: {goal}
|
||||
Target: {target}
|
||||
Constraints: {json.dumps(constraints, indent=2)}
|
||||
|
||||
{tool_info}
|
||||
|
||||
TASK:
|
||||
Analyze this goal and determine what structure and initial tasks are needed to accomplish it efficiently.
|
||||
|
||||
DO NOT assume any predefined phases or structure. Instead:
|
||||
1. Analyze what the goal actually requires
|
||||
2. Determine if you need phases/categories or if direct tasks are better
|
||||
3. Create an appropriate initial structure
|
||||
4. Define specific actionable tasks to start with
|
||||
|
||||
Consider:
|
||||
- What does this specific goal require?
|
||||
- What's the minimal viable approach?
|
||||
- How can available tools be leveraged?
|
||||
- What structure makes sense for THIS goal?
|
||||
|
||||
IMPORTANT: When suggesting tool usage, be specific about commands and modules. For example:
|
||||
|
||||
Provide your analysis and initial structure in JSON format:
|
||||
|
||||
{{
|
||||
"analysis": "Your assessment of what this goal requires and approach",
|
||||
"structure": [
|
||||
{{
|
||||
"type": "phase/category/direct",
|
||||
"name": "Name of organizational structure",
|
||||
"description": "What this encompasses",
|
||||
"justification": "Why this structure element is needed for this goal"
|
||||
}}
|
||||
],
|
||||
"initial_tasks": [
|
||||
{{
|
||||
"description": "Specific actionable task",
|
||||
"parent": "Which structure element this belongs to, or 'root' for direct tasks",
|
||||
"tool_suggestion": "Which available tool to use, or 'manual' if no suitable tool",
|
||||
"priority": 1-10,
|
||||
"risk_level": "low/medium/high",
|
||||
"rationale": "Why this task is necessary for the goal"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
BE INTELLIGENT: If the goal is simple, don't create complex multi-phase structures. If it's complex, then structure appropriately. Let the goal drive the structure, not the other way around."""
|
||||
|
||||
return prompt
|
||||
|
||||
def get_tree_update_prompt(self, tool_output: str, command: str, node: TaskNode) -> str:
|
||||
"""
|
||||
Generate prompt for updating the tree based on tool output.
|
||||
|
||||
Args:
|
||||
tool_output: Output from the executed tool
|
||||
command: The command that was executed
|
||||
node: The node being updated
|
||||
|
||||
Returns:
|
||||
Update prompt
|
||||
"""
|
||||
current_tree = self.tree_manager.to_natural_language()
|
||||
|
||||
prompt = f"""You are managing a Pentesting Task Tree (PTT). A task has been executed and you need to update the tree based on the results.
|
||||
|
||||
Current PTT State:
|
||||
{current_tree}
|
||||
|
||||
Executed Task: {node.description}
|
||||
Command: {command}
|
||||
Tool Output:
|
||||
{tool_output[:2000]} # Limit output length
|
||||
|
||||
Based on this output, provide updates in the following JSON format:
|
||||
|
||||
{{
|
||||
"node_updates": {{
|
||||
"status": "completed/failed/vulnerable/not_vulnerable",
|
||||
"findings": "Summary of key findings from the output",
|
||||
"output_summary": "Brief technical summary"
|
||||
}},
|
||||
"new_tasks": [
|
||||
{{
|
||||
"description": "New task based on findings",
|
||||
"parent_phase": "Phase 1/2/3/4",
|
||||
"tool_suggestion": "Suggested tool",
|
||||
"priority": 1-10,
|
||||
"risk_level": "low/medium/high",
|
||||
"rationale": "Why this task is important"
|
||||
}}
|
||||
],
|
||||
"insights": "Any strategic insights or patterns noticed"
|
||||
}}
|
||||
|
||||
Consider:
|
||||
1. What vulnerabilities or opportunities were discovered?
|
||||
2. What follow-up actions are needed based on the findings?
|
||||
3. Should any new attack vectors be explored?
|
||||
4. Are there any security misconfigurations evident?"""
|
||||
|
||||
return prompt
|
||||
|
||||
def get_next_action_prompt(self, available_tools: List[str]) -> str:
|
||||
"""
|
||||
Generate prompt for selecting the next action.
|
||||
|
||||
Args:
|
||||
available_tools: List of available MCP tools
|
||||
|
||||
Returns:
|
||||
Next action selection prompt
|
||||
"""
|
||||
current_tree = self.tree_manager.to_natural_language()
|
||||
candidates = self.tree_manager.get_candidate_tasks()
|
||||
|
||||
# Prepare candidate descriptions
|
||||
candidate_desc = []
|
||||
for i, task in enumerate(candidates[:10]): # Limit to top 10
|
||||
desc = f"{i+1}. {task.description}"
|
||||
if task.priority:
|
||||
desc += f" (Priority: {task.priority})"
|
||||
candidate_desc.append(desc)
|
||||
|
||||
# Generate tool context
|
||||
if available_tools:
|
||||
tool_context = f"""
|
||||
Connected MCP Tools: {', '.join(available_tools)}
|
||||
|
||||
Think about how to leverage these tools for the selected task. Each tool has its own capabilities -
|
||||
be creative and intelligent about how to accomplish penetration testing objectives with available tools.
|
||||
If a tool doesn't directly support a traditional approach, consider alternative methods that achieve the same goal.
|
||||
"""
|
||||
else:
|
||||
tool_context = """
|
||||
No MCP tools are currently connected. Select tasks that can be performed manually or recommend connecting appropriate tools.
|
||||
"""
|
||||
|
||||
prompt = f"""You are managing a Pentesting Task Tree (PTT) and need to select the next action.
|
||||
|
||||
Goal: {self.tree_manager.goal}
|
||||
Target: {self.tree_manager.target}
|
||||
|
||||
Current PTT State:
|
||||
{current_tree}
|
||||
|
||||
{tool_context}
|
||||
|
||||
Candidate Tasks:
|
||||
{chr(10).join(candidate_desc)}
|
||||
|
||||
Statistics:
|
||||
- Total tasks: {len(self.tree_manager.nodes)}
|
||||
- Completed: {sum(1 for n in self.tree_manager.nodes.values() if n.status == NodeStatus.COMPLETED)}
|
||||
- In Progress: {sum(1 for n in self.tree_manager.nodes.values() if n.status == NodeStatus.IN_PROGRESS)}
|
||||
- Pending: {sum(1 for n in self.tree_manager.nodes.values() if n.status == NodeStatus.PENDING)}
|
||||
|
||||
Select the most strategic next action and provide your response in JSON format:
|
||||
|
||||
{{
|
||||
"selected_task_index": 1-based index from candidate list,
|
||||
"rationale": "Why this task is the best next step",
|
||||
"command": "Intelligent request that leverages available tools effectively",
|
||||
"tool": "Which available tool to use, or 'manual' if no suitable tool",
|
||||
"expected_outcome": "What we hope to discover/achieve",
|
||||
"alternative_if_blocked": "Backup task index if this fails"
|
||||
}}
|
||||
|
||||
Consider:
|
||||
1. Logical progression through the penetration testing methodology
|
||||
2. Task dependencies and prerequisites
|
||||
3. Risk vs reward of different approaches
|
||||
4. How to best utilize available tools for maximum effectiveness
|
||||
5. Strategic value of each potential action
|
||||
|
||||
Be intelligent about tool selection - think about what each available tool can accomplish."""
|
||||
|
||||
return prompt
|
||||
|
||||
def get_goal_check_prompt(self) -> str:
|
||||
"""
|
||||
Generate prompt to check if the goal has been achieved.
|
||||
|
||||
Returns:
|
||||
Goal achievement check prompt
|
||||
"""
|
||||
current_tree = self.tree_manager.to_natural_language()
|
||||
goal = self.tree_manager.goal
|
||||
|
||||
# Extract completed tasks and findings for better context
|
||||
completed_tasks_with_findings = []
|
||||
for node in self.tree_manager.nodes.values():
|
||||
if node.status == NodeStatus.COMPLETED and node.findings:
|
||||
completed_tasks_with_findings.append(f"✓ {node.description}: {node.findings}")
|
||||
|
||||
completed_context = "\n".join(completed_tasks_with_findings) if completed_tasks_with_findings else "No completed tasks with findings yet."
|
||||
|
||||
prompt = f"""Analyze the current Pentesting Task Tree (PTT) to determine if the PRIMARY GOAL has been achieved.
|
||||
|
||||
IMPORTANT: Focus ONLY on whether the specific goal stated has been accomplished. Do not suggest additional scope or activities beyond the original goal.
|
||||
|
||||
PRIMARY GOAL: {goal}
|
||||
Target: {self.tree_manager.target}
|
||||
|
||||
COMPLETED TASKS WITH FINDINGS:
|
||||
{completed_context}
|
||||
|
||||
Current PTT State:
|
||||
{current_tree}
|
||||
|
||||
GOAL ACHIEVEMENT CRITERIA:
|
||||
- For information gathering goals, the goal is achieved when that specific information is obtained
|
||||
- For vulnerability assessment goals, the goal is achieved when vulnerabilities are identified and documented
|
||||
- For exploitation goals, the goal is achieved when successful exploitation is demonstrated
|
||||
- For access goals, the goal is achieved when the specified access level is obtained
|
||||
|
||||
Provide your analysis in JSON format:
|
||||
|
||||
{{
|
||||
"goal_achieved": true/false,
|
||||
"confidence": 0-100,
|
||||
"evidence": "Specific evidence that the PRIMARY GOAL has been met (quote actual findings)",
|
||||
"remaining_objectives": "What still needs to be done if goal not achieved (related to the ORIGINAL goal only)",
|
||||
"recommendations": "Next steps ONLY if they relate to the original goal - do not expand scope",
|
||||
"scope_warning": "Flag if any tasks seem to exceed the original goal scope"
|
||||
}}
|
||||
|
||||
Consider:
|
||||
1. Has the SPECIFIC goal been demonstrably achieved?
|
||||
2. Is there sufficient evidence/proof in the completed tasks?
|
||||
3. Are there critical paths unexplored that are NECESSARY for the original goal?
|
||||
4. Would additional testing strengthen the results for the ORIGINAL goal only?
|
||||
|
||||
DO NOT recommend expanding the scope beyond the original goal. If the goal is completed, mark it as achieved regardless of what other security activities could be performed."""
|
||||
|
||||
return prompt
|
||||
|
||||
def parse_tree_initialization_response(self, llm_response: str) -> Dict[str, Any]:
|
||||
"""Parse LLM response for tree initialization."""
|
||||
try:
|
||||
print(f"{Fore.CYAN}Parsing initialization response...{Style.RESET_ALL}")
|
||||
# Extract JSON from response
|
||||
response_json = self._extract_json(llm_response)
|
||||
|
||||
analysis = response_json.get('analysis', 'No analysis provided')
|
||||
structure = response_json.get('structure', [])
|
||||
initial_tasks = response_json.get('initial_tasks', [])
|
||||
|
||||
print(f"{Fore.GREEN}LLM Analysis: {analysis}{Style.RESET_ALL}")
|
||||
print(f"{Fore.GREEN}Successfully parsed {len(structure)} structure elements and {len(initial_tasks)} tasks{Style.RESET_ALL}")
|
||||
|
||||
return {
|
||||
'analysis': analysis,
|
||||
'structure': structure,
|
||||
'initial_tasks': initial_tasks
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Failed to parse initialization response: {e}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Response text (first 500 chars): {llm_response[:500]}{Style.RESET_ALL}")
|
||||
return {
|
||||
'analysis': 'Failed to parse LLM response',
|
||||
'structure': [],
|
||||
'initial_tasks': []
|
||||
}
|
||||
|
||||
def parse_tree_update_response(self, llm_response: str) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
|
||||
"""Parse LLM response for tree updates."""
|
||||
try:
|
||||
response_json = self._extract_json(llm_response)
|
||||
node_updates = response_json.get('node_updates', {})
|
||||
new_tasks = response_json.get('new_tasks', [])
|
||||
return node_updates, new_tasks
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Failed to parse update response: {e}{Style.RESET_ALL}")
|
||||
return {}, []
|
||||
|
||||
def parse_next_action_response(self, llm_response: str, available_tools: List[str] = None) -> Optional[Dict[str, Any]]:
|
||||
"""Parse LLM response for next action selection."""
|
||||
try:
|
||||
response_json = self._extract_json(llm_response)
|
||||
return response_json
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Failed to parse next action response: {e}{Style.RESET_ALL}")
|
||||
return None
|
||||
|
||||
def parse_goal_check_response(self, llm_response: str) -> Dict[str, Any]:
|
||||
"""Parse LLM response for goal achievement check."""
|
||||
try:
|
||||
response_json = self._extract_json(llm_response)
|
||||
return response_json
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Failed to parse goal check response: {e}{Style.RESET_ALL}")
|
||||
return {"goal_achieved": False, "confidence": 0}
|
||||
|
||||
def _extract_json(self, text: str) -> Dict[str, Any]:
|
||||
"""Extract JSON from LLM response text."""
|
||||
if not text:
|
||||
raise ValueError("Empty response text")
|
||||
|
||||
print(f"{Fore.CYAN}Attempting to extract JSON from {len(text)} character response{Style.RESET_ALL}")
|
||||
|
||||
# Try multiple strategies to extract JSON
|
||||
strategies = [
|
||||
self._extract_json_code_block,
|
||||
self._extract_json_braces,
|
||||
self._extract_json_fuzzy,
|
||||
self._create_fallback_json
|
||||
]
|
||||
|
||||
for i, strategy in enumerate(strategies):
|
||||
try:
|
||||
result = strategy(text)
|
||||
if result:
|
||||
print(f"{Fore.GREEN}Successfully extracted JSON using strategy {i+1}{Style.RESET_ALL}")
|
||||
return result
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Strategy {i+1} failed: {e}{Style.RESET_ALL}")
|
||||
continue
|
||||
|
||||
raise ValueError("Could not extract valid JSON from response")
|
||||
|
||||
def _extract_json_code_block(self, text: str) -> Dict[str, Any]:
|
||||
"""Extract JSON from code blocks."""
|
||||
# Look for JSON between ```json and ``` or just ```
|
||||
patterns = [
|
||||
r'```json\s*(\{.*?\})\s*```',
|
||||
r'```\s*(\{.*?\})\s*```'
|
||||
]
|
||||
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
json_str = match.group(1)
|
||||
return json.loads(json_str)
|
||||
|
||||
raise ValueError("No JSON code block found")
|
||||
|
||||
def _extract_json_braces(self, text: str) -> Dict[str, Any]:
|
||||
"""Extract JSON by finding brace boundaries."""
|
||||
# Find the first { and last }
|
||||
json_start = text.find('{')
|
||||
json_end = text.rfind('}')
|
||||
|
||||
if json_start != -1 and json_end != -1 and json_end > json_start:
|
||||
json_str = text[json_start:json_end + 1]
|
||||
return json.loads(json_str)
|
||||
|
||||
raise ValueError("No valid JSON braces found")
|
||||
|
||||
def _extract_json_fuzzy(self, text: str) -> Dict[str, Any]:
|
||||
"""Try to extract JSON with more flexible matching."""
|
||||
# Look for task-like patterns and try to construct JSON
|
||||
if "tasks" in text.lower():
|
||||
# Try to find task descriptions
|
||||
task_patterns = [
|
||||
r'"description":\s*"([^"]+)"',
|
||||
r'"tool_suggestion":\s*"([^"]+)"',
|
||||
r'"priority":\s*(\d+)',
|
||||
r'"risk_level":\s*"([^"]+)"'
|
||||
]
|
||||
|
||||
# This is a simplified approach - could be enhanced
|
||||
# For now, fall through to the next strategy
|
||||
pass
|
||||
|
||||
raise ValueError("Fuzzy JSON extraction failed")
|
||||
|
||||
def _create_fallback_json(self, text: str) -> Dict[str, Any]:
|
||||
"""Create fallback JSON if no valid JSON is found."""
|
||||
print(f"{Fore.YELLOW}Creating fallback JSON structure{Style.RESET_ALL}")
|
||||
|
||||
# Return an empty but valid structure
|
||||
return {
|
||||
"tasks": [],
|
||||
"node_updates": {"status": "completed"},
|
||||
"new_tasks": [],
|
||||
"selected_task_index": 1,
|
||||
"goal_achieved": False,
|
||||
"confidence": 0
|
||||
}
|
||||
|
||||
def verify_tree_update(self, old_tree_state: str, new_tree_state: str) -> bool:
|
||||
"""
|
||||
Verify that tree updates maintain integrity.
|
||||
|
||||
Args:
|
||||
old_tree_state: Tree state before update
|
||||
new_tree_state: Tree state after update
|
||||
|
||||
Returns:
|
||||
True if update is valid
|
||||
"""
|
||||
# For now, basic verification - can be enhanced
|
||||
# Check that only leaf nodes were modified (as per PentestGPT approach)
|
||||
# This is simplified - in practice would need more sophisticated checks
|
||||
|
||||
return True # Placeholder - implement actual verification logic
|
||||
|
||||
def generate_strategic_summary(self) -> str:
|
||||
"""Generate a strategic summary of the current PTT state."""
|
||||
stats = self.tree_manager.get_statistics()
|
||||
|
||||
summary = f"""
|
||||
=== PTT Strategic Summary ===
|
||||
Goal: {self.tree_manager.goal}
|
||||
Target: {self.tree_manager.target}
|
||||
|
||||
Progress Overview:
|
||||
- Total Tasks: {stats['total_nodes']}
|
||||
- Completed: {stats['status_counts'].get('completed', 0)}
|
||||
- In Progress: {stats['status_counts'].get('in_progress', 0)}
|
||||
- Failed: {stats['status_counts'].get('failed', 0)}
|
||||
- Vulnerabilities Found: {stats['status_counts'].get('vulnerable', 0)}
|
||||
|
||||
Current Phase Focus:
|
||||
"""
|
||||
|
||||
# Identify which phase is most active
|
||||
phase_activity = {}
|
||||
for node in self.tree_manager.nodes.values():
|
||||
if node.node_type == "phase":
|
||||
completed_children = sum(
|
||||
1 for child_id in node.children_ids
|
||||
if child_id in self.tree_manager.nodes
|
||||
and self.tree_manager.nodes[child_id].status == NodeStatus.COMPLETED
|
||||
)
|
||||
total_children = len(node.children_ids)
|
||||
phase_activity[node.description] = (completed_children, total_children)
|
||||
|
||||
for phase, (completed, total) in phase_activity.items():
|
||||
if total > 0:
|
||||
progress = (completed / total) * 100
|
||||
summary += f"- {phase}: {completed}/{total} tasks ({progress:.0f}%)\n"
|
||||
|
||||
# Add key findings
|
||||
summary += "\nKey Findings:\n"
|
||||
vuln_count = 0
|
||||
for node in self.tree_manager.nodes.values():
|
||||
if node.status == NodeStatus.VULNERABLE and node.findings:
|
||||
vuln_count += 1
|
||||
summary += f"- {node.description}: {node.findings[:100]}...\n"
|
||||
if vuln_count >= 5: # Limit to top 5
|
||||
break
|
||||
|
||||
return summary
|
||||
|
||||
def validate_and_fix_tool_suggestions(self, tasks: List[Dict[str, Any]], available_tools: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Let the LLM re-evaluate tool suggestions if they don't match available tools."""
|
||||
if not available_tools:
|
||||
return tasks
|
||||
|
||||
# Check if any tasks use unavailable tools
|
||||
needs_fixing = []
|
||||
valid_tasks = []
|
||||
|
||||
for task in tasks:
|
||||
tool_suggestion = task.get('tool_suggestion', '')
|
||||
if tool_suggestion in available_tools or tool_suggestion in ['manual', 'generic']:
|
||||
valid_tasks.append(task)
|
||||
else:
|
||||
needs_fixing.append(task)
|
||||
|
||||
if needs_fixing:
|
||||
print(f"{Fore.YELLOW}Some tasks reference unavailable tools. Letting AI re-evaluate...{Style.RESET_ALL}")
|
||||
# Return all tasks - let the execution phase handle tool mismatches intelligently
|
||||
|
||||
return tasks
|
||||
@@ -1,357 +0,0 @@
|
||||
"""Task Tree Manager for PTT-based autonomous agent mode."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Dict, List, Optional, Any, Tuple
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class NodeStatus(Enum):
|
||||
"""Enumeration of possible node statuses."""
|
||||
PENDING = "pending"
|
||||
IN_PROGRESS = "in_progress"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
BLOCKED = "blocked"
|
||||
VULNERABLE = "vulnerable"
|
||||
NOT_VULNERABLE = "not_vulnerable"
|
||||
|
||||
|
||||
class RiskLevel(Enum):
|
||||
"""Enumeration of risk levels."""
|
||||
LOW = "low"
|
||||
MEDIUM = "medium"
|
||||
HIGH = "high"
|
||||
|
||||
|
||||
class TaskNode:
|
||||
"""Represents a single node in the task tree."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
description: str,
|
||||
parent_id: Optional[str] = None,
|
||||
node_type: str = "task",
|
||||
**kwargs
|
||||
):
|
||||
"""Initialize a task node."""
|
||||
self.id = kwargs.get('id', str(uuid.uuid4()))
|
||||
self.description = description
|
||||
self.status = NodeStatus(kwargs.get('status', NodeStatus.PENDING.value))
|
||||
self.node_type = node_type # task, phase, finding, objective
|
||||
self.parent_id = parent_id
|
||||
self.children_ids: List[str] = kwargs.get('children_ids', [])
|
||||
|
||||
# Task execution details
|
||||
self.tool_used = kwargs.get('tool_used', None)
|
||||
self.command_executed = kwargs.get('command_executed', None)
|
||||
self.output_summary = kwargs.get('output_summary', None)
|
||||
self.findings = kwargs.get('findings', None)
|
||||
|
||||
# Metadata
|
||||
self.priority = kwargs.get('priority', 5) # 1-10, higher is more important
|
||||
self.risk_level = RiskLevel(kwargs.get('risk_level', RiskLevel.LOW.value))
|
||||
self.timestamp = kwargs.get('timestamp', None)
|
||||
self.kb_references = kwargs.get('kb_references', [])
|
||||
self.dependencies = kwargs.get('dependencies', [])
|
||||
|
||||
# Additional attributes
|
||||
self.attributes = kwargs.get('attributes', {})
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert node to dictionary representation."""
|
||||
return {
|
||||
'id': self.id,
|
||||
'description': self.description,
|
||||
'status': self.status.value,
|
||||
'node_type': self.node_type,
|
||||
'parent_id': self.parent_id,
|
||||
'children_ids': self.children_ids,
|
||||
'tool_used': self.tool_used,
|
||||
'command_executed': self.command_executed,
|
||||
'output_summary': self.output_summary,
|
||||
'findings': self.findings,
|
||||
'priority': self.priority,
|
||||
'risk_level': self.risk_level.value,
|
||||
'timestamp': self.timestamp,
|
||||
'kb_references': self.kb_references,
|
||||
'dependencies': self.dependencies,
|
||||
'attributes': self.attributes
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'TaskNode':
|
||||
"""Create node from dictionary representation."""
|
||||
return cls(
|
||||
description=data['description'],
|
||||
**data
|
||||
)
|
||||
|
||||
|
||||
class TaskTreeManager:
|
||||
"""Manages the Pentesting Task Tree (PTT) structure and operations."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the task tree manager."""
|
||||
self.nodes: Dict[str, TaskNode] = {}
|
||||
self.root_id: Optional[str] = None
|
||||
self.goal: Optional[str] = None
|
||||
self.target: Optional[str] = None
|
||||
self.constraints: Dict[str, Any] = {}
|
||||
self.creation_time = datetime.now()
|
||||
|
||||
def initialize_tree(self, goal: str, target: str, constraints: Dict[str, Any] = None) -> str:
|
||||
"""
|
||||
Initialize the task tree with a goal and target.
|
||||
|
||||
Args:
|
||||
goal: The primary objective
|
||||
target: The target system/network
|
||||
constraints: Any constraints or scope limitations
|
||||
|
||||
Returns:
|
||||
The root node ID
|
||||
"""
|
||||
self.goal = goal
|
||||
self.target = target
|
||||
self.constraints = constraints or {}
|
||||
|
||||
# Create root node - let the LLM determine what structure is needed
|
||||
root_node = TaskNode(
|
||||
description=f"Goal: {goal}",
|
||||
node_type="objective"
|
||||
)
|
||||
self.root_id = root_node.id
|
||||
self.nodes[root_node.id] = root_node
|
||||
|
||||
return self.root_id
|
||||
|
||||
def add_node(self, node: TaskNode) -> str:
|
||||
"""
|
||||
Add a node to the tree.
|
||||
|
||||
Args:
|
||||
node: The TaskNode to add
|
||||
|
||||
Returns:
|
||||
The node ID
|
||||
"""
|
||||
self.nodes[node.id] = node
|
||||
|
||||
# Update parent's children list
|
||||
if node.parent_id and node.parent_id in self.nodes:
|
||||
parent = self.nodes[node.parent_id]
|
||||
if node.id not in parent.children_ids:
|
||||
parent.children_ids.append(node.id)
|
||||
|
||||
return node.id
|
||||
|
||||
def update_node(self, node_id: str, updates: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Update a node's attributes.
|
||||
|
||||
Args:
|
||||
node_id: The ID of the node to update
|
||||
updates: Dictionary of attributes to update
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if node_id not in self.nodes:
|
||||
return False
|
||||
|
||||
node = self.nodes[node_id]
|
||||
|
||||
# Update allowed fields
|
||||
allowed_fields = {
|
||||
'status', 'tool_used', 'command_executed', 'output_summary',
|
||||
'findings', 'priority', 'risk_level', 'timestamp', 'kb_references'
|
||||
}
|
||||
|
||||
for field, value in updates.items():
|
||||
if field in allowed_fields:
|
||||
if field == 'status':
|
||||
node.status = NodeStatus(value)
|
||||
elif field == 'risk_level':
|
||||
node.risk_level = RiskLevel(value)
|
||||
else:
|
||||
setattr(node, field, value)
|
||||
elif field == 'attributes':
|
||||
node.attributes.update(value)
|
||||
|
||||
return True
|
||||
|
||||
def get_node(self, node_id: str) -> Optional[TaskNode]:
|
||||
"""Get a node by ID."""
|
||||
return self.nodes.get(node_id)
|
||||
|
||||
def get_children(self, node_id: str) -> List[TaskNode]:
|
||||
"""Get all children of a node."""
|
||||
if node_id not in self.nodes:
|
||||
return []
|
||||
|
||||
parent = self.nodes[node_id]
|
||||
return [self.nodes[child_id] for child_id in parent.children_ids if child_id in self.nodes]
|
||||
|
||||
def get_leaf_nodes(self) -> List[TaskNode]:
|
||||
"""Get all leaf nodes (nodes without children)."""
|
||||
return [node for node in self.nodes.values() if not node.children_ids]
|
||||
|
||||
def get_candidate_tasks(self) -> List[TaskNode]:
|
||||
"""
|
||||
Get candidate tasks for next action.
|
||||
|
||||
Returns tasks that are:
|
||||
- Leaf nodes
|
||||
- Status is PENDING or FAILED
|
||||
- All dependencies are completed
|
||||
"""
|
||||
candidates = []
|
||||
|
||||
for node in self.get_leaf_nodes():
|
||||
if node.status in [NodeStatus.PENDING, NodeStatus.FAILED]:
|
||||
# Check dependencies
|
||||
deps_satisfied = all(
|
||||
self.nodes.get(dep_id, TaskNode("")).status == NodeStatus.COMPLETED
|
||||
for dep_id in node.dependencies
|
||||
)
|
||||
|
||||
if deps_satisfied:
|
||||
candidates.append(node)
|
||||
|
||||
return candidates
|
||||
|
||||
def prioritize_tasks(self, tasks: List[TaskNode]) -> List[TaskNode]:
|
||||
"""
|
||||
Prioritize tasks based on various factors.
|
||||
|
||||
Args:
|
||||
tasks: List of candidate tasks
|
||||
|
||||
Returns:
|
||||
Sorted list of tasks (highest priority first)
|
||||
"""
|
||||
def task_score(task: TaskNode) -> float:
|
||||
# Base score from priority
|
||||
score = task.priority
|
||||
|
||||
# Boost for reconnaissance tasks in early stages
|
||||
if "recon" in task.description.lower() or "scan" in task.description.lower():
|
||||
completed_count = sum(1 for n in self.nodes.values() if n.status == NodeStatus.COMPLETED)
|
||||
if completed_count < 5:
|
||||
score += 3
|
||||
|
||||
# Boost for vulnerability assessment after recon
|
||||
if "vuln" in task.description.lower() and self._has_completed_recon():
|
||||
score += 2
|
||||
|
||||
# Penalty for high-risk tasks early on
|
||||
if task.risk_level == RiskLevel.HIGH:
|
||||
score -= 2
|
||||
|
||||
return score
|
||||
|
||||
return sorted(tasks, key=task_score, reverse=True)
|
||||
|
||||
def _has_completed_recon(self) -> bool:
|
||||
"""Check if basic reconnaissance has been completed."""
|
||||
recon_keywords = ["scan", "recon", "enumerat", "discover"]
|
||||
completed_recon = any(
|
||||
any(keyword in node.description.lower() for keyword in recon_keywords)
|
||||
and node.status == NodeStatus.COMPLETED
|
||||
for node in self.nodes.values()
|
||||
)
|
||||
return completed_recon
|
||||
|
||||
def to_natural_language(self, node_id: Optional[str] = None, indent: int = 0) -> str:
|
||||
"""
|
||||
Convert the tree (or subtree) to natural language representation.
|
||||
|
||||
Args:
|
||||
node_id: Starting node ID (None for root)
|
||||
indent: Indentation level
|
||||
|
||||
Returns:
|
||||
Natural language representation of the tree
|
||||
"""
|
||||
if node_id is None:
|
||||
node_id = self.root_id
|
||||
|
||||
if node_id not in self.nodes:
|
||||
return ""
|
||||
|
||||
node = self.nodes[node_id]
|
||||
indent_str = " " * indent
|
||||
|
||||
# Format node information
|
||||
status_symbol = {
|
||||
NodeStatus.PENDING: "○",
|
||||
NodeStatus.IN_PROGRESS: "◐",
|
||||
NodeStatus.COMPLETED: "●",
|
||||
NodeStatus.FAILED: "✗",
|
||||
NodeStatus.BLOCKED: "□",
|
||||
NodeStatus.VULNERABLE: "⚠",
|
||||
NodeStatus.NOT_VULNERABLE: "✓"
|
||||
}.get(node.status, "?")
|
||||
|
||||
lines = [f"{indent_str}{status_symbol} {node.description}"]
|
||||
|
||||
# Add findings if present
|
||||
if node.findings:
|
||||
lines.append(f"{indent_str} → Findings: {node.findings}")
|
||||
|
||||
# Add tool/command info if present
|
||||
if node.tool_used:
|
||||
lines.append(f"{indent_str} → Tool: {node.tool_used}")
|
||||
|
||||
# Process children
|
||||
for child_id in node.children_ids:
|
||||
lines.append(self.to_natural_language(child_id, indent + 1))
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Serialize the tree to JSON."""
|
||||
data = {
|
||||
'goal': self.goal,
|
||||
'target': self.target,
|
||||
'constraints': self.constraints,
|
||||
'root_id': self.root_id,
|
||||
'creation_time': self.creation_time.isoformat(),
|
||||
'nodes': {node_id: node.to_dict() for node_id, node in self.nodes.items()}
|
||||
}
|
||||
return json.dumps(data, indent=2)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str: str) -> 'TaskTreeManager':
|
||||
"""Deserialize a tree from JSON."""
|
||||
data = json.loads(json_str)
|
||||
|
||||
manager = cls()
|
||||
manager.goal = data['goal']
|
||||
manager.target = data['target']
|
||||
manager.constraints = data['constraints']
|
||||
manager.root_id = data['root_id']
|
||||
manager.creation_time = datetime.fromisoformat(data['creation_time'])
|
||||
|
||||
# Recreate nodes
|
||||
for node_id, node_data in data['nodes'].items():
|
||||
node = TaskNode.from_dict(node_data)
|
||||
manager.nodes[node_id] = node
|
||||
|
||||
return manager
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get tree statistics."""
|
||||
status_counts = {}
|
||||
for node in self.nodes.values():
|
||||
status = node.status.value
|
||||
status_counts[status] = status_counts.get(status, 0) + 1
|
||||
|
||||
return {
|
||||
'total_nodes': len(self.nodes),
|
||||
'status_counts': status_counts,
|
||||
'leaf_nodes': len(self.get_leaf_nodes()),
|
||||
'candidate_tasks': len(self.get_candidate_tasks())
|
||||
}
|
||||
45
docker-compose.yml
Normal file
45
docker-compose.yml
Normal file
@@ -0,0 +1,45 @@
|
||||
services:
|
||||
ghostcrew:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: ghostcrew
|
||||
environment:
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY}
|
||||
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
|
||||
- GHOSTCREW_MODEL=${GHOSTCREW_MODEL}
|
||||
- GHOSTCREW_DEBUG=${GHOSTCREW_DEBUG:-false}
|
||||
volumes:
|
||||
- ./loot:/app/loot
|
||||
networks:
|
||||
- ghostcrew-net
|
||||
stdin_open: true
|
||||
tty: true
|
||||
|
||||
ghostcrew-kali:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.kali
|
||||
container_name: ghostcrew-kali
|
||||
privileged: true # Required for VPN and some tools
|
||||
cap_add:
|
||||
- NET_ADMIN
|
||||
- SYS_ADMIN
|
||||
environment:
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY}
|
||||
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
|
||||
- GHOSTCREW_MODEL=${GHOSTCREW_MODEL}
|
||||
- ENABLE_TOR=${ENABLE_TOR:-false}
|
||||
- INIT_METASPLOIT=${INIT_METASPLOIT:-false}
|
||||
volumes:
|
||||
- ./loot:/app/loot
|
||||
networks:
|
||||
- ghostcrew-net
|
||||
stdin_open: true
|
||||
tty: true
|
||||
profiles:
|
||||
- kali
|
||||
|
||||
networks:
|
||||
ghostcrew-net:
|
||||
driver: bridge
|
||||
50
docker-entrypoint.sh
Normal file
50
docker-entrypoint.sh
Normal file
@@ -0,0 +1,50 @@
|
||||
#!/bin/bash
|
||||
# GhostCrew Docker Entrypoint
|
||||
|
||||
set -e
|
||||
|
||||
# Colors for output
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
NC='\033[0m'
|
||||
|
||||
echo -e "${GREEN}🔧 GhostCrew Container Starting...${NC}"
|
||||
|
||||
# Start VPN if config provided
|
||||
if [ -f "/vpn/config.ovpn" ]; then
|
||||
echo -e "${YELLOW}📡 Starting VPN connection...${NC}"
|
||||
openvpn --config /vpn/config.ovpn --daemon
|
||||
sleep 5
|
||||
|
||||
# Check VPN connection
|
||||
if ip a show tun0 &>/dev/null; then
|
||||
echo -e "${GREEN}✅ VPN connected${NC}"
|
||||
else
|
||||
echo -e "${RED}⚠️ VPN connection may have failed${NC}"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Start Tor if enabled
|
||||
if [ "$ENABLE_TOR" = "true" ]; then
|
||||
echo -e "${YELLOW}🧅 Starting Tor...${NC}"
|
||||
service tor start
|
||||
sleep 3
|
||||
fi
|
||||
|
||||
# Initialize any databases
|
||||
if [ "$INIT_METASPLOIT" = "true" ]; then
|
||||
echo -e "${YELLOW}🗄️ Initializing Metasploit database...${NC}"
|
||||
msfdb init 2>/dev/null || true
|
||||
fi
|
||||
|
||||
# Create output directory with timestamp
|
||||
OUTPUT_DIR="/output/$(date +%Y%m%d_%H%M%S)"
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
export GHOSTCREW_OUTPUT_DIR="$OUTPUT_DIR"
|
||||
|
||||
echo -e "${GREEN}📁 Output directory: $OUTPUT_DIR${NC}"
|
||||
echo -e "${GREEN}🚀 Starting GhostCrew...${NC}"
|
||||
|
||||
# Execute the main command
|
||||
exec "$@"
|
||||
4
ghostcrew/__init__.py
Normal file
4
ghostcrew/__init__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
"""GhostCrew - AI penetration testing."""
|
||||
|
||||
__version__ = "0.2.0"
|
||||
__author__ = "Masic"
|
||||
6
ghostcrew/__main__.py
Normal file
6
ghostcrew/__main__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
"""GhostCrew entry point for `python -m ghostcrew`."""
|
||||
|
||||
from ghostcrew.interface.main import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
15
ghostcrew/agents/__init__.py
Normal file
15
ghostcrew/agents/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
"""Agent system for GhostCrew."""
|
||||
|
||||
from .base_agent import AgentMessage, BaseAgent
|
||||
from .crew import AgentStatus, AgentWorker, CrewOrchestrator, CrewState
|
||||
from .state import AgentState
|
||||
|
||||
__all__ = [
|
||||
"BaseAgent",
|
||||
"AgentMessage",
|
||||
"AgentState",
|
||||
"CrewOrchestrator",
|
||||
"CrewState",
|
||||
"AgentWorker",
|
||||
"AgentStatus",
|
||||
]
|
||||
528
ghostcrew/agents/base_agent.py
Normal file
528
ghostcrew/agents/base_agent.py
Normal file
@@ -0,0 +1,528 @@
|
||||
"""Base agent class for GhostCrew."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, List, Optional
|
||||
|
||||
from ..config.constants import DEFAULT_MAX_ITERATIONS
|
||||
from .state import AgentState, AgentStateManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..llm import LLM
|
||||
from ..runtime import Runtime
|
||||
from ..tools import Tool
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""Represents a tool call from the LLM."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
arguments: dict
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResult:
|
||||
"""Result from a tool execution."""
|
||||
|
||||
tool_call_id: str
|
||||
tool_name: str
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
success: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentMessage:
|
||||
"""A message in the agent conversation."""
|
||||
|
||||
role: str # "user", "assistant", "tool_result", "system"
|
||||
content: str
|
||||
tool_calls: Optional[List[ToolCall]] = None
|
||||
tool_results: Optional[List[ToolResult]] = None
|
||||
metadata: dict = field(default_factory=dict)
|
||||
usage: Optional[dict] = None # Token usage from LLM response
|
||||
|
||||
def to_llm_format(self) -> dict:
|
||||
"""Convert to LLM message format."""
|
||||
import json
|
||||
|
||||
msg = {"role": self.role, "content": self.content}
|
||||
|
||||
if self.tool_calls:
|
||||
msg["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": (
|
||||
json.dumps(tc.arguments)
|
||||
if isinstance(tc.arguments, dict)
|
||||
else tc.arguments
|
||||
),
|
||||
},
|
||||
}
|
||||
for tc in self.tool_calls
|
||||
]
|
||||
|
||||
return msg
|
||||
|
||||
|
||||
class BaseAgent(ABC):
|
||||
"""Base class for all agents."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm: "LLM",
|
||||
tools: List["Tool"],
|
||||
runtime: "Runtime",
|
||||
max_iterations: int = DEFAULT_MAX_ITERATIONS,
|
||||
):
|
||||
"""
|
||||
Initialize the base agent.
|
||||
|
||||
Args:
|
||||
llm: The LLM instance for generating responses
|
||||
tools: List of tools available to the agent
|
||||
runtime: The runtime environment for tool execution
|
||||
max_iterations: Maximum iterations before forcing stop (safety limit)
|
||||
"""
|
||||
self.llm = llm
|
||||
self.tools = tools
|
||||
self.runtime = runtime
|
||||
self.max_iterations = max_iterations
|
||||
self.state_manager = AgentStateManager()
|
||||
self.conversation_history: List[AgentMessage] = []
|
||||
|
||||
@property
|
||||
def state(self) -> AgentState:
|
||||
"""Get current agent state."""
|
||||
return self.state_manager.current_state
|
||||
|
||||
@state.setter
|
||||
def state(self, value: AgentState):
|
||||
"""Set agent state."""
|
||||
self.state_manager.transition_to(value)
|
||||
|
||||
def cleanup_after_cancel(self) -> None:
|
||||
"""
|
||||
Clean up agent state after a cancellation.
|
||||
|
||||
Removes the cancelled request and any pending tool calls from
|
||||
conversation history to prevent stale responses from contaminating
|
||||
the next conversation.
|
||||
"""
|
||||
# Remove incomplete messages from the end of conversation
|
||||
while self.conversation_history:
|
||||
last_msg = self.conversation_history[-1]
|
||||
# Remove assistant message with tool calls (incomplete tool execution)
|
||||
if last_msg.role == "assistant" and last_msg.tool_calls:
|
||||
self.conversation_history.pop()
|
||||
# Remove orphaned tool_result messages
|
||||
elif last_msg.role == "tool":
|
||||
self.conversation_history.pop()
|
||||
# Remove the user message that triggered the cancelled request
|
||||
elif last_msg.role == "user":
|
||||
self.conversation_history.pop()
|
||||
break # Stop after removing the user message
|
||||
else:
|
||||
break
|
||||
|
||||
# Reset state to idle
|
||||
self.state_manager.transition_to(AgentState.IDLE)
|
||||
|
||||
@abstractmethod
|
||||
def get_system_prompt(self) -> str:
|
||||
"""Return the system prompt for this agent."""
|
||||
pass
|
||||
|
||||
async def agent_loop(self, initial_message: str) -> AsyncIterator[AgentMessage]:
|
||||
"""
|
||||
Main agent execution loop.
|
||||
|
||||
Simple control flow:
|
||||
- Tool calls: Execute tools, continue loop
|
||||
- Text response (no tools): Done
|
||||
- Max iterations reached: Force stop with warning
|
||||
|
||||
Args:
|
||||
initial_message: The initial user message to process
|
||||
|
||||
Yields:
|
||||
AgentMessage objects as the agent processes
|
||||
"""
|
||||
self.state_manager.transition_to(AgentState.THINKING)
|
||||
self.conversation_history.append(
|
||||
AgentMessage(role="user", content=initial_message)
|
||||
)
|
||||
|
||||
async for msg in self._run_loop():
|
||||
yield msg
|
||||
|
||||
async def continue_conversation(
|
||||
self, user_message: str
|
||||
) -> AsyncIterator[AgentMessage]:
|
||||
"""
|
||||
Continue the conversation with a new user message.
|
||||
|
||||
Args:
|
||||
user_message: The new user message
|
||||
|
||||
Yields:
|
||||
AgentMessage objects as the agent processes
|
||||
"""
|
||||
self.conversation_history.append(
|
||||
AgentMessage(role="user", content=user_message)
|
||||
)
|
||||
self.state_manager.transition_to(AgentState.THINKING)
|
||||
|
||||
async for msg in self._run_loop():
|
||||
yield msg
|
||||
|
||||
async def _run_loop(self) -> AsyncIterator[AgentMessage]:
|
||||
"""
|
||||
Core agent loop logic - shared by agent_loop and continue_conversation.
|
||||
|
||||
Termination conditions:
|
||||
1. finish tool is called -> clean exit with summary
|
||||
2. max_iterations reached -> forced exit with warning
|
||||
3. error -> exit with error state
|
||||
|
||||
Text responses WITHOUT tool calls are treated as "thinking out loud"
|
||||
and do NOT terminate the loop. This prevents premature stopping.
|
||||
|
||||
Yields:
|
||||
AgentMessage objects as the agent processes
|
||||
"""
|
||||
from ..tools.completion import extract_completion_summary, is_task_complete
|
||||
|
||||
iteration = 0
|
||||
|
||||
while iteration < self.max_iterations:
|
||||
iteration += 1
|
||||
|
||||
response = await self.llm.generate(
|
||||
system_prompt=self.get_system_prompt(),
|
||||
messages=self._format_messages_for_llm(),
|
||||
tools=self.tools,
|
||||
)
|
||||
|
||||
if response.tool_calls:
|
||||
# Build tool calls list FIRST (before execution)
|
||||
tool_calls = [
|
||||
ToolCall(
|
||||
id=tc.id if hasattr(tc, "id") else str(i),
|
||||
name=(
|
||||
tc.function.name
|
||||
if hasattr(tc, "function")
|
||||
else tc.get("name", "")
|
||||
),
|
||||
arguments=self._parse_arguments(tc),
|
||||
)
|
||||
for i, tc in enumerate(response.tool_calls)
|
||||
]
|
||||
|
||||
# Yield early - show tool calls before execution starts
|
||||
early_msg = AgentMessage(
|
||||
role="assistant",
|
||||
content=response.content or "",
|
||||
tool_calls=tool_calls,
|
||||
tool_results=[], # No results yet
|
||||
usage=response.usage,
|
||||
)
|
||||
yield early_msg
|
||||
|
||||
# Now execute tools
|
||||
self.state_manager.transition_to(AgentState.EXECUTING)
|
||||
tool_results = await self._execute_tools(response.tool_calls)
|
||||
|
||||
# Record in history
|
||||
assistant_msg = AgentMessage(
|
||||
role="assistant",
|
||||
content=response.content or "",
|
||||
tool_calls=tool_calls,
|
||||
usage=response.usage,
|
||||
)
|
||||
self.conversation_history.append(assistant_msg)
|
||||
|
||||
tool_result_msg = AgentMessage(
|
||||
role="tool_result", content="", tool_results=tool_results
|
||||
)
|
||||
self.conversation_history.append(tool_result_msg)
|
||||
|
||||
# Check for explicit task_complete signal
|
||||
for result in tool_results:
|
||||
if result.success and result.result and is_task_complete(result.result):
|
||||
summary = extract_completion_summary(result.result)
|
||||
# Yield results with completion summary
|
||||
display_msg = AgentMessage(
|
||||
role="assistant",
|
||||
content=summary,
|
||||
tool_calls=tool_calls,
|
||||
tool_results=tool_results,
|
||||
usage=response.usage,
|
||||
metadata={"task_complete": True},
|
||||
)
|
||||
yield display_msg
|
||||
self.state_manager.transition_to(AgentState.COMPLETE)
|
||||
return
|
||||
|
||||
# Yield results for display update (no completion yet)
|
||||
display_msg = AgentMessage(
|
||||
role="assistant",
|
||||
content=response.content or "",
|
||||
tool_calls=tool_calls,
|
||||
tool_results=tool_results,
|
||||
usage=response.usage,
|
||||
)
|
||||
yield display_msg
|
||||
self.state_manager.transition_to(AgentState.THINKING)
|
||||
else:
|
||||
# Text response WITHOUT tool calls = thinking/intermediate output
|
||||
# Store it but DON'T terminate - wait for task_complete
|
||||
if response.content:
|
||||
thinking_msg = AgentMessage(
|
||||
role="assistant",
|
||||
content=response.content,
|
||||
usage=response.usage,
|
||||
metadata={"intermediate": True},
|
||||
)
|
||||
self.conversation_history.append(thinking_msg)
|
||||
yield thinking_msg
|
||||
# Continue loop - only task_complete or max_iterations stops us
|
||||
|
||||
# Max iterations reached - force stop
|
||||
warning_msg = AgentMessage(
|
||||
role="assistant",
|
||||
content=f"[!] Reached maximum iterations ({self.max_iterations}). Stopping to prevent infinite loop. You can continue the conversation if needed.",
|
||||
metadata={"max_iterations_reached": True},
|
||||
)
|
||||
self.conversation_history.append(warning_msg)
|
||||
yield warning_msg
|
||||
self.state_manager.transition_to(AgentState.COMPLETE)
|
||||
|
||||
def _format_messages_for_llm(self) -> List[dict]:
|
||||
"""Format conversation history for LLM."""
|
||||
messages = []
|
||||
|
||||
for msg in self.conversation_history:
|
||||
if msg.role == "tool_result" and msg.tool_results:
|
||||
# Format tool results as tool response messages
|
||||
for result in msg.tool_results:
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": (
|
||||
result.result
|
||||
if result.success
|
||||
else f"Error: {result.error}"
|
||||
),
|
||||
"tool_call_id": result.tool_call_id,
|
||||
}
|
||||
)
|
||||
else:
|
||||
messages.append(msg.to_llm_format())
|
||||
|
||||
return messages
|
||||
|
||||
def _parse_arguments(self, tool_call: Any) -> dict:
|
||||
"""Parse tool call arguments."""
|
||||
import json
|
||||
|
||||
if hasattr(tool_call, "function"):
|
||||
args = tool_call.function.arguments
|
||||
elif isinstance(tool_call, dict):
|
||||
args = tool_call.get("arguments", {})
|
||||
else:
|
||||
args = {}
|
||||
|
||||
if isinstance(args, str):
|
||||
try:
|
||||
return json.loads(args)
|
||||
except json.JSONDecodeError:
|
||||
return {"raw": args}
|
||||
return args
|
||||
|
||||
async def _execute_tools(self, tool_calls: List[Any]) -> List[ToolResult]:
|
||||
"""
|
||||
Execute tool calls and return results.
|
||||
|
||||
Args:
|
||||
tool_calls: List of tool calls from the LLM
|
||||
|
||||
Returns:
|
||||
List of ToolResult objects
|
||||
"""
|
||||
results = []
|
||||
|
||||
for i, call in enumerate(tool_calls):
|
||||
# Extract tool call id, name and arguments
|
||||
if hasattr(call, "id"):
|
||||
tool_call_id = call.id
|
||||
elif isinstance(call, dict) and "id" in call:
|
||||
tool_call_id = call["id"]
|
||||
else:
|
||||
tool_call_id = f"call_{i}"
|
||||
|
||||
if hasattr(call, "function"):
|
||||
name = call.function.name
|
||||
arguments = self._parse_arguments(call)
|
||||
elif isinstance(call, dict):
|
||||
name = call.get("name", "")
|
||||
arguments = call.get("arguments", {})
|
||||
else:
|
||||
continue
|
||||
|
||||
tool = self._find_tool(name)
|
||||
|
||||
if tool:
|
||||
try:
|
||||
result = await tool.execute(arguments, self.runtime)
|
||||
results.append(
|
||||
ToolResult(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_name=name,
|
||||
result=result,
|
||||
success=True,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
results.append(
|
||||
ToolResult(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_name=name,
|
||||
error=str(e),
|
||||
success=False,
|
||||
)
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
ToolResult(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_name=name,
|
||||
error=f"Tool '{name}' not found",
|
||||
success=False,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
def _find_tool(self, name: str) -> Optional["Tool"]:
|
||||
"""
|
||||
Find a tool by name.
|
||||
|
||||
Args:
|
||||
name: The tool name to find
|
||||
|
||||
Returns:
|
||||
The Tool if found, None otherwise
|
||||
"""
|
||||
for tool in self.tools:
|
||||
if tool.name == name:
|
||||
return tool
|
||||
return None
|
||||
|
||||
def reset(self):
|
||||
"""Reset the agent state for a new conversation."""
|
||||
self.state_manager.reset()
|
||||
self.conversation_history.clear()
|
||||
|
||||
async def assist(self, message: str) -> AsyncIterator[AgentMessage]:
|
||||
"""
|
||||
Assist mode - single LLM call, single tool execution if needed.
|
||||
|
||||
Simple flow: LLM responds, optionally calls one tool, returns result.
|
||||
No looping, no retries. User can follow up if needed.
|
||||
|
||||
Note: 'finish' tool is excluded - assist mode doesn't need explicit
|
||||
termination since it's single-shot by design.
|
||||
|
||||
Args:
|
||||
message: The user message to respond to
|
||||
|
||||
Yields:
|
||||
AgentMessage objects
|
||||
"""
|
||||
self.state_manager.transition_to(AgentState.THINKING)
|
||||
self.conversation_history.append(AgentMessage(role="user", content=message))
|
||||
|
||||
# Filter out 'finish' tool - not needed for single-shot assist mode
|
||||
assist_tools = [t for t in self.tools if t.name != "finish"]
|
||||
|
||||
# Single LLM call with tools available
|
||||
response = await self.llm.generate(
|
||||
system_prompt=self.get_system_prompt(),
|
||||
messages=self._format_messages_for_llm(),
|
||||
tools=assist_tools,
|
||||
)
|
||||
|
||||
# If LLM wants to use tools, execute and return result
|
||||
if response.tool_calls:
|
||||
# Build tool calls list
|
||||
tool_calls = [
|
||||
ToolCall(
|
||||
id=tc.id if hasattr(tc, "id") else str(i),
|
||||
name=(
|
||||
tc.function.name
|
||||
if hasattr(tc, "function")
|
||||
else tc.get("name", "")
|
||||
),
|
||||
arguments=self._parse_arguments(tc),
|
||||
)
|
||||
for i, tc in enumerate(response.tool_calls)
|
||||
]
|
||||
|
||||
# Yield tool calls IMMEDIATELY (before execution) for UI display
|
||||
# Include any thinking/planning content from the LLM
|
||||
thinking_msg = AgentMessage(
|
||||
role="assistant", content=response.content or "", tool_calls=tool_calls
|
||||
)
|
||||
yield thinking_msg
|
||||
|
||||
# NOW execute the tools (this can take a while)
|
||||
self.state_manager.transition_to(AgentState.EXECUTING)
|
||||
tool_results = await self._execute_tools(response.tool_calls)
|
||||
|
||||
# Store in history (minimal content to save tokens)
|
||||
assistant_msg = AgentMessage(
|
||||
role="assistant", content="", tool_calls=tool_calls
|
||||
)
|
||||
self.conversation_history.append(assistant_msg)
|
||||
|
||||
tool_result_msg = AgentMessage(
|
||||
role="tool_result", content="", tool_results=tool_results
|
||||
)
|
||||
self.conversation_history.append(tool_result_msg)
|
||||
|
||||
# Yield tool results for display
|
||||
results_msg = AgentMessage(
|
||||
role="assistant", content="", tool_results=tool_results
|
||||
)
|
||||
yield results_msg
|
||||
|
||||
# Format tool results as final response
|
||||
result_text = self._format_tool_results(tool_results)
|
||||
final_msg = AgentMessage(role="assistant", content=result_text)
|
||||
self.conversation_history.append(final_msg)
|
||||
yield final_msg
|
||||
else:
|
||||
# Direct response, no tools needed
|
||||
assistant_msg = AgentMessage(
|
||||
role="assistant", content=response.content or ""
|
||||
)
|
||||
self.conversation_history.append(assistant_msg)
|
||||
yield assistant_msg
|
||||
|
||||
self.state_manager.transition_to(AgentState.COMPLETE)
|
||||
|
||||
def _format_tool_results(self, results: List[ToolResult]) -> str:
|
||||
"""Format tool results as a simple response."""
|
||||
parts = []
|
||||
for r in results:
|
||||
if r.success:
|
||||
parts.append(r.result or "Done.")
|
||||
else:
|
||||
parts.append(f"Error: {r.error}")
|
||||
return "\n".join(parts)
|
||||
17
ghostcrew/agents/crew/__init__.py
Normal file
17
ghostcrew/agents/crew/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""Crew orchestration module."""
|
||||
|
||||
from .models import AgentStatus, AgentWorker, CrewState, Finding, WorkerCallback
|
||||
from .orchestrator import CrewOrchestrator
|
||||
from .tools import create_crew_tools
|
||||
from .worker_pool import WorkerPool
|
||||
|
||||
__all__ = [
|
||||
"CrewOrchestrator",
|
||||
"CrewState",
|
||||
"AgentStatus",
|
||||
"AgentWorker",
|
||||
"Finding",
|
||||
"WorkerCallback",
|
||||
"WorkerPool",
|
||||
"create_crew_tools",
|
||||
]
|
||||
78
ghostcrew/agents/crew/models.py
Normal file
78
ghostcrew/agents/crew/models.py
Normal file
@@ -0,0 +1,78 @@
|
||||
"""Data models for crew orchestration."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Literal, Optional
|
||||
|
||||
|
||||
class CrewState(Enum):
|
||||
"""State of the crew orchestrator."""
|
||||
|
||||
IDLE = "idle"
|
||||
RUNNING = "running"
|
||||
COMPLETE = "complete"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
class AgentStatus(Enum):
|
||||
"""Status of a worker agent."""
|
||||
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETE = "complete"
|
||||
ERROR = "error"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentWorker:
|
||||
"""A worker agent managed by the crew."""
|
||||
|
||||
id: str
|
||||
task: str
|
||||
status: AgentStatus = AgentStatus.PENDING
|
||||
priority: int = 1
|
||||
depends_on: List[str] = field(default_factory=list)
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
tools_used: List[str] = field(default_factory=list)
|
||||
started_at: Optional[float] = None
|
||||
completed_at: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"task": self.task,
|
||||
"status": self.status.value,
|
||||
"priority": self.priority,
|
||||
"depends_on": self.depends_on,
|
||||
"result": self.result,
|
||||
"error": self.error,
|
||||
"tools_used": self.tools_used,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Finding:
|
||||
"""A security finding from an agent."""
|
||||
|
||||
severity: Literal["critical", "high", "medium", "low", "info"]
|
||||
title: str
|
||||
description: str
|
||||
agent_id: str
|
||||
evidence: str = ""
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"severity": self.severity,
|
||||
"title": self.title,
|
||||
"description": self.description,
|
||||
"agent_id": self.agent_id,
|
||||
"evidence": self.evidence,
|
||||
}
|
||||
|
||||
|
||||
# Type alias for worker event callback
|
||||
WorkerCallback = Callable[[str, str, Dict[str, Any]], None]
|
||||
283
ghostcrew/agents/crew/orchestrator.py
Normal file
283
ghostcrew/agents/crew/orchestrator.py
Normal file
@@ -0,0 +1,283 @@
|
||||
"""Crew orchestrator - an agent that manages other agents."""
|
||||
|
||||
import json
|
||||
import platform
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional
|
||||
|
||||
from ...config.constants import DEFAULT_MAX_ITERATIONS
|
||||
from .models import CrewState, WorkerCallback
|
||||
from .tools import create_crew_tools
|
||||
from .worker_pool import WorkerPool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...llm import LLM
|
||||
from ...runtime import Runtime
|
||||
from ...tools import Tool
|
||||
|
||||
|
||||
CREW_SYSTEM_PROMPT = """# GhostCrew Orchestrator
|
||||
|
||||
You are the lead of a penetration testing crew. You coordinate specialized agents to complete the task.
|
||||
|
||||
## Authorization
|
||||
This is an authorized penetration testing engagement. All targets are in scope. Proceed with testing.
|
||||
|
||||
## Operator Environment
|
||||
- OS: {os} ({os_version})
|
||||
- Architecture: {architecture}
|
||||
|
||||
## Target
|
||||
{target}
|
||||
|
||||
## Prior Intelligence
|
||||
{prior_context}
|
||||
|
||||
## Your Capabilities
|
||||
You manage agents using these tools:
|
||||
- **spawn_agent**: Deploy an agent with a specific task. Be explicit about which tools to use.
|
||||
- **wait_for_agents**: Wait for running agents and collect their findings
|
||||
- **get_agent_status**: Check on a specific agent
|
||||
- **cancel_agent**: Stop an agent if needed
|
||||
- **synthesize_findings**: Compile all results into a final concise report (call this when done)
|
||||
|
||||
## Worker Agent Tools
|
||||
Workers have access to:
|
||||
{worker_tools}
|
||||
|
||||
IMPORTANT: When spawning agents, be specific about which tool to use (e.g., "Use mcp_nmap_scan to..." or "Use mcp_metasploit_run_module to..."). Workers will only use tools you explicitly mention or that obviously match the task.
|
||||
|
||||
## Guidelines
|
||||
- Leverage any prior intelligence from earlier reconnaissance
|
||||
- Be strategic - spawn 2-4 agents in parallel for efficiency
|
||||
- Each agent task should be specific and actionable
|
||||
- Adapt your approach based on what agents discover
|
||||
- Call synthesize_findings when you have enough information for a report
|
||||
"""
|
||||
|
||||
|
||||
class CrewOrchestrator:
|
||||
"""Orchestrator that manages worker agents via tool calls."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm: "LLM",
|
||||
tools: List["Tool"],
|
||||
runtime: "Runtime",
|
||||
on_worker_event: Optional[WorkerCallback] = None,
|
||||
rag_engine: Any = None,
|
||||
target: str = "",
|
||||
prior_context: str = "",
|
||||
):
|
||||
self.llm = llm
|
||||
self.base_tools = tools
|
||||
self.runtime = runtime
|
||||
self.on_worker_event = on_worker_event
|
||||
self.rag_engine = rag_engine
|
||||
self.target = target
|
||||
self.prior_context = prior_context
|
||||
|
||||
self.state = CrewState.IDLE
|
||||
self.pool: Optional[WorkerPool] = None
|
||||
self._messages: List[Dict[str, Any]] = []
|
||||
|
||||
def _get_system_prompt(self) -> str:
|
||||
"""Build the system prompt with target info and context."""
|
||||
tool_lines = []
|
||||
for t in self.base_tools:
|
||||
desc = (
|
||||
t.description[:80] + "..." if len(t.description) > 80 else t.description
|
||||
)
|
||||
tool_lines.append(f"- **{t.name}**: {desc}")
|
||||
worker_tools_formatted = (
|
||||
"\n".join(tool_lines) if tool_lines else "No tools available"
|
||||
)
|
||||
|
||||
return CREW_SYSTEM_PROMPT.format(
|
||||
target=self.target or "Not specified",
|
||||
prior_context=self.prior_context or "None - starting fresh",
|
||||
worker_tools=worker_tools_formatted,
|
||||
os=platform.system(),
|
||||
os_version=platform.release(),
|
||||
architecture=platform.machine(),
|
||||
)
|
||||
|
||||
async def run(self, task: str) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Run the crew on a task."""
|
||||
self.state = CrewState.RUNNING
|
||||
yield {"phase": "starting"}
|
||||
|
||||
self.pool = WorkerPool(
|
||||
llm=self.llm,
|
||||
tools=self.base_tools,
|
||||
runtime=self.runtime,
|
||||
target=self.target,
|
||||
rag_engine=self.rag_engine,
|
||||
on_worker_event=self.on_worker_event,
|
||||
)
|
||||
|
||||
crew_tools = create_crew_tools(self.pool, self.llm)
|
||||
|
||||
self._messages = [
|
||||
{"role": "user", "content": f"Target: {self.target}\n\nTask: {task}"}
|
||||
]
|
||||
|
||||
iteration = 0
|
||||
final_report = ""
|
||||
|
||||
try:
|
||||
while iteration < DEFAULT_MAX_ITERATIONS:
|
||||
iteration += 1
|
||||
|
||||
response = await self.llm.generate(
|
||||
system_prompt=self._get_system_prompt(),
|
||||
messages=self._messages,
|
||||
tools=crew_tools,
|
||||
)
|
||||
|
||||
if response.content:
|
||||
yield {"phase": "thinking", "content": response.content}
|
||||
self._messages.append(
|
||||
{"role": "assistant", "content": response.content}
|
||||
)
|
||||
|
||||
if response.tool_calls:
|
||||
def get_tc_name(tc):
|
||||
if hasattr(tc, "function"):
|
||||
return tc.function.name
|
||||
return (
|
||||
tc.get("function", {}).get("name", "")
|
||||
if isinstance(tc, dict)
|
||||
else ""
|
||||
)
|
||||
|
||||
def get_tc_args(tc):
|
||||
if hasattr(tc, "function"):
|
||||
args = tc.function.arguments
|
||||
else:
|
||||
args = (
|
||||
tc.get("function", {}).get("arguments", "{}")
|
||||
if isinstance(tc, dict)
|
||||
else "{}"
|
||||
)
|
||||
if isinstance(args, str):
|
||||
try:
|
||||
return json.loads(args)
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
return args if isinstance(args, dict) else {}
|
||||
|
||||
def get_tc_id(tc):
|
||||
if hasattr(tc, "id"):
|
||||
return tc.id
|
||||
return tc.get("id", "") if isinstance(tc, dict) else ""
|
||||
|
||||
self._messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": response.content or "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": get_tc_id(tc),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": get_tc_name(tc),
|
||||
"arguments": json.dumps(get_tc_args(tc)),
|
||||
},
|
||||
}
|
||||
for tc in response.tool_calls
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
for tc in response.tool_calls:
|
||||
tc_name = get_tc_name(tc)
|
||||
tc_args = get_tc_args(tc)
|
||||
tc_id = get_tc_id(tc)
|
||||
|
||||
yield {"phase": "tool_call", "tool": tc_name, "args": tc_args}
|
||||
|
||||
tool = next((t for t in crew_tools if t.name == tc_name), None)
|
||||
if tool:
|
||||
try:
|
||||
result = await tool.execute(tc_args, self.runtime)
|
||||
|
||||
yield {
|
||||
"phase": "tool_result",
|
||||
"tool": tc_name,
|
||||
"result": result,
|
||||
}
|
||||
|
||||
self._messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
"content": str(result),
|
||||
}
|
||||
)
|
||||
|
||||
if tc_name == "synthesize_findings":
|
||||
final_report = result
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error: {e}"
|
||||
yield {
|
||||
"phase": "tool_result",
|
||||
"tool": tc_name,
|
||||
"result": error_msg,
|
||||
}
|
||||
self._messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
"content": error_msg,
|
||||
}
|
||||
)
|
||||
else:
|
||||
error_msg = f"Unknown tool: {tc_name}"
|
||||
self._messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
"content": error_msg,
|
||||
}
|
||||
)
|
||||
|
||||
if final_report:
|
||||
break
|
||||
else:
|
||||
content = response.content or ""
|
||||
if content:
|
||||
final_report = content
|
||||
break
|
||||
|
||||
self.state = CrewState.COMPLETE
|
||||
yield {"phase": "complete", "report": final_report}
|
||||
|
||||
except Exception as e:
|
||||
self.state = CrewState.ERROR
|
||||
yield {"phase": "error", "error": str(e)}
|
||||
|
||||
finally:
|
||||
if self.pool:
|
||||
await self.pool.cancel_all()
|
||||
|
||||
async def cancel(self) -> None:
|
||||
"""Cancel the crew run."""
|
||||
if self.pool:
|
||||
await self.pool.cancel_all()
|
||||
self._cleanup_pending_calls()
|
||||
self.state = CrewState.IDLE
|
||||
|
||||
def _cleanup_pending_calls(self) -> None:
|
||||
"""Remove incomplete tool calls from message history."""
|
||||
while self._messages:
|
||||
last_msg = self._messages[-1]
|
||||
if last_msg.get("role") == "assistant" and last_msg.get("tool_calls"):
|
||||
self._messages.pop()
|
||||
elif last_msg.get("role") == "tool":
|
||||
self._messages.pop()
|
||||
elif last_msg.get("role") == "user":
|
||||
self._messages.pop()
|
||||
break
|
||||
else:
|
||||
break
|
||||
197
ghostcrew/agents/crew/tools.py
Normal file
197
ghostcrew/agents/crew/tools.py
Normal file
@@ -0,0 +1,197 @@
|
||||
"""Orchestration tools for the crew agent."""
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, List
|
||||
|
||||
from ...tools.registry import Tool, ToolSchema
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...llm import LLM
|
||||
from ...runtime import Runtime
|
||||
from .worker_pool import WorkerPool
|
||||
|
||||
|
||||
def create_crew_tools(pool: "WorkerPool", llm: "LLM") -> List[Tool]:
|
||||
"""Create orchestration tools bound to a worker pool."""
|
||||
|
||||
async def spawn_agent_fn(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""Spawn a new agent to work on a task."""
|
||||
task = arguments.get("task", "")
|
||||
priority = arguments.get("priority", 1)
|
||||
depends_on = arguments.get("depends_on", [])
|
||||
|
||||
if not task:
|
||||
return "Error: task is required"
|
||||
|
||||
agent_id = await pool.spawn(task, priority, depends_on)
|
||||
return f"Spawned {agent_id}: {task}"
|
||||
|
||||
async def wait_for_agents_fn(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""Wait for agents to complete and get their results."""
|
||||
agent_ids = arguments.get("agent_ids", None)
|
||||
|
||||
results = await pool.wait_for(agent_ids)
|
||||
|
||||
if not results:
|
||||
return "No agents to wait for."
|
||||
|
||||
output = []
|
||||
for agent_id, data in results.items():
|
||||
status = data.get("status", "unknown")
|
||||
task = data.get("task", "")
|
||||
result = data.get("result", "")
|
||||
error = data.get("error", "")
|
||||
tools = data.get("tools_used", [])
|
||||
|
||||
output.append(f"## {agent_id}: {task}")
|
||||
output.append(f"Status: {status}")
|
||||
if tools:
|
||||
output.append(f"Tools used: {', '.join(tools)}")
|
||||
if result:
|
||||
output.append(f"Result:\n{result}")
|
||||
if error:
|
||||
output.append(f"Error: {error}")
|
||||
output.append("")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
async def get_agent_status_fn(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""Check the current status of an agent."""
|
||||
agent_id = arguments.get("agent_id", "")
|
||||
|
||||
if not agent_id:
|
||||
return "Error: agent_id is required"
|
||||
|
||||
status = pool.get_status(agent_id)
|
||||
if not status:
|
||||
return f"Agent {agent_id} not found."
|
||||
|
||||
return json.dumps(status, indent=2)
|
||||
|
||||
async def cancel_agent_fn(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""Cancel a running agent."""
|
||||
agent_id = arguments.get("agent_id", "")
|
||||
|
||||
if not agent_id:
|
||||
return "Error: agent_id is required"
|
||||
|
||||
success = await pool.cancel(agent_id)
|
||||
if success:
|
||||
return f"Cancelled {agent_id}"
|
||||
return f"Could not cancel {agent_id} (not running or not found)"
|
||||
|
||||
async def synthesize_findings_fn(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""Compile all agent results into a unified report."""
|
||||
workers = pool.get_workers()
|
||||
if not workers:
|
||||
return "No agents have run yet."
|
||||
|
||||
results_text = []
|
||||
for w in workers:
|
||||
if w.result:
|
||||
results_text.append(f"## {w.task}\n{w.result}")
|
||||
elif w.error:
|
||||
results_text.append(f"## {w.task}\nError: {w.error}")
|
||||
|
||||
if not results_text:
|
||||
return "No results to synthesize."
|
||||
|
||||
prompt = f"""Synthesize these agent findings into a unified penetration test report.
|
||||
Present concrete findings. Be factual and concise about what was discovered.
|
||||
|
||||
{chr(10).join(results_text)}"""
|
||||
|
||||
response = await llm.generate(
|
||||
system_prompt="Synthesize penetration test findings into a clear, actionable report.",
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
tools=[],
|
||||
)
|
||||
|
||||
return response.content
|
||||
|
||||
# Create Tool objects
|
||||
tools = [
|
||||
Tool(
|
||||
name="spawn_agent",
|
||||
description="Spawn a new agent to work on a specific task. Use for delegating work like port scanning, service enumeration, or vulnerability testing. Each agent runs independently with access to all pentest tools.",
|
||||
schema=ToolSchema(
|
||||
type="object",
|
||||
properties={
|
||||
"task": {
|
||||
"type": "string",
|
||||
"description": "Clear, action-oriented task description. Be specific about what to scan/test and the target.",
|
||||
},
|
||||
"priority": {
|
||||
"type": "integer",
|
||||
"description": "Execution priority (higher = runs sooner). Default 1.",
|
||||
},
|
||||
"depends_on": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Agent IDs that must complete before this agent starts. Use for sequential workflows.",
|
||||
},
|
||||
},
|
||||
required=["task"],
|
||||
),
|
||||
execute_fn=spawn_agent_fn,
|
||||
category="orchestration",
|
||||
),
|
||||
Tool(
|
||||
name="wait_for_agents",
|
||||
description="Wait for spawned agents to complete and retrieve their results. Call this after spawning agents to get findings before proceeding.",
|
||||
schema=ToolSchema(
|
||||
type="object",
|
||||
properties={
|
||||
"agent_ids": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "List of agent IDs to wait for. Omit to wait for all spawned agents.",
|
||||
}
|
||||
},
|
||||
required=[],
|
||||
),
|
||||
execute_fn=wait_for_agents_fn,
|
||||
category="orchestration",
|
||||
),
|
||||
Tool(
|
||||
name="get_agent_status",
|
||||
description="Check the current status of a specific agent. Useful for monitoring long-running tasks.",
|
||||
schema=ToolSchema(
|
||||
type="object",
|
||||
properties={
|
||||
"agent_id": {
|
||||
"type": "string",
|
||||
"description": "The agent ID to check (e.g., 'agent-0')",
|
||||
}
|
||||
},
|
||||
required=["agent_id"],
|
||||
),
|
||||
execute_fn=get_agent_status_fn,
|
||||
category="orchestration",
|
||||
),
|
||||
Tool(
|
||||
name="cancel_agent",
|
||||
description="Cancel a running agent. Use if an agent is taking too long or is no longer needed.",
|
||||
schema=ToolSchema(
|
||||
type="object",
|
||||
properties={
|
||||
"agent_id": {
|
||||
"type": "string",
|
||||
"description": "The agent ID to cancel (e.g., 'agent-0')",
|
||||
}
|
||||
},
|
||||
required=["agent_id"],
|
||||
),
|
||||
execute_fn=cancel_agent_fn,
|
||||
category="orchestration",
|
||||
),
|
||||
Tool(
|
||||
name="synthesize_findings",
|
||||
description="Compile all agent results into a unified penetration test report. Call this after all agents have completed.",
|
||||
schema=ToolSchema(type="object", properties={}, required=[]),
|
||||
execute_fn=synthesize_findings_fn,
|
||||
category="orchestration",
|
||||
),
|
||||
]
|
||||
|
||||
return tools
|
||||
250
ghostcrew/agents/crew/worker_pool.py
Normal file
250
ghostcrew/agents/crew/worker_pool.py
Normal file
@@ -0,0 +1,250 @@
|
||||
"""Worker pool for managing concurrent agent execution."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from .models import AgentStatus, AgentWorker, WorkerCallback
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...llm import LLM
|
||||
from ...runtime import Runtime
|
||||
from ...tools import Tool
|
||||
|
||||
|
||||
class WorkerPool:
|
||||
"""Manages concurrent execution of worker agents."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm: "LLM",
|
||||
tools: List["Tool"],
|
||||
runtime: "Runtime",
|
||||
target: str = "",
|
||||
rag_engine: Any = None,
|
||||
on_worker_event: Optional[WorkerCallback] = None,
|
||||
):
|
||||
self.llm = llm
|
||||
self.tools = tools
|
||||
self.runtime = runtime
|
||||
self.target = target
|
||||
self.rag_engine = rag_engine
|
||||
self.on_worker_event = on_worker_event
|
||||
|
||||
self._workers: Dict[str, AgentWorker] = {}
|
||||
self._tasks: Dict[str, asyncio.Task] = {}
|
||||
self._results: Dict[str, str] = {}
|
||||
self._next_id = 0
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _emit(self, worker_id: str, event: str, data: Dict[str, Any]) -> None:
|
||||
"""Emit event to callback if registered."""
|
||||
if self.on_worker_event:
|
||||
self.on_worker_event(worker_id, event, data)
|
||||
|
||||
def _generate_id(self) -> str:
|
||||
"""Generate unique worker ID."""
|
||||
worker_id = f"agent-{self._next_id}"
|
||||
self._next_id += 1
|
||||
return worker_id
|
||||
|
||||
async def spawn(
|
||||
self,
|
||||
task: str,
|
||||
priority: int = 1,
|
||||
depends_on: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Spawn a new worker agent.
|
||||
|
||||
Args:
|
||||
task: The task description for the agent
|
||||
priority: Higher priority runs first (for future use)
|
||||
depends_on: List of agent IDs that must complete first
|
||||
|
||||
Returns:
|
||||
The worker ID
|
||||
"""
|
||||
async with self._lock:
|
||||
worker_id = self._generate_id()
|
||||
|
||||
worker = AgentWorker(
|
||||
id=worker_id,
|
||||
task=task,
|
||||
priority=priority,
|
||||
depends_on=depends_on or [],
|
||||
)
|
||||
self._workers[worker_id] = worker
|
||||
|
||||
# Emit spawn event for UI
|
||||
self._emit(
|
||||
worker_id,
|
||||
"spawn",
|
||||
{
|
||||
"worker_type": worker_id,
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
# Start the agent task
|
||||
self._tasks[worker_id] = asyncio.create_task(self._run_worker(worker))
|
||||
|
||||
return worker_id
|
||||
|
||||
async def _run_worker(self, worker: AgentWorker) -> None:
|
||||
"""Run a single worker agent."""
|
||||
from ..ghostcrew_agent import GhostCrewAgent
|
||||
|
||||
# Wait for dependencies
|
||||
if worker.depends_on:
|
||||
await self._wait_for_dependencies(worker.depends_on)
|
||||
|
||||
worker.status = AgentStatus.RUNNING
|
||||
worker.started_at = time.time()
|
||||
self._emit(worker.id, "status", {"status": "running"})
|
||||
|
||||
agent = GhostCrewAgent(
|
||||
llm=self.llm,
|
||||
tools=self.tools,
|
||||
runtime=self.runtime,
|
||||
target=self.target,
|
||||
rag_engine=self.rag_engine,
|
||||
)
|
||||
|
||||
try:
|
||||
final_response = ""
|
||||
async for response in agent.agent_loop(worker.task):
|
||||
# Track tool calls
|
||||
if response.tool_calls:
|
||||
for tc in response.tool_calls:
|
||||
if tc.name not in worker.tools_used:
|
||||
worker.tools_used.append(tc.name)
|
||||
self._emit(worker.id, "tool", {"tool": tc.name})
|
||||
|
||||
# Track tokens
|
||||
if response.usage:
|
||||
total = response.usage.get("total_tokens", 0)
|
||||
if total > 0:
|
||||
self._emit(worker.id, "tokens", {"tokens": total})
|
||||
|
||||
# Capture final response (text without tool calls)
|
||||
if response.content and not response.tool_calls:
|
||||
final_response = response.content
|
||||
|
||||
worker.result = final_response or "No findings."
|
||||
worker.status = AgentStatus.COMPLETE
|
||||
worker.completed_at = time.time()
|
||||
self._results[worker.id] = worker.result
|
||||
|
||||
self._emit(
|
||||
worker.id,
|
||||
"complete",
|
||||
{
|
||||
"summary": worker.result[:200],
|
||||
},
|
||||
)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
worker.status = AgentStatus.CANCELLED
|
||||
worker.completed_at = time.time()
|
||||
self._emit(worker.id, "cancelled", {})
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
worker.error = str(e)
|
||||
worker.status = AgentStatus.ERROR
|
||||
worker.completed_at = time.time()
|
||||
self._emit(worker.id, "error", {"error": str(e)})
|
||||
|
||||
async def _wait_for_dependencies(self, depends_on: List[str]) -> None:
|
||||
"""Wait for dependent workers to complete."""
|
||||
for dep_id in depends_on:
|
||||
if dep_id in self._tasks:
|
||||
try:
|
||||
await self._tasks[dep_id]
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass # Dependency failed, but we continue
|
||||
|
||||
async def wait_for(self, agent_ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Wait for specified agents (or all) to complete.
|
||||
|
||||
Args:
|
||||
agent_ids: List of agent IDs to wait for. None = wait for all.
|
||||
|
||||
Returns:
|
||||
Dict mapping agent_id to result/error
|
||||
"""
|
||||
if agent_ids is None:
|
||||
agent_ids = list(self._tasks.keys())
|
||||
|
||||
results = {}
|
||||
for agent_id in agent_ids:
|
||||
if agent_id in self._tasks:
|
||||
try:
|
||||
await self._tasks[agent_id]
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
worker = self._workers.get(agent_id)
|
||||
if worker:
|
||||
results[agent_id] = {
|
||||
"task": worker.task,
|
||||
"status": worker.status.value,
|
||||
"result": worker.result,
|
||||
"error": worker.error,
|
||||
"tools_used": worker.tools_used,
|
||||
}
|
||||
|
||||
return results
|
||||
|
||||
def get_status(self, agent_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get status of a specific agent."""
|
||||
worker = self._workers.get(agent_id)
|
||||
if not worker:
|
||||
return None
|
||||
return worker.to_dict()
|
||||
|
||||
def get_all_status(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get status of all agents."""
|
||||
return {wid: w.to_dict() for wid, w in self._workers.items()}
|
||||
|
||||
async def cancel(self, agent_id: str) -> bool:
|
||||
"""Cancel a running agent."""
|
||||
if agent_id not in self._tasks:
|
||||
return False
|
||||
|
||||
task = self._tasks[agent_id]
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
return True
|
||||
return False
|
||||
|
||||
async def cancel_all(self) -> None:
|
||||
"""Cancel all running agents."""
|
||||
for task in self._tasks.values():
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
|
||||
# Wait for all to finish
|
||||
if self._tasks:
|
||||
await asyncio.gather(*self._tasks.values(), return_exceptions=True)
|
||||
|
||||
def get_results(self) -> Dict[str, str]:
|
||||
"""Get results from all completed agents."""
|
||||
return dict(self._results)
|
||||
|
||||
def get_workers(self) -> List[AgentWorker]:
|
||||
"""Get all workers."""
|
||||
return list(self._workers.values())
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the pool for a new task."""
|
||||
self._workers.clear()
|
||||
self._tasks.clear()
|
||||
self._results.clear()
|
||||
self._next_id = 0
|
||||
5
ghostcrew/agents/ghostcrew_agent/__init__.py
Normal file
5
ghostcrew/agents/ghostcrew_agent/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""GhostCrew main agent implementation."""
|
||||
|
||||
from .ghostcrew_agent import GhostCrewAgent
|
||||
|
||||
__all__ = ["GhostCrewAgent"]
|
||||
101
ghostcrew/agents/ghostcrew_agent/ghostcrew_agent.py
Normal file
101
ghostcrew/agents/ghostcrew_agent/ghostcrew_agent.py
Normal file
@@ -0,0 +1,101 @@
|
||||
"""GhostCrew main pentesting agent."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
from jinja2 import Template
|
||||
|
||||
from ..base_agent import BaseAgent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...knowledge import RAGEngine
|
||||
from ...llm import LLM
|
||||
from ...runtime import Runtime
|
||||
from ...tools import Tool
|
||||
|
||||
|
||||
class GhostCrewAgent(BaseAgent):
|
||||
"""Main pentesting agent for GhostCrew."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm: "LLM",
|
||||
tools: List["Tool"],
|
||||
runtime: "Runtime",
|
||||
target: Optional[str] = None,
|
||||
scope: Optional[List[str]] = None,
|
||||
rag_engine: Optional["RAGEngine"] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the GhostCrew agent.
|
||||
|
||||
Args:
|
||||
llm: The LLM instance for generating responses
|
||||
tools: List of tools available to the agent
|
||||
runtime: The runtime environment for tool execution
|
||||
target: Primary target for penetration testing
|
||||
scope: List of in-scope targets/networks
|
||||
rag_engine: RAG engine for knowledge retrieval
|
||||
**kwargs: Additional arguments passed to BaseAgent
|
||||
"""
|
||||
super().__init__(llm, tools, runtime, **kwargs)
|
||||
self.target = target
|
||||
self.scope = scope or []
|
||||
self.rag_engine = rag_engine
|
||||
self._system_prompt_template = self._load_prompt_template()
|
||||
|
||||
def _load_prompt_template(self) -> Template:
|
||||
"""Load the Jinja2 system prompt template."""
|
||||
template_path = Path(__file__).parent / "system_prompt.jinja"
|
||||
return Template(template_path.read_text(encoding="utf-8"))
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
"""Generate system prompt with context."""
|
||||
# Get RAG context if available
|
||||
rag_context = ""
|
||||
if self.rag_engine and self.conversation_history:
|
||||
last_msg = self.conversation_history[-1].content
|
||||
# Ensure content is a string (could be list for multimodal)
|
||||
if isinstance(last_msg, list):
|
||||
last_msg = " ".join(
|
||||
str(part.get("text", ""))
|
||||
for part in last_msg
|
||||
if isinstance(part, dict)
|
||||
)
|
||||
if last_msg:
|
||||
relevant = self.rag_engine.search(last_msg)
|
||||
if relevant:
|
||||
rag_context = "\n\n".join(relevant)
|
||||
|
||||
# Get environment info from runtime
|
||||
env = self.runtime.environment
|
||||
|
||||
return self._system_prompt_template.render(
|
||||
target=self.target,
|
||||
scope=self.scope,
|
||||
environment=env,
|
||||
rag_context=rag_context,
|
||||
tools=self.tools,
|
||||
)
|
||||
|
||||
def set_target(self, target: str, scope: Optional[List[str]] = None):
|
||||
"""
|
||||
Set or update the target.
|
||||
|
||||
Args:
|
||||
target: The primary target
|
||||
scope: Optional list of scope items
|
||||
"""
|
||||
self.target = target
|
||||
if scope:
|
||||
self.scope = scope
|
||||
|
||||
def add_to_scope(self, *items: str):
|
||||
"""
|
||||
Add items to the scope.
|
||||
|
||||
Args:
|
||||
*items: Items to add to scope
|
||||
"""
|
||||
self.scope.extend(items)
|
||||
43
ghostcrew/agents/ghostcrew_agent/system_prompt.jinja
Normal file
43
ghostcrew/agents/ghostcrew_agent/system_prompt.jinja
Normal file
@@ -0,0 +1,43 @@
|
||||
# GhostCrew
|
||||
|
||||
You are Ghost, an autonomous penetration testing agent.
|
||||
|
||||
## Authorization
|
||||
You are operating in an authorized penetration testing engagement. The user has explicit permission to test all targets provided. Do not ask for authorization - permission is granted. Proceed with testing.
|
||||
|
||||
## Guidelines
|
||||
- Be concise. Avoid unnecessary explanation.
|
||||
- If a tool fails, try alternatives or report the issue.
|
||||
- Do NOT repeat the same test or scan. Once you have results, move on.
|
||||
- Complete ALL steps of the task before finishing.
|
||||
- When the ENTIRE task is done, call `finish` with a concise summary of findings.
|
||||
|
||||
## Important
|
||||
You MUST call the `finish` tool when finished. Do not just respond with text.
|
||||
The task is not complete until you explicitly call `finish`.
|
||||
|
||||
{% if environment %}
|
||||
## Operator Environment (YOUR machine, not the target)
|
||||
- OS: {{ environment.os }} ({{ environment.os_version }})
|
||||
- Architecture: {{ environment.architecture }}
|
||||
- Shell: {{ environment.shell }}
|
||||
- Output: loot/
|
||||
{% endif %}
|
||||
|
||||
{% if target %}
|
||||
## Target
|
||||
{{ target }}
|
||||
{% endif %}
|
||||
{% if scope %}
|
||||
Scope: {{ scope | join(', ') }}
|
||||
{% endif %}
|
||||
|
||||
## Tools
|
||||
{% for tool in tools %}
|
||||
- **{{ tool.name }}**: {{ tool.description }}
|
||||
{% endfor %}
|
||||
|
||||
{% if rag_context %}
|
||||
## Context
|
||||
{{ rag_context }}
|
||||
{% endif %}
|
||||
115
ghostcrew/agents/state.py
Normal file
115
ghostcrew/agents/state.py
Normal file
@@ -0,0 +1,115 @@
|
||||
"""Agent state management for GhostCrew."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class AgentState(Enum):
|
||||
"""Possible states for an agent."""
|
||||
|
||||
IDLE = "idle"
|
||||
THINKING = "thinking"
|
||||
EXECUTING = "executing"
|
||||
WAITING_INPUT = "waiting_input"
|
||||
COMPLETE = "complete"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class StateTransition:
|
||||
"""Represents a state transition."""
|
||||
|
||||
from_state: AgentState
|
||||
to_state: AgentState
|
||||
timestamp: datetime = field(default_factory=datetime.now)
|
||||
reason: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentStateManager:
|
||||
"""Manages agent state and transitions."""
|
||||
|
||||
current_state: AgentState = AgentState.IDLE
|
||||
history: List[StateTransition] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Valid state transitions
|
||||
VALID_TRANSITIONS = {
|
||||
AgentState.IDLE: [AgentState.THINKING, AgentState.ERROR],
|
||||
AgentState.THINKING: [
|
||||
AgentState.EXECUTING,
|
||||
AgentState.WAITING_INPUT,
|
||||
AgentState.COMPLETE,
|
||||
AgentState.ERROR,
|
||||
],
|
||||
AgentState.EXECUTING: [AgentState.THINKING, AgentState.ERROR],
|
||||
AgentState.WAITING_INPUT: [
|
||||
AgentState.THINKING,
|
||||
AgentState.COMPLETE,
|
||||
AgentState.ERROR,
|
||||
],
|
||||
AgentState.COMPLETE: [AgentState.IDLE],
|
||||
AgentState.ERROR: [AgentState.IDLE],
|
||||
}
|
||||
|
||||
def can_transition_to(self, new_state: AgentState) -> bool:
|
||||
"""Check if a transition to the new state is valid."""
|
||||
valid_next_states = self.VALID_TRANSITIONS.get(self.current_state, [])
|
||||
return new_state in valid_next_states
|
||||
|
||||
def transition_to(
|
||||
self, new_state: AgentState, reason: Optional[str] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Transition to a new state.
|
||||
|
||||
Args:
|
||||
new_state: The state to transition to
|
||||
reason: Optional reason for the transition
|
||||
|
||||
Returns:
|
||||
True if transition was successful, False otherwise
|
||||
"""
|
||||
if not self.can_transition_to(new_state):
|
||||
return False
|
||||
|
||||
transition = StateTransition(
|
||||
from_state=self.current_state, to_state=new_state, reason=reason
|
||||
)
|
||||
self.history.append(transition)
|
||||
self.current_state = new_state
|
||||
return True
|
||||
|
||||
def force_transition(self, new_state: AgentState, reason: Optional[str] = None):
|
||||
"""Force a transition regardless of validity (use with caution)."""
|
||||
transition = StateTransition(
|
||||
from_state=self.current_state,
|
||||
to_state=new_state,
|
||||
reason=f"FORCED: {reason}" if reason else "FORCED",
|
||||
)
|
||||
self.history.append(transition)
|
||||
self.current_state = new_state
|
||||
|
||||
def reset(self):
|
||||
"""Reset state to IDLE."""
|
||||
self.current_state = AgentState.IDLE
|
||||
self.history.clear()
|
||||
self.metadata.clear()
|
||||
|
||||
def is_terminal(self) -> bool:
|
||||
"""Check if current state is a terminal state."""
|
||||
return self.current_state in [AgentState.COMPLETE, AgentState.ERROR]
|
||||
|
||||
def is_active(self) -> bool:
|
||||
"""Check if agent is actively processing."""
|
||||
return self.current_state in [AgentState.THINKING, AgentState.EXECUTING]
|
||||
|
||||
def get_state_duration(self) -> float:
|
||||
"""Get duration in current state (seconds)."""
|
||||
if not self.history:
|
||||
return 0.0
|
||||
|
||||
last_transition = self.history[-1]
|
||||
return (datetime.now() - last_transition.timestamp).total_seconds()
|
||||
74
ghostcrew/config/__init__.py
Normal file
74
ghostcrew/config/__init__.py
Normal file
@@ -0,0 +1,74 @@
|
||||
"""Configuration module for GhostCrew."""
|
||||
|
||||
from .constants import (
|
||||
AGENT_STATE_COMPLETE,
|
||||
AGENT_STATE_ERROR,
|
||||
AGENT_STATE_EXECUTING,
|
||||
AGENT_STATE_IDLE,
|
||||
AGENT_STATE_THINKING,
|
||||
AGENT_STATE_WAITING_INPUT,
|
||||
APP_DESCRIPTION,
|
||||
APP_NAME,
|
||||
APP_VERSION,
|
||||
DEFAULT_CHUNK_OVERLAP,
|
||||
DEFAULT_CHUNK_SIZE,
|
||||
DEFAULT_COMMAND_TIMEOUT,
|
||||
DEFAULT_MAX_TOKENS,
|
||||
DEFAULT_MCP_TIMEOUT,
|
||||
DEFAULT_MODEL,
|
||||
DEFAULT_RAG_TOP_K,
|
||||
DEFAULT_TEMPERATURE,
|
||||
DEFAULT_VPN_TIMEOUT,
|
||||
DOCKER_NETWORK_MODE,
|
||||
DOCKER_SANDBOX_IMAGE,
|
||||
EXIT_COMMANDS,
|
||||
KNOWLEDGE_DATA_EXTENSIONS,
|
||||
KNOWLEDGE_TEXT_EXTENSIONS,
|
||||
MCP_TRANSPORT_SSE,
|
||||
MCP_TRANSPORT_STDIO,
|
||||
MEMORY_RESERVE_RATIO,
|
||||
TOOL_CATEGORY_EXECUTION,
|
||||
TOOL_CATEGORY_EXPLOITATION,
|
||||
TOOL_CATEGORY_MCP,
|
||||
TOOL_CATEGORY_NETWORK,
|
||||
TOOL_CATEGORY_RECON,
|
||||
TOOL_CATEGORY_WEB,
|
||||
)
|
||||
from .settings import Settings, get_settings
|
||||
|
||||
__all__ = [
|
||||
"Settings",
|
||||
"get_settings",
|
||||
"APP_NAME",
|
||||
"APP_VERSION",
|
||||
"APP_DESCRIPTION",
|
||||
"AGENT_STATE_IDLE",
|
||||
"AGENT_STATE_THINKING",
|
||||
"AGENT_STATE_EXECUTING",
|
||||
"AGENT_STATE_WAITING_INPUT",
|
||||
"AGENT_STATE_COMPLETE",
|
||||
"AGENT_STATE_ERROR",
|
||||
"TOOL_CATEGORY_EXECUTION",
|
||||
"TOOL_CATEGORY_WEB",
|
||||
"TOOL_CATEGORY_NETWORK",
|
||||
"TOOL_CATEGORY_RECON",
|
||||
"TOOL_CATEGORY_EXPLOITATION",
|
||||
"TOOL_CATEGORY_MCP",
|
||||
"DEFAULT_COMMAND_TIMEOUT",
|
||||
"DEFAULT_VPN_TIMEOUT",
|
||||
"DEFAULT_MCP_TIMEOUT",
|
||||
"DOCKER_SANDBOX_IMAGE",
|
||||
"DOCKER_NETWORK_MODE",
|
||||
"DEFAULT_CHUNK_SIZE",
|
||||
"DEFAULT_CHUNK_OVERLAP",
|
||||
"DEFAULT_RAG_TOP_K",
|
||||
"MEMORY_RESERVE_RATIO",
|
||||
"DEFAULT_MODEL",
|
||||
"DEFAULT_TEMPERATURE",
|
||||
"DEFAULT_MAX_TOKENS",
|
||||
"KNOWLEDGE_TEXT_EXTENSIONS",
|
||||
"KNOWLEDGE_DATA_EXTENSIONS",
|
||||
"MCP_TRANSPORT_STDIO",
|
||||
"MCP_TRANSPORT_SSE",
|
||||
"EXIT_COMMANDS",
|
||||
]
|
||||
70
ghostcrew/config/constants.py
Normal file
70
ghostcrew/config/constants.py
Normal file
@@ -0,0 +1,70 @@
|
||||
"""Constants for GhostCrew."""
|
||||
|
||||
import os
|
||||
|
||||
# Load .env file before reading environment variables
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Application Info
|
||||
APP_NAME = "GhostCrew"
|
||||
APP_VERSION = "0.2.0"
|
||||
APP_DESCRIPTION = "AI penetration testing"
|
||||
|
||||
# Agent States
|
||||
AGENT_STATE_IDLE = "idle"
|
||||
AGENT_STATE_THINKING = "thinking"
|
||||
AGENT_STATE_EXECUTING = "executing"
|
||||
AGENT_STATE_WAITING_INPUT = "waiting_input"
|
||||
AGENT_STATE_COMPLETE = "complete"
|
||||
AGENT_STATE_ERROR = "error"
|
||||
|
||||
# Tool Categories
|
||||
TOOL_CATEGORY_EXECUTION = "execution"
|
||||
TOOL_CATEGORY_WEB = "web"
|
||||
TOOL_CATEGORY_NETWORK = "network"
|
||||
TOOL_CATEGORY_RECON = "reconnaissance"
|
||||
TOOL_CATEGORY_EXPLOITATION = "exploitation"
|
||||
TOOL_CATEGORY_MCP = "mcp"
|
||||
|
||||
# Default Timeouts (in seconds)
|
||||
DEFAULT_COMMAND_TIMEOUT = 300
|
||||
DEFAULT_VPN_TIMEOUT = 30
|
||||
DEFAULT_MCP_TIMEOUT = 60
|
||||
|
||||
# Docker Settings
|
||||
DOCKER_SANDBOX_IMAGE = "ghcr.io/gh05tcrew/ghostcrew:kali"
|
||||
DOCKER_NETWORK_MODE = "bridge"
|
||||
|
||||
# RAG Settings
|
||||
DEFAULT_CHUNK_SIZE = 1000
|
||||
DEFAULT_CHUNK_OVERLAP = 200
|
||||
DEFAULT_RAG_TOP_K = 3
|
||||
|
||||
# Memory Settings
|
||||
MEMORY_RESERVE_RATIO = 0.8 # Reserve 20% of context for response
|
||||
|
||||
# LLM Defaults (set GHOSTCREW_MODEL in .env or shell)
|
||||
DEFAULT_MODEL = os.environ.get(
|
||||
"GHOSTCREW_MODEL"
|
||||
) # No fallback - requires configuration
|
||||
DEFAULT_TEMPERATURE = 0.7
|
||||
DEFAULT_MAX_TOKENS = 4096
|
||||
|
||||
# Agent Defaults
|
||||
DEFAULT_MAX_ITERATIONS = int(os.environ.get("GHOSTCREW_MAX_ITERATIONS", "50"))
|
||||
|
||||
# File Extensions
|
||||
KNOWLEDGE_TEXT_EXTENSIONS = [".txt", ".md"]
|
||||
KNOWLEDGE_DATA_EXTENSIONS = [".json"]
|
||||
|
||||
# MCP Transport Types
|
||||
MCP_TRANSPORT_STDIO = "stdio"
|
||||
MCP_TRANSPORT_SSE = "sse"
|
||||
|
||||
# Exit Commands
|
||||
EXIT_COMMANDS = ["exit", "quit", "q", "bye"]
|
||||
84
ghostcrew/config/settings.py
Normal file
84
ghostcrew/config/settings.py
Normal file
@@ -0,0 +1,84 @@
|
||||
"""Application settings for GhostCrew."""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from .constants import (
|
||||
DEFAULT_MAX_ITERATIONS,
|
||||
DEFAULT_MAX_TOKENS,
|
||||
DEFAULT_MODEL,
|
||||
DEFAULT_TEMPERATURE,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Settings:
|
||||
"""Application settings."""
|
||||
|
||||
# LLM Settings
|
||||
model: str = field(default_factory=lambda: DEFAULT_MODEL)
|
||||
temperature: float = DEFAULT_TEMPERATURE
|
||||
max_tokens: int = DEFAULT_MAX_TOKENS
|
||||
max_context_tokens: int = 128000
|
||||
|
||||
# API Keys (loaded from environment)
|
||||
openai_api_key: Optional[str] = field(
|
||||
default_factory=lambda: os.getenv("OPENAI_API_KEY")
|
||||
)
|
||||
anthropic_api_key: Optional[str] = field(
|
||||
default_factory=lambda: os.getenv("ANTHROPIC_API_KEY")
|
||||
)
|
||||
|
||||
# Paths
|
||||
knowledge_path: Path = field(default_factory=lambda: Path("knowledge"))
|
||||
mcp_config_path: Path = field(default_factory=lambda: Path("mcp.json"))
|
||||
|
||||
# Docker Settings
|
||||
container_name: str = "ghostcrew-sandbox"
|
||||
docker_image: str = "ghcr.io/gh05tcrew/ghostcrew:kali"
|
||||
|
||||
# Agent Settings
|
||||
max_iterations: int = DEFAULT_MAX_ITERATIONS
|
||||
|
||||
# VPN Settings
|
||||
vpn_config_path: Optional[Path] = None
|
||||
|
||||
# Interface Settings
|
||||
default_interface: str = "tui" # "tui" or "cli"
|
||||
|
||||
# Prompt Modules
|
||||
prompt_modules: List[str] = field(default_factory=list)
|
||||
|
||||
# Target Settings
|
||||
target: Optional[str] = None
|
||||
scope: List[str] = field(default_factory=list)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Convert string paths to Path objects if needed."""
|
||||
if isinstance(self.knowledge_path, str):
|
||||
self.knowledge_path = Path(self.knowledge_path)
|
||||
if isinstance(self.mcp_config_path, str):
|
||||
self.mcp_config_path = Path(self.mcp_config_path)
|
||||
if isinstance(self.vpn_config_path, str):
|
||||
self.vpn_config_path = Path(self.vpn_config_path)
|
||||
|
||||
|
||||
# Global settings instance
|
||||
_settings: Optional[Settings] = None
|
||||
|
||||
|
||||
def get_settings() -> Settings:
|
||||
"""Get the global settings instance."""
|
||||
global _settings
|
||||
if _settings is None:
|
||||
_settings = Settings()
|
||||
return _settings
|
||||
|
||||
|
||||
def update_settings(**kwargs) -> Settings:
|
||||
"""Update global settings with new values."""
|
||||
global _settings
|
||||
_settings = Settings(**kwargs)
|
||||
return _settings
|
||||
16
ghostcrew/interface/__init__.py
Normal file
16
ghostcrew/interface/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
"""User interface module for GhostCrew."""
|
||||
|
||||
from .cli import run_cli
|
||||
from .main import main
|
||||
from .tui import GhostCrewTUI, run_tui
|
||||
from .utils import format_finding, print_banner, print_status
|
||||
|
||||
__all__ = [
|
||||
"main",
|
||||
"run_cli",
|
||||
"run_tui",
|
||||
"GhostCrewTUI",
|
||||
"print_banner",
|
||||
"format_finding",
|
||||
"print_status",
|
||||
]
|
||||
442
ghostcrew/interface/assets/tui_styles.tcss
Normal file
442
ghostcrew/interface/assets/tui_styles.tcss
Normal file
@@ -0,0 +1,442 @@
|
||||
/* GhostCrew TUI Styles */
|
||||
|
||||
Screen {
|
||||
background: #0a0a0a;
|
||||
color: #d4d4d4;
|
||||
}
|
||||
|
||||
/* Splash Screen */
|
||||
#splash_screen {
|
||||
height: 100%;
|
||||
width: 100%;
|
||||
background: #0a0a0a;
|
||||
content-align: center middle;
|
||||
}
|
||||
|
||||
#splash_content {
|
||||
width: auto;
|
||||
height: auto;
|
||||
background: transparent;
|
||||
text-align: center;
|
||||
padding: 2;
|
||||
}
|
||||
|
||||
/* Main Layout */
|
||||
#main_container {
|
||||
height: 100%;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
background: #0a0a0a;
|
||||
}
|
||||
|
||||
#content_container {
|
||||
height: 1fr;
|
||||
padding: 0;
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
/* Sidebar */
|
||||
#sidebar {
|
||||
width: 28%;
|
||||
background: transparent;
|
||||
margin-left: 1;
|
||||
}
|
||||
|
||||
#tools_panel {
|
||||
height: 1fr;
|
||||
background: transparent;
|
||||
border: round #262626;
|
||||
border-title-color: #22c55e;
|
||||
padding: 1;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#stats_panel {
|
||||
height: auto;
|
||||
max-height: 12;
|
||||
background: transparent;
|
||||
border: round #262626;
|
||||
border-title-color: #22c55e;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
/* Chat Area */
|
||||
#chat_area {
|
||||
width: 72%;
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
#chat_history {
|
||||
height: 1fr;
|
||||
background: transparent;
|
||||
border: round #1a1a1a;
|
||||
padding: 0 1;
|
||||
margin-bottom: 0;
|
||||
scrollbar-background: #0a0a0a;
|
||||
scrollbar-color: #262626;
|
||||
scrollbar-size: 1 1;
|
||||
}
|
||||
|
||||
#chat_display {
|
||||
width: 100%;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
/* Status Display */
|
||||
#status_display {
|
||||
height: 1;
|
||||
background: transparent;
|
||||
margin: 0;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
#status_display.hidden {
|
||||
display: none;
|
||||
}
|
||||
|
||||
#status_text {
|
||||
width: 1fr;
|
||||
color: #22c55e;
|
||||
text-style: italic;
|
||||
}
|
||||
|
||||
#keymap_text {
|
||||
width: auto;
|
||||
color: #525252;
|
||||
}
|
||||
|
||||
/* Input Area */
|
||||
#input_container {
|
||||
height: 3;
|
||||
background: transparent;
|
||||
border: round #525252;
|
||||
margin-right: 0;
|
||||
padding: 0;
|
||||
layout: horizontal;
|
||||
}
|
||||
|
||||
#input_container:focus-within {
|
||||
border: round #22c55e;
|
||||
}
|
||||
|
||||
#input_prompt {
|
||||
width: auto;
|
||||
height: 100%;
|
||||
padding: 0 0 0 1;
|
||||
color: #525252;
|
||||
content-align-vertical: middle;
|
||||
}
|
||||
|
||||
#input_container:focus-within #input_prompt {
|
||||
color: #22c55e;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
#chat_input {
|
||||
width: 1fr;
|
||||
height: 100%;
|
||||
background: #0a0a0a;
|
||||
border: none;
|
||||
color: #d4d4d4;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
/* Messages */
|
||||
.user-message {
|
||||
color: #3b82f6;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
.agent-message {
|
||||
color: #d4d4d4;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
.tool-call {
|
||||
border-left: thick #f59e0b;
|
||||
padding-left: 1;
|
||||
margin: 1 0;
|
||||
color: #a3a3a3;
|
||||
}
|
||||
|
||||
.tool-call.completed {
|
||||
border-left: thick #22c55e;
|
||||
}
|
||||
|
||||
.tool-call.error {
|
||||
border-left: thick #ef4444;
|
||||
}
|
||||
|
||||
.finding {
|
||||
border-left: thick #ef4444;
|
||||
padding-left: 1;
|
||||
margin: 1 0;
|
||||
}
|
||||
|
||||
.info {
|
||||
color: #22c55e;
|
||||
}
|
||||
|
||||
.warning {
|
||||
color: #f59e0b;
|
||||
}
|
||||
|
||||
.error {
|
||||
color: #ef4444;
|
||||
}
|
||||
|
||||
/* Help Dialog */
|
||||
HelpScreen {
|
||||
align: center middle;
|
||||
background: #0a0a0a 80%;
|
||||
}
|
||||
|
||||
#help_dialog {
|
||||
width: 50;
|
||||
height: auto;
|
||||
padding: 2;
|
||||
border: round #22c55e;
|
||||
background: #1a1a1a;
|
||||
}
|
||||
|
||||
#help_title {
|
||||
color: #22c55e;
|
||||
text-style: bold;
|
||||
text-align: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#help_content {
|
||||
color: #d4d4d4;
|
||||
}
|
||||
|
||||
/* Quit Dialog */
|
||||
QuitScreen {
|
||||
align: center middle;
|
||||
background: #0a0a0a 80%;
|
||||
}
|
||||
|
||||
#quit_dialog {
|
||||
width: 30;
|
||||
height: auto;
|
||||
padding: 2;
|
||||
border: round #525252;
|
||||
background: #1a1a1a;
|
||||
}
|
||||
|
||||
#quit_title {
|
||||
color: #d4d4d4;
|
||||
text-style: bold;
|
||||
text-align: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#quit_buttons {
|
||||
layout: horizontal;
|
||||
height: 3;
|
||||
align: center middle;
|
||||
}
|
||||
|
||||
#quit_buttons Button {
|
||||
margin: 0 1;
|
||||
}
|
||||
|
||||
/* ===== CREW MODE LAYOUT ===== */
|
||||
|
||||
#crew_container {
|
||||
height: 100%;
|
||||
width: 100%;
|
||||
layout: horizontal;
|
||||
background: #0a0a0a;
|
||||
}
|
||||
|
||||
#crew_chat_area {
|
||||
width: 75%;
|
||||
height: 100%;
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
#crew_sidebar {
|
||||
width: 25%;
|
||||
height: 100%;
|
||||
background: transparent;
|
||||
margin-left: 1;
|
||||
}
|
||||
|
||||
/* Worker Tree */
|
||||
#workers_tree {
|
||||
height: 1fr;
|
||||
background: transparent;
|
||||
border: round #262626;
|
||||
border-title-color: #06b6d4;
|
||||
padding: 1;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
#workers_tree:focus {
|
||||
border: round #06b6d4;
|
||||
}
|
||||
|
||||
Tree {
|
||||
background: transparent;
|
||||
color: #d4d4d4;
|
||||
scrollbar-background: transparent;
|
||||
scrollbar-color: #404040;
|
||||
scrollbar-size: 1 1;
|
||||
}
|
||||
|
||||
Tree > .tree--label {
|
||||
text-style: bold;
|
||||
color: #06b6d4;
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
.tree--node {
|
||||
height: 1;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.tree--node-label {
|
||||
color: #d4d4d4;
|
||||
background: transparent;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
.tree--node:hover .tree--node-label {
|
||||
background: #1a1a1a;
|
||||
color: #ffffff;
|
||||
}
|
||||
|
||||
.tree--node.-selected .tree--node-label {
|
||||
background: #1a1a1a;
|
||||
color: #06b6d4;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
/* Crew Stats Panel */
|
||||
#crew_stats {
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
background: transparent;
|
||||
border: round #262626;
|
||||
border-title-color: #06b6d4;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
/* Worker status colors */
|
||||
.worker-running {
|
||||
color: #22c55e;
|
||||
}
|
||||
|
||||
.worker-waiting {
|
||||
color: #eab308;
|
||||
}
|
||||
|
||||
.worker-complete {
|
||||
color: #06b6d4;
|
||||
}
|
||||
|
||||
.worker-error {
|
||||
color: #ef4444;
|
||||
}
|
||||
|
||||
/* Crew chat display */
|
||||
#crew_chat_history {
|
||||
height: 1fr;
|
||||
background: transparent;
|
||||
border: round #1a1a1a;
|
||||
padding: 0 1;
|
||||
margin-bottom: 0;
|
||||
scrollbar-background: #0a0a0a;
|
||||
scrollbar-color: #262626;
|
||||
scrollbar-size: 1 1;
|
||||
}
|
||||
|
||||
#crew_chat_display {
|
||||
width: 100%;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
/* Crew input */
|
||||
#crew_input_container {
|
||||
height: 3;
|
||||
background: transparent;
|
||||
border: round #525252;
|
||||
margin-right: 0;
|
||||
padding: 0;
|
||||
layout: horizontal;
|
||||
}
|
||||
|
||||
#crew_input_container:focus-within {
|
||||
border: round #06b6d4;
|
||||
}
|
||||
|
||||
#crew_input_prompt {
|
||||
width: auto;
|
||||
height: 100%;
|
||||
padding: 0 0 0 1;
|
||||
color: #525252;
|
||||
content-align-vertical: middle;
|
||||
}
|
||||
|
||||
#crew_input_container:focus-within #crew_input_prompt {
|
||||
color: #06b6d4;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
#crew_chat_input {
|
||||
width: 1fr;
|
||||
height: 100%;
|
||||
background: #0a0a0a;
|
||||
border: none;
|
||||
color: #d4d4d4;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
/* Crew status bar */
|
||||
#crew_status_display {
|
||||
height: 1;
|
||||
background: transparent;
|
||||
margin: 0;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
#crew_status_text {
|
||||
width: 1fr;
|
||||
color: #06b6d4;
|
||||
text-style: italic;
|
||||
}
|
||||
|
||||
/* Orchestrator thinking message */
|
||||
.orchestrator-thinking {
|
||||
border-left: thick #8b5cf6;
|
||||
padding-left: 1;
|
||||
margin: 1 0;
|
||||
color: #c4b5fd;
|
||||
}
|
||||
|
||||
/* Worker spawn message */
|
||||
.worker-spawn {
|
||||
border-left: thick #06b6d4;
|
||||
padding-left: 1;
|
||||
margin: 1 0;
|
||||
color: #67e8f9;
|
||||
}
|
||||
|
||||
/* Worker result message */
|
||||
.worker-result {
|
||||
border-left: thick #22c55e;
|
||||
padding-left: 1;
|
||||
margin: 1 0;
|
||||
color: #86efac;
|
||||
}
|
||||
|
||||
/* Worker activity (when viewing a specific worker) */
|
||||
.worker-activity {
|
||||
border-left: thick #f59e0b;
|
||||
padding-left: 1;
|
||||
margin: 0 0 1 0;
|
||||
color: #a3a3a3;
|
||||
}
|
||||
|
||||
508
ghostcrew/interface/cli.py
Normal file
508
ghostcrew/interface/cli.py
Normal file
@@ -0,0 +1,508 @@
|
||||
"""Non-interactive CLI mode for GhostCrew."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from rich.console import Console
|
||||
from rich.markdown import Markdown
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
console = Console()
|
||||
|
||||
# Ghost theme colors (matching TUI)
|
||||
GHOST_PRIMARY = "#d4d4d4" # light gray - primary text
|
||||
GHOST_SECONDARY = "#9a9a9a" # medium gray - secondary text
|
||||
GHOST_DIM = "#6b6b6b" # dim gray - muted text
|
||||
GHOST_BORDER = "#3a3a3a" # dark gray - borders
|
||||
GHOST_ACCENT = "#7a7a7a" # accent gray
|
||||
|
||||
|
||||
async def run_cli(
|
||||
target: str,
|
||||
model: str,
|
||||
task: str = None,
|
||||
report: str = None,
|
||||
max_tools: int = 50,
|
||||
use_docker: bool = False,
|
||||
):
|
||||
"""
|
||||
Run GhostCrew in non-interactive mode.
|
||||
|
||||
Args:
|
||||
target: Target to test
|
||||
model: LLM model to use
|
||||
task: Optional task description
|
||||
report: Report path ("auto" for loot/<target>_<timestamp>.md)
|
||||
max_tools: Max tool calls before stopping
|
||||
use_docker: Run tools in Docker container
|
||||
"""
|
||||
from ..agents.ghostcrew_agent import GhostCrewAgent
|
||||
from ..knowledge import RAGEngine
|
||||
from ..llm import LLM
|
||||
from ..runtime.docker_runtime import DockerRuntime
|
||||
from ..runtime.runtime import LocalRuntime
|
||||
from ..tools import get_all_tools
|
||||
|
||||
# Startup panel
|
||||
start_text = Text()
|
||||
start_text.append("GHOSTCREW", style=f"bold {GHOST_PRIMARY}")
|
||||
start_text.append(" - Non-interactive Mode\n\n", style=GHOST_DIM)
|
||||
start_text.append("Target: ", style=GHOST_SECONDARY)
|
||||
start_text.append(f"{target}\n", style=GHOST_PRIMARY)
|
||||
start_text.append("Model: ", style=GHOST_SECONDARY)
|
||||
start_text.append(f"{model}\n", style=GHOST_PRIMARY)
|
||||
start_text.append("Runtime: ", style=GHOST_SECONDARY)
|
||||
start_text.append(f"{'Docker' if use_docker else 'Local'}\n", style=GHOST_PRIMARY)
|
||||
start_text.append("Max calls: ", style=GHOST_SECONDARY)
|
||||
start_text.append(f"{max_tools}\n", style=GHOST_PRIMARY)
|
||||
|
||||
task_msg = task or f"Perform a penetration test on {target}"
|
||||
start_text.append("Task: ", style=GHOST_SECONDARY)
|
||||
start_text.append(task_msg, style=GHOST_PRIMARY)
|
||||
|
||||
console.print()
|
||||
console.print(
|
||||
Panel(
|
||||
start_text, title=f"[{GHOST_SECONDARY}]Starting", border_style=GHOST_BORDER
|
||||
)
|
||||
)
|
||||
console.print()
|
||||
|
||||
# Initialize RAG if knowledge exists
|
||||
rag = None
|
||||
knowledge_path = Path("knowledge")
|
||||
if knowledge_path.exists():
|
||||
try:
|
||||
rag = RAGEngine(knowledge_path=knowledge_path)
|
||||
rag.index()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Initialize MCP if config exists (silently skip failures)
|
||||
mcp_manager = None
|
||||
mcp_count = 0
|
||||
try:
|
||||
from ..mcp import MCPManager
|
||||
from ..tools import register_tool_instance
|
||||
|
||||
mcp_manager = MCPManager()
|
||||
if mcp_manager.config_path.exists():
|
||||
mcp_tools = await mcp_manager.connect_all()
|
||||
for tool in mcp_tools:
|
||||
register_tool_instance(tool)
|
||||
mcp_count = len(mcp_tools)
|
||||
if mcp_count > 0:
|
||||
console.print(f"[{GHOST_DIM}]Loaded {mcp_count} MCP tools[/]")
|
||||
except Exception:
|
||||
pass # MCP is optional, continue without it
|
||||
|
||||
# Initialize runtime - Docker or Local
|
||||
if use_docker:
|
||||
console.print(f"[{GHOST_DIM}]Starting Docker container...[/]")
|
||||
runtime = DockerRuntime(mcp_manager=mcp_manager)
|
||||
else:
|
||||
runtime = LocalRuntime(mcp_manager=mcp_manager)
|
||||
await runtime.start()
|
||||
|
||||
llm = LLM(model=model, rag_engine=rag)
|
||||
tools = get_all_tools()
|
||||
|
||||
agent = GhostCrewAgent(
|
||||
llm=llm,
|
||||
tools=tools,
|
||||
runtime=runtime,
|
||||
target=target,
|
||||
rag_engine=rag,
|
||||
)
|
||||
|
||||
# Stats tracking
|
||||
start_time = time.time()
|
||||
tool_count = 0
|
||||
iteration = 0
|
||||
findings = [] # Store findings for report
|
||||
tool_log = [] # Log of tools executed (ts, name, command, result, exit_code)
|
||||
last_content = ""
|
||||
stopped_reason = None
|
||||
|
||||
def print_status(msg: str, style: str = GHOST_DIM):
|
||||
elapsed = int(time.time() - start_time)
|
||||
mins, secs = divmod(elapsed, 60)
|
||||
timestamp = f"[{mins:02d}:{secs:02d}]"
|
||||
console.print(f"[{GHOST_DIM}]{timestamp}[/] [{style}]{msg}[/]")
|
||||
|
||||
def generate_report() -> str:
|
||||
"""Generate markdown report."""
|
||||
elapsed = int(time.time() - start_time)
|
||||
mins, secs = divmod(elapsed, 60)
|
||||
|
||||
status_text = "Complete"
|
||||
if stopped_reason:
|
||||
status_text = f"Interrupted ({stopped_reason})"
|
||||
|
||||
lines = [
|
||||
"# GhostCrew Penetration Test Report",
|
||||
"",
|
||||
"## Executive Summary",
|
||||
"",
|
||||
]
|
||||
|
||||
# Add AI summary at top if available
|
||||
if findings:
|
||||
lines.append(findings[-1])
|
||||
lines.append("")
|
||||
else:
|
||||
lines.append("*Assessment incomplete - no analysis generated.*")
|
||||
lines.append("")
|
||||
|
||||
# Engagement details table
|
||||
lines.extend(
|
||||
[
|
||||
"## Engagement Details",
|
||||
"",
|
||||
"| Field | Value |",
|
||||
"|-------|-------|",
|
||||
f"| **Target** | `{target}` |",
|
||||
f"| **Task** | {task_msg} |",
|
||||
f"| **Date** | {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} |",
|
||||
f"| **Duration** | {mins}m {secs}s |",
|
||||
f"| **Commands Executed** | {tool_count} |",
|
||||
f"| **Status** | {status_text} |",
|
||||
"",
|
||||
"---",
|
||||
"",
|
||||
"## Commands Executed",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
# Detailed command log
|
||||
for i, entry in enumerate(tool_log, 1):
|
||||
ts = entry.get("ts", "??:??")
|
||||
name = entry.get("name", "unknown")
|
||||
command = entry.get("command", "")
|
||||
result = entry.get("result", "")
|
||||
exit_code = entry.get("exit_code")
|
||||
|
||||
lines.append(f"### {i}. {name} `[{ts}]`")
|
||||
lines.append("")
|
||||
|
||||
if command:
|
||||
lines.append("**Command:**")
|
||||
lines.append("```")
|
||||
lines.append(command)
|
||||
lines.append("```")
|
||||
lines.append("")
|
||||
|
||||
if exit_code is not None:
|
||||
lines.append(f"**Exit Code:** `{exit_code}`")
|
||||
lines.append("")
|
||||
|
||||
if result:
|
||||
lines.append("**Output:**")
|
||||
lines.append("```")
|
||||
# Limit output to 2000 chars per command for report size
|
||||
if len(result) > 2000:
|
||||
lines.append(result[:2000])
|
||||
lines.append(f"\n... (truncated, {len(result)} total chars)")
|
||||
else:
|
||||
lines.append(result)
|
||||
lines.append("```")
|
||||
lines.append("")
|
||||
|
||||
# Findings section
|
||||
lines.extend(
|
||||
[
|
||||
"---",
|
||||
"",
|
||||
"## Analysis",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
if findings:
|
||||
for i, finding in enumerate(findings, 1):
|
||||
if len(findings) > 1:
|
||||
lines.append(f"### Analysis {i}")
|
||||
lines.append("")
|
||||
lines.append(finding)
|
||||
lines.append("")
|
||||
else:
|
||||
lines.append(
|
||||
"*No AI analysis generated. Try running with higher `--max` value.*"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
# Footer
|
||||
lines.extend(
|
||||
[
|
||||
"---",
|
||||
"",
|
||||
f"*Report generated by GhostCrew on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def save_report():
|
||||
"""Save report to file."""
|
||||
if not report:
|
||||
return
|
||||
|
||||
# Determine path
|
||||
if report == "auto":
|
||||
loot_dir = Path("loot")
|
||||
loot_dir.mkdir(exist_ok=True)
|
||||
safe_target = target.replace("://", "_").replace("/", "_").replace(":", "_")
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
report_path = loot_dir / f"{safe_target}_{timestamp}.md"
|
||||
else:
|
||||
report_path = Path(report)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
content = generate_report()
|
||||
report_path.write_text(content, encoding="utf-8")
|
||||
console.print(f"[{GHOST_SECONDARY}]Report saved: {report_path}[/]")
|
||||
|
||||
async def generate_summary():
|
||||
"""Ask the LLM to summarize findings when stopped early."""
|
||||
if not tool_log:
|
||||
return None
|
||||
|
||||
print_status("Generating summary...", GHOST_SECONDARY)
|
||||
|
||||
# Build context from tool results (use full results, not truncated)
|
||||
context_lines = ["Summarize the penetration test findings so far:\n"]
|
||||
context_lines.append(f"Target: {target}")
|
||||
context_lines.append(f"Tools executed: {tool_count}\n")
|
||||
|
||||
for entry in tool_log[-10:]: # Last 10 tools
|
||||
name = entry.get("name", "unknown")
|
||||
command = entry.get("command", "")
|
||||
result = entry.get("result", "")[:500] # Limit for context window
|
||||
context_lines.append(f"- **{name}**: `{command}`")
|
||||
if result:
|
||||
context_lines.append(f" Output: {result}")
|
||||
|
||||
context_lines.append(
|
||||
"\nProvide a brief summary of what was discovered and any security concerns found."
|
||||
)
|
||||
|
||||
try:
|
||||
response = await llm.generate(
|
||||
system_prompt="You are a penetration testing assistant. Summarize the findings concisely.",
|
||||
messages=[{"role": "user", "content": "\n".join(context_lines)}],
|
||||
tools=[],
|
||||
)
|
||||
return response.content
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def print_summary(interrupted: bool = False):
|
||||
nonlocal findings
|
||||
|
||||
# Generate summary if we don't have findings yet
|
||||
if not findings and tool_log:
|
||||
summary = await generate_summary()
|
||||
if summary:
|
||||
findings.append(summary)
|
||||
|
||||
elapsed = int(time.time() - start_time)
|
||||
mins, secs = divmod(elapsed, 60)
|
||||
|
||||
title = "Interrupted" if interrupted else "Finished"
|
||||
status = "PARTIAL RESULTS" if interrupted else "COMPLETE"
|
||||
if stopped_reason:
|
||||
status = f"STOPPED ({stopped_reason})"
|
||||
|
||||
final_text = Text()
|
||||
final_text.append(f"{status}\n\n", style=f"bold {GHOST_PRIMARY}")
|
||||
final_text.append("Duration: ", style=GHOST_DIM)
|
||||
final_text.append(f"{mins}m {secs}s\n", style=GHOST_SECONDARY)
|
||||
final_text.append("Iterations: ", style=GHOST_DIM)
|
||||
final_text.append(f"{iteration}\n", style=GHOST_SECONDARY)
|
||||
final_text.append("Tools: ", style=GHOST_DIM)
|
||||
final_text.append(f"{tool_count}/{max_tools}\n", style=GHOST_SECONDARY)
|
||||
|
||||
if findings:
|
||||
final_text.append("Findings: ", style=GHOST_DIM)
|
||||
final_text.append(f"{len(findings)}", style=GHOST_SECONDARY)
|
||||
|
||||
console.print()
|
||||
console.print(
|
||||
Panel(
|
||||
final_text,
|
||||
title=f"[{GHOST_SECONDARY}]{title}",
|
||||
border_style=GHOST_BORDER,
|
||||
)
|
||||
)
|
||||
|
||||
# Show summary/findings
|
||||
if findings:
|
||||
console.print()
|
||||
console.print(
|
||||
Panel(
|
||||
Markdown(findings[-1]),
|
||||
title=f"[{GHOST_PRIMARY}]Summary",
|
||||
border_style=GHOST_BORDER,
|
||||
)
|
||||
)
|
||||
|
||||
# Save report
|
||||
save_report()
|
||||
|
||||
print_status("Initializing agent...")
|
||||
|
||||
try:
|
||||
async for response in agent.agent_loop(task_msg):
|
||||
iteration += 1
|
||||
|
||||
# Show tool calls and results as they happen
|
||||
if response.tool_calls:
|
||||
for i, call in enumerate(response.tool_calls):
|
||||
tool_count += 1
|
||||
name = getattr(call, "name", None) or getattr(
|
||||
call.function, "name", "tool"
|
||||
)
|
||||
|
||||
elapsed = int(time.time() - start_time)
|
||||
mins, secs = divmod(elapsed, 60)
|
||||
ts = f"{mins:02d}:{secs:02d}"
|
||||
|
||||
# Get result if available
|
||||
if response.tool_results and i < len(response.tool_results):
|
||||
tr = response.tool_results[i]
|
||||
result_text = tr.result or tr.error or ""
|
||||
if result_text:
|
||||
# Truncate for display
|
||||
preview = result_text[:200].replace("\n", " ")
|
||||
if len(result_text) > 200:
|
||||
preview += "..."
|
||||
|
||||
# Parse args for command extraction
|
||||
command_text = ""
|
||||
exit_code = None
|
||||
try:
|
||||
args = getattr(call, "arguments", None) or getattr(
|
||||
call.function, "arguments", "{}"
|
||||
)
|
||||
if isinstance(args, str):
|
||||
import json
|
||||
|
||||
args = json.loads(args)
|
||||
if isinstance(args, dict):
|
||||
command_text = args.get("command", "")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Extract exit code from result
|
||||
if response.tool_results and i < len(response.tool_results):
|
||||
tr = response.tool_results[i]
|
||||
full_result = tr.result or tr.error or ""
|
||||
# Try to parse exit code
|
||||
if "Exit Code:" in full_result:
|
||||
try:
|
||||
import re
|
||||
|
||||
match = re.search(r"Exit Code:\s*(\d+)", full_result)
|
||||
if match:
|
||||
exit_code = int(match.group(1))
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
full_result = ""
|
||||
|
||||
# Store full data for report (not truncated)
|
||||
tool_log.append(
|
||||
{
|
||||
"ts": ts,
|
||||
"name": name,
|
||||
"command": command_text,
|
||||
"result": full_result,
|
||||
"exit_code": exit_code,
|
||||
}
|
||||
)
|
||||
|
||||
# Metasploit-style output with better spacing
|
||||
console.print() # Blank line before each tool
|
||||
print_status(f"$ {name} ({tool_count}/{max_tools})", GHOST_ACCENT)
|
||||
|
||||
# Show command/args on separate indented line (truncated for display)
|
||||
if command_text:
|
||||
display_cmd = command_text[:80]
|
||||
if len(command_text) > 80:
|
||||
display_cmd += "..."
|
||||
console.print(f" [{GHOST_DIM}]{display_cmd}[/]")
|
||||
|
||||
# Show result on separate line with status indicator
|
||||
if response.tool_results and i < len(response.tool_results):
|
||||
tr = response.tool_results[i]
|
||||
if tr.error:
|
||||
console.print(
|
||||
f" [{GHOST_DIM}][!] {tr.error[:100]}[/]"
|
||||
)
|
||||
elif tr.result:
|
||||
# Show exit code or brief result
|
||||
result_line = tr.result[:100].replace("\n", " ")
|
||||
if exit_code == 0 or "success" in result_line.lower():
|
||||
console.print(f" [{GHOST_DIM}][+] OK[/]")
|
||||
elif exit_code is not None and exit_code != 0:
|
||||
console.print(
|
||||
f" [{GHOST_DIM}][-] Exit {exit_code}[/]"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
f" [{GHOST_DIM}][*] {result_line[:60]}...[/]"
|
||||
)
|
||||
|
||||
# Check max tools limit
|
||||
if tool_count >= max_tools:
|
||||
stopped_reason = "max calls reached"
|
||||
console.print()
|
||||
print_status(f"Max calls limit reached ({max_tools})", "yellow")
|
||||
raise StopIteration()
|
||||
|
||||
# Print assistant content immediately (analysis/findings)
|
||||
if response.content and response.content != last_content:
|
||||
last_content = response.content
|
||||
findings.append(response.content)
|
||||
|
||||
console.print()
|
||||
console.print(
|
||||
Panel(
|
||||
Markdown(response.content),
|
||||
title=f"[{GHOST_PRIMARY}]GhostCrew",
|
||||
border_style=GHOST_BORDER,
|
||||
)
|
||||
)
|
||||
console.print()
|
||||
|
||||
await print_summary(interrupted=False)
|
||||
|
||||
except StopIteration:
|
||||
await print_summary(interrupted=True)
|
||||
except (KeyboardInterrupt, asyncio.CancelledError):
|
||||
stopped_reason = "user interrupt"
|
||||
await print_summary(interrupted=True)
|
||||
except Exception as e:
|
||||
console.print(f"\n[red]Error: {e}[/]")
|
||||
stopped_reason = f"error: {e}"
|
||||
await print_summary(interrupted=True)
|
||||
|
||||
finally:
|
||||
# Cleanup MCP connections first
|
||||
if mcp_manager:
|
||||
try:
|
||||
await mcp_manager.disconnect_all()
|
||||
await asyncio.sleep(0.1) # Allow transports to close cleanly
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Then stop runtime
|
||||
if runtime:
|
||||
try:
|
||||
await runtime.stop()
|
||||
except Exception:
|
||||
pass
|
||||
286
ghostcrew/interface/main.py
Normal file
286
ghostcrew/interface/main.py
Normal file
@@ -0,0 +1,286 @@
|
||||
"""Main entry point for GhostCrew."""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
|
||||
from ..config.constants import DEFAULT_MODEL
|
||||
from .cli import run_cli
|
||||
from .tui import run_tui
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
"""Parse command line arguments."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="GhostCrew - AI Penetration Testing",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
ghostcrew Launch TUI
|
||||
ghostcrew -t 192.168.1.1 Launch TUI with target
|
||||
ghostcrew -n -t example.com Non-interactive run
|
||||
ghostcrew tools list List available tools
|
||||
ghostcrew mcp list List MCP servers
|
||||
""",
|
||||
)
|
||||
|
||||
# Subcommands
|
||||
subparsers = parser.add_subparsers(dest="command", help="Commands")
|
||||
|
||||
# Tools subcommand
|
||||
tools_parser = subparsers.add_parser("tools", help="Manage tools")
|
||||
tools_subparsers = tools_parser.add_subparsers(
|
||||
dest="tools_command", help="Tool commands"
|
||||
)
|
||||
|
||||
# tools list
|
||||
tools_subparsers.add_parser("list", help="List all available tools")
|
||||
|
||||
# tools info
|
||||
tools_info = tools_subparsers.add_parser("info", help="Show tool details")
|
||||
tools_info.add_argument("name", help="Tool name")
|
||||
|
||||
# MCP subcommand
|
||||
mcp_parser = subparsers.add_parser("mcp", help="Manage MCP servers")
|
||||
mcp_subparsers = mcp_parser.add_subparsers(dest="mcp_command", help="MCP commands")
|
||||
|
||||
# mcp list
|
||||
mcp_subparsers.add_parser("list", help="List configured MCP servers")
|
||||
|
||||
# mcp add
|
||||
mcp_add = mcp_subparsers.add_parser("add", help="Add an MCP server")
|
||||
mcp_add.add_argument("name", help="Server name")
|
||||
mcp_add.add_argument("command", help="Command to run (e.g., npx)")
|
||||
mcp_add.add_argument("args", nargs="*", help="Command arguments")
|
||||
mcp_add.add_argument("--description", "-d", default="", help="Server description")
|
||||
|
||||
# mcp remove
|
||||
mcp_remove = mcp_subparsers.add_parser("remove", help="Remove an MCP server")
|
||||
mcp_remove.add_argument("name", help="Server name to remove")
|
||||
|
||||
# mcp test
|
||||
mcp_test = mcp_subparsers.add_parser("test", help="Test MCP server connection")
|
||||
mcp_test.add_argument("name", help="Server name to test")
|
||||
|
||||
# Target option
|
||||
parser.add_argument("--target", "-t", help="Target (IP, hostname, or URL)")
|
||||
|
||||
# Non-interactive mode
|
||||
parser.add_argument(
|
||||
"-n",
|
||||
"--headless",
|
||||
action="store_true",
|
||||
help="Run without TUI (requires --target)",
|
||||
)
|
||||
|
||||
# Task for non-interactive mode
|
||||
parser.add_argument("--task", help="Task to run in non-interactive mode")
|
||||
|
||||
# Report output (saves to loot/ by default)
|
||||
parser.add_argument(
|
||||
"--report",
|
||||
"-r",
|
||||
nargs="?",
|
||||
const="auto",
|
||||
help="Generate report (default: loot/<target>_<timestamp>.md)",
|
||||
)
|
||||
|
||||
# Max tool calls limit
|
||||
parser.add_argument(
|
||||
"--max", type=int, default=50, help="Max calls before stopping (default: 50)"
|
||||
)
|
||||
|
||||
# Model options
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
default=DEFAULT_MODEL,
|
||||
help="LLM model (set GHOSTCREW_MODEL in .env)",
|
||||
)
|
||||
|
||||
# Docker mode
|
||||
parser.add_argument(
|
||||
"--docker",
|
||||
"-d",
|
||||
action="store_true",
|
||||
help="Run tools inside Docker container (requires Docker)",
|
||||
)
|
||||
|
||||
# Version
|
||||
parser.add_argument("--version", action="version", version="GhostCrew 0.2.0")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def handle_tools_command(args: argparse.Namespace):
|
||||
"""Handle tools subcommand."""
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from ..tools import get_all_tools, get_tool
|
||||
|
||||
console = Console()
|
||||
|
||||
if args.tools_command == "list":
|
||||
tools = get_all_tools()
|
||||
|
||||
if not tools:
|
||||
console.print("[yellow]No tools found[/]")
|
||||
return
|
||||
|
||||
table = Table(title="Available Tools")
|
||||
table.add_column("Name", style="cyan")
|
||||
table.add_column("Category", style="green")
|
||||
table.add_column("Description")
|
||||
|
||||
for tool in sorted(tools, key=lambda t: t.name):
|
||||
desc = (
|
||||
tool.description[:50] + "..."
|
||||
if len(tool.description) > 50
|
||||
else tool.description
|
||||
)
|
||||
table.add_row(tool.name, tool.category, desc)
|
||||
|
||||
console.print(table)
|
||||
console.print(f"\nTotal: {len(tools)} tools")
|
||||
|
||||
elif args.tools_command == "info":
|
||||
tool = get_tool(args.name)
|
||||
if not tool:
|
||||
console.print(f"[red]Tool not found: {args.name}[/]")
|
||||
return
|
||||
|
||||
console.print(f"\n[bold cyan]{tool.name}[/]")
|
||||
console.print(f"[dim]Category:[/] {tool.category}")
|
||||
console.print(f"\n{tool.description}")
|
||||
|
||||
if tool.schema.properties:
|
||||
console.print("\n[bold]Parameters:[/]")
|
||||
for name, props in tool.schema.properties.items():
|
||||
required = (
|
||||
"required" if name in (tool.schema.required or []) else "optional"
|
||||
)
|
||||
ptype = props.get("type", "any")
|
||||
desc = props.get("description", "")
|
||||
console.print(f" [cyan]{name}[/] ({ptype}, {required}): {desc}")
|
||||
|
||||
else:
|
||||
console.print("[yellow]Use 'ghostcrew tools --help' for commands[/]")
|
||||
|
||||
|
||||
def handle_mcp_command(args: argparse.Namespace):
|
||||
"""Handle MCP subcommand."""
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from ..mcp.manager import MCPManager
|
||||
|
||||
console = Console()
|
||||
manager = MCPManager()
|
||||
|
||||
if args.mcp_command == "list":
|
||||
servers = manager.list_configured_servers()
|
||||
|
||||
if not servers:
|
||||
console.print("[yellow]No MCP servers configured[/]")
|
||||
console.print(
|
||||
"\nAdd a server with: ghostcrew mcp add <name> <command> <args...>"
|
||||
)
|
||||
return
|
||||
|
||||
table = Table(title="Configured MCP Servers")
|
||||
table.add_column("Name", style="cyan")
|
||||
table.add_column("Command", style="green")
|
||||
table.add_column("Args")
|
||||
table.add_column("Connected", style="yellow")
|
||||
|
||||
for server in servers:
|
||||
args_str = " ".join(server["args"][:3])
|
||||
if len(server["args"]) > 3:
|
||||
args_str += "..."
|
||||
connected = "+" if server.get("connected") else "-"
|
||||
table.add_row(server["name"], server["command"], args_str, connected)
|
||||
|
||||
console.print(table)
|
||||
console.print(f"\nConfig file: {manager.config_path}")
|
||||
|
||||
elif args.mcp_command == "add":
|
||||
manager.add_server(
|
||||
name=args.name,
|
||||
command=args.command,
|
||||
args=args.args or [],
|
||||
description=args.description,
|
||||
)
|
||||
console.print(f"[green]Added MCP server: {args.name}[/]")
|
||||
console.print(f" Command: {args.command} {' '.join(args.args or [])}")
|
||||
|
||||
elif args.mcp_command == "remove":
|
||||
if manager.remove_server(args.name):
|
||||
console.print(f"[yellow]Removed MCP server: {args.name}[/]")
|
||||
else:
|
||||
console.print(f"[red]Server not found: {args.name}[/]")
|
||||
|
||||
elif args.mcp_command == "test":
|
||||
console.print(f"[bold]Testing MCP server: {args.name}[/]\n")
|
||||
|
||||
async def test_server():
|
||||
server = await manager.connect_server(args.name)
|
||||
if server and server.connected:
|
||||
console.print("[green]+ Connected successfully![/]")
|
||||
console.print(f"\n[bold]Available tools ({len(server.tools)}):[/]")
|
||||
for tool in server.tools:
|
||||
desc = tool.get("description", "No description")[:60]
|
||||
console.print(f" [cyan]{tool['name']}[/]: {desc}")
|
||||
await manager.disconnect_all()
|
||||
else:
|
||||
console.print("[red]x Failed to connect[/]")
|
||||
|
||||
asyncio.run(test_server())
|
||||
|
||||
else:
|
||||
console.print("[yellow]Use 'ghostcrew mcp --help' for available commands[/]")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point."""
|
||||
args = parse_arguments()
|
||||
|
||||
# Handle subcommands
|
||||
if args.command == "tools":
|
||||
handle_tools_command(args)
|
||||
return
|
||||
|
||||
if args.command == "mcp":
|
||||
handle_mcp_command(args)
|
||||
return
|
||||
|
||||
# Check model configuration
|
||||
if not args.model:
|
||||
print("Error: No model configured.")
|
||||
print("Set GHOSTCREW_MODEL in .env file or use --model flag.")
|
||||
print(
|
||||
"Example: GHOSTCREW_MODEL=gpt-5 or GHOSTCREW_MODEL=claude-sonnet-4-20250514"
|
||||
)
|
||||
return
|
||||
|
||||
# Determine interface mode
|
||||
if args.headless:
|
||||
if not args.target:
|
||||
print("Error: --target is required for headless mode")
|
||||
return
|
||||
asyncio.run(
|
||||
run_cli(
|
||||
target=args.target,
|
||||
model=args.model,
|
||||
task=args.task,
|
||||
report=args.report,
|
||||
max_tools=args.max,
|
||||
use_docker=args.docker,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# TUI doesn't need asyncio.run - it runs its own event loop
|
||||
run_tui(target=args.target, model=args.model, use_docker=args.docker)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1819
ghostcrew/interface/tui.py
Normal file
1819
ghostcrew/interface/tui.py
Normal file
File diff suppressed because it is too large
Load Diff
237
ghostcrew/interface/utils.py
Normal file
237
ghostcrew/interface/utils.py
Normal file
@@ -0,0 +1,237 @@
|
||||
"""Interface utilities for GhostCrew."""
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
# ASCII Art Banner
|
||||
ASCII_BANNER = r"""
|
||||
('-. .-. .-') .-') _ _ .-') ('-. (`\ .-') /`
|
||||
( OO ) / ( OO ). ( OO) ) ( \( -O ) _( OO) `.( OO ),'
|
||||
,----. ,--. ,--. .-'),-----. (_)---\_)/ '._ .-----. ,------. (,------.,--./ .--.
|
||||
' .-./-') | | | |( OO' .-. '/ _ | |'--...__)' .--./ | /`. ' | .---'| | |
|
||||
| |_( O- )| .| |/ | | | |\ :` `. '--. .--'| |('-. | / | | | | | | | |,
|
||||
| | .--, \| |\_) | |\| | '..`''.) | | /_) |OO )| |_.' |(| '--. | |.'.| |_)
|
||||
(| | '. (_/| .-. | \ | | | |.-._) \ | | || |`-'| | . '.' | .--' | |
|
||||
| '--' | | | | | `' '-' '\ / | | (_' '--'\ | |\ \ | `---.| ,'. |
|
||||
`------' `--' `--' `-----' `-----' `--' `-----' `--' '--' `------''--' '--'
|
||||
"""
|
||||
|
||||
|
||||
def print_banner():
|
||||
"""Print the GhostCrew banner."""
|
||||
console.print(f"[bold white]{ASCII_BANNER}[/]")
|
||||
console.print(
|
||||
"[bold white]====================== GHOSTCREW =======================[/]"
|
||||
)
|
||||
console.print(
|
||||
"[dim white] AI Penetration Testing Agents v0.2.0[/dim white]\n"
|
||||
)
|
||||
|
||||
|
||||
def format_finding(
|
||||
title: str,
|
||||
severity: str,
|
||||
target: str,
|
||||
description: str,
|
||||
evidence: str = "",
|
||||
impact: str = "",
|
||||
remediation: str = "",
|
||||
) -> Panel:
|
||||
"""
|
||||
Format a security finding for display.
|
||||
|
||||
Args:
|
||||
title: Finding title
|
||||
severity: Severity level
|
||||
target: Affected target
|
||||
description: Description of the finding
|
||||
evidence: Proof/evidence
|
||||
impact: Potential impact
|
||||
remediation: How to fix
|
||||
|
||||
Returns:
|
||||
Rich Panel with formatted finding
|
||||
"""
|
||||
severity_colors = {
|
||||
"critical": "red bold",
|
||||
"high": "red",
|
||||
"medium": "yellow",
|
||||
"low": "blue",
|
||||
"informational": "dim",
|
||||
}
|
||||
|
||||
color = severity_colors.get(severity.lower(), "white")
|
||||
|
||||
content = f"""
|
||||
[bold]Target:[/] {target}
|
||||
[{color}]Severity:[/{color}] [{color}]{severity.upper()}[/{color}]
|
||||
|
||||
[bold]Description:[/]
|
||||
{description}
|
||||
"""
|
||||
|
||||
if evidence:
|
||||
content += f"\n[bold]Evidence:[/]\n{evidence}\n"
|
||||
|
||||
if impact:
|
||||
content += f"\n[bold]Impact:[/]\n{impact}\n"
|
||||
|
||||
if remediation:
|
||||
content += f"\n[bold]Remediation:[/]\n{remediation}\n"
|
||||
|
||||
return Panel(content, title=f"[bold]{title}[/]", border_style=color)
|
||||
|
||||
|
||||
def format_tool_call(tool_call: Any) -> str:
|
||||
"""
|
||||
Format a tool call for display.
|
||||
|
||||
Args:
|
||||
tool_call: The tool call object
|
||||
|
||||
Returns:
|
||||
Formatted string
|
||||
"""
|
||||
name = tool_call.name if hasattr(tool_call, "name") else str(tool_call)
|
||||
args = tool_call.arguments if hasattr(tool_call, "arguments") else {}
|
||||
|
||||
# Truncate long arguments
|
||||
args_str = str(args)
|
||||
if len(args_str) > 100:
|
||||
args_str = args_str[:100] + "..."
|
||||
|
||||
return f"[bold yellow]⚡ Tool:[/] {name}\n[dim]{args_str}[/dim]"
|
||||
|
||||
|
||||
def print_status(
|
||||
target: Optional[str] = None,
|
||||
scope: Optional[list] = None,
|
||||
agent_state: str = "idle",
|
||||
tools_count: int = 0,
|
||||
findings_count: int = 0,
|
||||
):
|
||||
"""
|
||||
Print current status information.
|
||||
|
||||
Args:
|
||||
target: Current target
|
||||
scope: Current scope
|
||||
agent_state: Agent state
|
||||
tools_count: Number of loaded tools
|
||||
findings_count: Number of findings
|
||||
"""
|
||||
table = Table(title="GhostCrew Status", show_header=False)
|
||||
table.add_column("Property", style="cyan")
|
||||
table.add_column("Value", style="white")
|
||||
|
||||
table.add_row("Target", target or "Not set")
|
||||
table.add_row("Scope", ", ".join(scope) if scope else "Not set")
|
||||
table.add_row("Agent State", agent_state)
|
||||
table.add_row("Tools Loaded", str(tools_count))
|
||||
table.add_row("Findings", str(findings_count))
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
def format_scan_progress(current: int, total: int, current_item: str) -> str:
|
||||
"""
|
||||
Format scan progress for display.
|
||||
|
||||
Args:
|
||||
current: Current item number
|
||||
total: Total items
|
||||
current_item: Current item being scanned
|
||||
|
||||
Returns:
|
||||
Formatted progress string
|
||||
"""
|
||||
percentage = (current / total * 100) if total > 0 else 0
|
||||
bar_width = 30
|
||||
filled = int(bar_width * current / total) if total > 0 else 0
|
||||
bar = "█" * filled + "░" * (bar_width - filled)
|
||||
|
||||
return f"[{bar}] {percentage:.1f}% ({current}/{total}) - {current_item}"
|
||||
|
||||
|
||||
def truncate_output(output: str, max_lines: int = 50) -> str:
|
||||
"""
|
||||
Truncate long output for display.
|
||||
|
||||
Args:
|
||||
output: Output to truncate
|
||||
max_lines: Maximum number of lines
|
||||
|
||||
Returns:
|
||||
Truncated output
|
||||
"""
|
||||
lines = output.split("\n")
|
||||
|
||||
if len(lines) <= max_lines:
|
||||
return output
|
||||
|
||||
half = max_lines // 2
|
||||
truncated = (
|
||||
lines[:half]
|
||||
+ [f"\n... ({len(lines) - max_lines} lines omitted) ...\n"]
|
||||
+ lines[-half:]
|
||||
)
|
||||
|
||||
return "\n".join(truncated)
|
||||
|
||||
|
||||
def colorize_severity(severity: str) -> str:
|
||||
"""
|
||||
Add color to severity text.
|
||||
|
||||
Args:
|
||||
severity: Severity level
|
||||
|
||||
Returns:
|
||||
Colorized severity string
|
||||
"""
|
||||
colors = {
|
||||
"critical": "[red bold]CRITICAL[/]",
|
||||
"high": "[red]HIGH[/]",
|
||||
"medium": "[yellow]MEDIUM[/]",
|
||||
"low": "[blue]LOW[/]",
|
||||
"informational": "[dim]INFO[/]",
|
||||
"info": "[dim]INFO[/]",
|
||||
}
|
||||
|
||||
return colors.get(severity.lower(), severity)
|
||||
|
||||
|
||||
def format_command_output(
|
||||
command: str, exit_code: int, stdout: str, stderr: str
|
||||
) -> Panel:
|
||||
"""
|
||||
Format command output for display.
|
||||
|
||||
Args:
|
||||
command: The command that was run
|
||||
exit_code: Exit code
|
||||
stdout: Standard output
|
||||
stderr: Standard error
|
||||
|
||||
Returns:
|
||||
Rich Panel with formatted output
|
||||
"""
|
||||
success = exit_code == 0
|
||||
border_color = "green" if success else "red"
|
||||
|
||||
content = f"[bold]Command:[/] {command}\n"
|
||||
content += f"[bold]Exit Code:[/] {exit_code}\n"
|
||||
|
||||
if stdout:
|
||||
content += f"\n[bold]Output:[/]\n{truncate_output(stdout)}"
|
||||
|
||||
if stderr:
|
||||
content += f"\n[bold red]Errors:[/]\n{truncate_output(stderr)}"
|
||||
|
||||
return Panel(content, title="Command Result", border_style=border_color)
|
||||
13
ghostcrew/knowledge/__init__.py
Normal file
13
ghostcrew/knowledge/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
"""Knowledge and RAG system for GhostCrew."""
|
||||
|
||||
from .embeddings import get_embeddings, get_embeddings_local
|
||||
from .indexer import KnowledgeIndexer
|
||||
from .rag import Document, RAGEngine
|
||||
|
||||
__all__ = [
|
||||
"RAGEngine",
|
||||
"Document",
|
||||
"get_embeddings",
|
||||
"get_embeddings_local",
|
||||
"KnowledgeIndexer",
|
||||
]
|
||||
146
ghostcrew/knowledge/embeddings.py
Normal file
146
ghostcrew/knowledge/embeddings.py
Normal file
@@ -0,0 +1,146 @@
|
||||
"""Embedding generation for GhostCrew."""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_embeddings(
|
||||
texts: List[str], model: str = "text-embedding-3-small"
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Generate embeddings for a list of texts using LiteLLM.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed
|
||||
model: The embedding model to use
|
||||
|
||||
Returns:
|
||||
NumPy array of embeddings
|
||||
"""
|
||||
try:
|
||||
import litellm
|
||||
|
||||
response = litellm.embedding(model=model, input=texts)
|
||||
|
||||
embeddings = [item["embedding"] for item in response.data]
|
||||
return np.array(embeddings, dtype=np.float32)
|
||||
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"litellm is required for embeddings. Install with: pip install litellm"
|
||||
) from e
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to generate embeddings: {e}") from e
|
||||
|
||||
|
||||
def get_embeddings_local(
|
||||
texts: List[str], model: str = "all-MiniLM-L6-v2"
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Generate embeddings locally using sentence-transformers.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed
|
||||
model: The sentence-transformer model to use
|
||||
|
||||
Returns:
|
||||
NumPy array of embeddings
|
||||
"""
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
encoder = SentenceTransformer(model)
|
||||
embeddings = encoder.encode(texts, show_progress_bar=False)
|
||||
return np.array(embeddings, dtype=np.float32)
|
||||
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"sentence-transformers is required for local embeddings. "
|
||||
"Install with: pip install sentence-transformers"
|
||||
) from e
|
||||
|
||||
|
||||
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
|
||||
"""
|
||||
Compute cosine similarity between two vectors.
|
||||
|
||||
Args:
|
||||
a: First vector
|
||||
b: Second vector
|
||||
|
||||
Returns:
|
||||
Cosine similarity score
|
||||
"""
|
||||
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-10)
|
||||
|
||||
|
||||
def batch_cosine_similarity(query: np.ndarray, embeddings: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Compute cosine similarity between a query and multiple embeddings.
|
||||
|
||||
Args:
|
||||
query: Query vector
|
||||
embeddings: Matrix of embeddings
|
||||
|
||||
Returns:
|
||||
Array of similarity scores
|
||||
"""
|
||||
query_norm = np.linalg.norm(query)
|
||||
embeddings_norm = np.linalg.norm(embeddings, axis=1)
|
||||
|
||||
return np.dot(embeddings, query) / (embeddings_norm * query_norm + 1e-10)
|
||||
|
||||
|
||||
class EmbeddingCache:
|
||||
"""Cache for embeddings to avoid recomputation."""
|
||||
|
||||
def __init__(self, max_size: int = 1000):
|
||||
"""
|
||||
Initialize the embedding cache.
|
||||
|
||||
Args:
|
||||
max_size: Maximum number of embeddings to cache
|
||||
"""
|
||||
self.max_size = max_size
|
||||
self._cache: dict[str, np.ndarray] = {}
|
||||
self._order: list[str] = []
|
||||
|
||||
def get(self, text: str) -> Optional[np.ndarray]:
|
||||
"""
|
||||
Get a cached embedding.
|
||||
|
||||
Args:
|
||||
text: The text to look up
|
||||
|
||||
Returns:
|
||||
The cached embedding or None
|
||||
"""
|
||||
return self._cache.get(text)
|
||||
|
||||
def set(self, text: str, embedding: np.ndarray):
|
||||
"""
|
||||
Cache an embedding.
|
||||
|
||||
Args:
|
||||
text: The text key
|
||||
embedding: The embedding to cache
|
||||
"""
|
||||
if text in self._cache:
|
||||
return
|
||||
|
||||
if len(self._cache) >= self.max_size:
|
||||
# Remove oldest entry
|
||||
oldest = self._order.pop(0)
|
||||
del self._cache[oldest]
|
||||
|
||||
self._cache[text] = embedding
|
||||
self._order.append(text)
|
||||
|
||||
def clear(self):
|
||||
"""Clear the cache."""
|
||||
self._cache.clear()
|
||||
self._order.clear()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._cache)
|
||||
249
ghostcrew/knowledge/indexer.py
Normal file
249
ghostcrew/knowledge/indexer.py
Normal file
@@ -0,0 +1,249 @@
|
||||
"""Knowledge indexer for GhostCrew."""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, List
|
||||
|
||||
from .rag import Document
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexingResult:
|
||||
"""Result of an indexing operation."""
|
||||
|
||||
total_files: int
|
||||
indexed_files: int
|
||||
total_chunks: int
|
||||
errors: List[str]
|
||||
|
||||
|
||||
class KnowledgeIndexer:
|
||||
"""Indexes knowledge sources for the RAG engine."""
|
||||
|
||||
# Supported file extensions
|
||||
TEXT_EXTENSIONS = [".txt", ".md", ".rst"]
|
||||
DATA_EXTENSIONS = [".json", ".yaml", ".yml"]
|
||||
|
||||
def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):
|
||||
"""
|
||||
Initialize the knowledge indexer.
|
||||
|
||||
Args:
|
||||
chunk_size: Maximum chunk size in characters
|
||||
chunk_overlap: Overlap between chunks
|
||||
"""
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_overlap = chunk_overlap
|
||||
|
||||
def index_directory(self, directory: Path) -> tuple[List[Document], IndexingResult]:
|
||||
"""
|
||||
Index all supported files in a directory.
|
||||
|
||||
Args:
|
||||
directory: The directory to index
|
||||
|
||||
Returns:
|
||||
Tuple of (documents, indexing_result)
|
||||
"""
|
||||
documents = []
|
||||
errors = []
|
||||
total_files = 0
|
||||
indexed_files = 0
|
||||
|
||||
if not directory.exists():
|
||||
return documents, IndexingResult(
|
||||
0, 0, 0, [f"Directory not found: {directory}"]
|
||||
)
|
||||
|
||||
for file_path in directory.rglob("*"):
|
||||
if not file_path.is_file():
|
||||
continue
|
||||
|
||||
total_files += 1
|
||||
|
||||
try:
|
||||
file_docs = self.index_file(file_path)
|
||||
if file_docs:
|
||||
documents.extend(file_docs)
|
||||
indexed_files += 1
|
||||
except Exception as e:
|
||||
errors.append(f"Error indexing {file_path}: {e}")
|
||||
|
||||
result = IndexingResult(
|
||||
total_files=total_files,
|
||||
indexed_files=indexed_files,
|
||||
total_chunks=len(documents),
|
||||
errors=errors,
|
||||
)
|
||||
|
||||
return documents, result
|
||||
|
||||
def index_file(self, file_path: Path) -> List[Document]:
|
||||
"""
|
||||
Index a single file.
|
||||
|
||||
Args:
|
||||
file_path: The file to index
|
||||
|
||||
Returns:
|
||||
List of Document objects
|
||||
"""
|
||||
suffix = file_path.suffix.lower()
|
||||
|
||||
if suffix in self.TEXT_EXTENSIONS:
|
||||
return self._index_text_file(file_path)
|
||||
elif suffix in self.DATA_EXTENSIONS:
|
||||
return self._index_data_file(file_path)
|
||||
else:
|
||||
return []
|
||||
|
||||
def _index_text_file(self, file_path: Path) -> List[Document]:
|
||||
"""Index a text file."""
|
||||
content = file_path.read_text(encoding="utf-8", errors="ignore")
|
||||
return self._chunk_text(content, str(file_path))
|
||||
|
||||
def _index_data_file(self, file_path: Path) -> List[Document]:
|
||||
"""Index a JSON/YAML file."""
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
|
||||
if file_path.suffix == ".json":
|
||||
data = json.loads(content)
|
||||
else:
|
||||
try:
|
||||
import yaml
|
||||
|
||||
data = yaml.safe_load(content)
|
||||
except ImportError:
|
||||
return []
|
||||
|
||||
return self._process_data(data, str(file_path))
|
||||
|
||||
def _chunk_text(self, text: str, source: str) -> List[Document]:
|
||||
"""Split text into chunks."""
|
||||
chunks = []
|
||||
|
||||
# Try to split by sections (headers in markdown)
|
||||
sections = self._split_by_sections(text)
|
||||
|
||||
for section in sections:
|
||||
if len(section) <= self.chunk_size:
|
||||
if section.strip():
|
||||
chunks.append(Document(content=section.strip(), source=source))
|
||||
else:
|
||||
# Further split large sections
|
||||
sub_chunks = self._split_by_paragraphs(section)
|
||||
for sub in sub_chunks:
|
||||
if sub.strip():
|
||||
chunks.append(Document(content=sub.strip(), source=source))
|
||||
|
||||
return chunks
|
||||
|
||||
def _split_by_sections(self, text: str) -> List[str]:
|
||||
"""Split text by markdown headers."""
|
||||
import re
|
||||
|
||||
# Split by headers (# Header)
|
||||
sections = re.split(r"\n(?=#{1,3}\s)", text)
|
||||
|
||||
if len(sections) == 1:
|
||||
# No headers found, return original
|
||||
return [text]
|
||||
|
||||
return sections
|
||||
|
||||
def _split_by_paragraphs(self, text: str) -> List[str]:
|
||||
"""Split text by paragraphs with overlap."""
|
||||
paragraphs = text.split("\n\n")
|
||||
chunks = []
|
||||
current_chunk = ""
|
||||
|
||||
for para in paragraphs:
|
||||
if len(current_chunk) + len(para) + 2 <= self.chunk_size:
|
||||
current_chunk += para + "\n\n"
|
||||
else:
|
||||
if current_chunk.strip():
|
||||
chunks.append(current_chunk.strip())
|
||||
# Start new chunk with overlap
|
||||
current_chunk = para + "\n\n"
|
||||
|
||||
if current_chunk.strip():
|
||||
chunks.append(current_chunk.strip())
|
||||
|
||||
return chunks
|
||||
|
||||
def _process_data(self, data: Any, source: str) -> List[Document]:
|
||||
"""Process JSON/YAML data into documents."""
|
||||
documents = []
|
||||
|
||||
if isinstance(data, list):
|
||||
for i, item in enumerate(data):
|
||||
doc = Document(
|
||||
content=json.dumps(item, indent=2),
|
||||
source=source,
|
||||
metadata={"index": i, "type": "array_item"},
|
||||
)
|
||||
documents.append(doc)
|
||||
|
||||
elif isinstance(data, dict):
|
||||
# Check if it has a specific structure
|
||||
if "entries" in data or "items" in data or "data" in data:
|
||||
items = data.get("entries") or data.get("items") or data.get("data")
|
||||
if isinstance(items, list):
|
||||
for i, item in enumerate(items):
|
||||
doc = Document(
|
||||
content=json.dumps(item, indent=2),
|
||||
source=source,
|
||||
metadata={"index": i, "type": "data_item"},
|
||||
)
|
||||
documents.append(doc)
|
||||
else:
|
||||
doc = Document(
|
||||
content=json.dumps(data, indent=2),
|
||||
source=source,
|
||||
metadata={"type": "object"},
|
||||
)
|
||||
documents.append(doc)
|
||||
else:
|
||||
doc = Document(
|
||||
content=json.dumps(data, indent=2),
|
||||
source=source,
|
||||
metadata={"type": "object"},
|
||||
)
|
||||
documents.append(doc)
|
||||
|
||||
else:
|
||||
doc = Document(
|
||||
content=str(data), source=source, metadata={"type": "primitive"}
|
||||
)
|
||||
documents.append(doc)
|
||||
|
||||
return documents
|
||||
|
||||
def create_knowledge_structure(self, base_path: Path):
|
||||
"""
|
||||
Create the default knowledge directory structure.
|
||||
|
||||
Args:
|
||||
base_path: Base path for knowledge directory
|
||||
"""
|
||||
directories = [
|
||||
base_path / "cves",
|
||||
base_path / "wordlists",
|
||||
base_path / "exploits",
|
||||
base_path / "methodologies",
|
||||
base_path / "custom",
|
||||
]
|
||||
|
||||
for directory in directories:
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Create placeholder files
|
||||
(base_path / "methodologies" / "README.md").write_text(
|
||||
"# Penetration Testing Methodologies\n\n"
|
||||
"Add methodology documents here.\n"
|
||||
)
|
||||
|
||||
(base_path / "wordlists" / "common.txt").write_text(
|
||||
"# Common wordlist\n" "admin\n" "password\n" "root\n" "user\n"
|
||||
)
|
||||
439
ghostcrew/knowledge/rag.py
Normal file
439
ghostcrew/knowledge/rag.py
Normal file
@@ -0,0 +1,439 @@
|
||||
"""RAG (Retrieval Augmented Generation) engine for GhostCrew."""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .embeddings import get_embeddings
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""A chunk of knowledge."""
|
||||
|
||||
content: str
|
||||
source: str
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
embedding: Optional[np.ndarray] = None
|
||||
doc_id: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.metadata is None:
|
||||
self.metadata = {}
|
||||
if self.doc_id is None:
|
||||
self.doc_id = f"{hash(self.content)}_{hash(self.source)}"
|
||||
|
||||
|
||||
class RAGEngine:
|
||||
"""Vector search over security knowledge."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
knowledge_path: Path = Path("knowledge"),
|
||||
embedding_model: str = "text-embedding-3-small",
|
||||
use_local_embeddings: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize the RAG engine.
|
||||
|
||||
Args:
|
||||
knowledge_path: Path to the knowledge directory
|
||||
embedding_model: Model to use for embeddings
|
||||
use_local_embeddings: Whether to use local embeddings (sentence-transformers)
|
||||
"""
|
||||
self.knowledge_path = knowledge_path
|
||||
self.embedding_model = embedding_model
|
||||
self.use_local_embeddings = use_local_embeddings
|
||||
self.documents: List[Document] = []
|
||||
self.embeddings: Optional[np.ndarray] = None
|
||||
self._indexed = False
|
||||
self._source_files: set = set() # Track unique source files
|
||||
|
||||
def index(self, force: bool = False):
|
||||
"""
|
||||
Index all documents in knowledge directory.
|
||||
|
||||
Args:
|
||||
force: Force re-indexing even if already indexed
|
||||
"""
|
||||
if self._indexed and not force:
|
||||
return
|
||||
|
||||
chunks = []
|
||||
self._source_files = set() # Reset source file tracking
|
||||
|
||||
# Process all files in knowledge directory
|
||||
if self.knowledge_path.exists():
|
||||
for file in self.knowledge_path.rglob("*"):
|
||||
if not file.is_file():
|
||||
continue
|
||||
|
||||
try:
|
||||
if file.suffix in [".txt", ".md"]:
|
||||
self._source_files.add(str(file))
|
||||
content = file.read_text(encoding="utf-8", errors="ignore")
|
||||
file_chunks = self._chunk_text(content, source=str(file))
|
||||
chunks.extend(file_chunks)
|
||||
|
||||
elif file.suffix == ".json":
|
||||
self._source_files.add(str(file))
|
||||
data = json.loads(file.read_text(encoding="utf-8"))
|
||||
if isinstance(data, list):
|
||||
for item in data:
|
||||
chunks.append(
|
||||
Document(
|
||||
content=json.dumps(item, indent=2),
|
||||
source=str(file),
|
||||
metadata=(
|
||||
item
|
||||
if isinstance(item, dict)
|
||||
else {"data": item}
|
||||
),
|
||||
)
|
||||
)
|
||||
else:
|
||||
chunks.append(
|
||||
Document(
|
||||
content=json.dumps(data, indent=2),
|
||||
source=str(file),
|
||||
metadata=(
|
||||
data
|
||||
if isinstance(data, dict)
|
||||
else {"data": data}
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[RAG] Error processing {file}: {e}")
|
||||
|
||||
self.documents = chunks
|
||||
|
||||
# Generate embeddings
|
||||
if chunks:
|
||||
texts = [doc.content for doc in chunks]
|
||||
|
||||
if self.use_local_embeddings:
|
||||
from .embeddings import get_embeddings_local
|
||||
|
||||
self.embeddings = get_embeddings_local(texts)
|
||||
else:
|
||||
self.embeddings = get_embeddings(texts, model=self.embedding_model)
|
||||
|
||||
# Store embeddings in documents
|
||||
for i, doc in enumerate(self.documents):
|
||||
doc.embedding = self.embeddings[i]
|
||||
|
||||
self._indexed = True
|
||||
|
||||
def _chunk_text(
|
||||
self, text: str, source: str, chunk_size: int = 1000, overlap: int = 200
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Split text into overlapping chunks.
|
||||
|
||||
Args:
|
||||
text: The text to split
|
||||
source: The source file path
|
||||
chunk_size: Maximum chunk size in characters
|
||||
overlap: Overlap between chunks
|
||||
|
||||
Returns:
|
||||
List of Document objects
|
||||
"""
|
||||
chunks = []
|
||||
|
||||
# Split by paragraphs first for better context
|
||||
paragraphs = text.split("\n\n")
|
||||
current_chunk = ""
|
||||
|
||||
for para in paragraphs:
|
||||
if len(current_chunk) + len(para) + 2 <= chunk_size:
|
||||
current_chunk += para + "\n\n"
|
||||
else:
|
||||
if current_chunk.strip():
|
||||
chunks.append(
|
||||
Document(content=current_chunk.strip(), source=source)
|
||||
)
|
||||
current_chunk = para + "\n\n"
|
||||
|
||||
# Add the last chunk
|
||||
if current_chunk.strip():
|
||||
chunks.append(Document(content=current_chunk.strip(), source=source))
|
||||
|
||||
# If no paragraphs were found, fall back to simple chunking
|
||||
if not chunks and text.strip():
|
||||
start = 0
|
||||
while start < len(text):
|
||||
end = start + chunk_size
|
||||
chunk = text[start:end]
|
||||
|
||||
if chunk.strip():
|
||||
chunks.append(Document(content=chunk.strip(), source=source))
|
||||
|
||||
start = end - overlap
|
||||
|
||||
return chunks
|
||||
|
||||
def search(
|
||||
self, query: str, k: int = 5, threshold: float = 0.35, max_tokens: int = 1500
|
||||
) -> List[str]:
|
||||
"""
|
||||
Find relevant documents for a query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
k: Maximum number of results to return
|
||||
threshold: Minimum similarity threshold
|
||||
max_tokens: Maximum total tokens to return (prevents context bloat)
|
||||
|
||||
Returns:
|
||||
List of relevant document contents
|
||||
"""
|
||||
# Guard against empty/invalid queries
|
||||
if not query or not isinstance(query, str) or not query.strip():
|
||||
return []
|
||||
|
||||
if not self._indexed:
|
||||
self.index()
|
||||
|
||||
if not self.documents or self.embeddings is None:
|
||||
return []
|
||||
|
||||
# Get query embedding
|
||||
if self.use_local_embeddings:
|
||||
from .embeddings import get_embeddings_local
|
||||
|
||||
query_embedding = get_embeddings_local([query])[0]
|
||||
else:
|
||||
query_embedding = get_embeddings([query], model=self.embedding_model)[0]
|
||||
|
||||
# Compute cosine similarities
|
||||
similarities = np.dot(self.embeddings, query_embedding) / (
|
||||
np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_embedding)
|
||||
+ 1e-10
|
||||
)
|
||||
|
||||
# Get top k indices above threshold
|
||||
indices_above_threshold = np.where(similarities >= threshold)[0]
|
||||
|
||||
if len(indices_above_threshold) > 0:
|
||||
# Sort by similarity (descending) and take top k
|
||||
sorted_indices = indices_above_threshold[
|
||||
np.argsort(similarities[indices_above_threshold])[::-1]
|
||||
]
|
||||
top_indices = sorted_indices[:k]
|
||||
else:
|
||||
# No results above threshold - return empty rather than irrelevant content
|
||||
return []
|
||||
|
||||
# Collect results up to max_tokens budget
|
||||
results = []
|
||||
total_tokens = 0
|
||||
for idx in top_indices:
|
||||
content = self.documents[idx].content
|
||||
# Rough token estimate: ~4 chars per token
|
||||
chunk_tokens = len(content) // 4
|
||||
if total_tokens + chunk_tokens > max_tokens and results:
|
||||
# Stop if we'd exceed budget (but always include at least one)
|
||||
break
|
||||
results.append(content)
|
||||
total_tokens += chunk_tokens
|
||||
|
||||
return results
|
||||
|
||||
def search_with_scores(
|
||||
self, query: str, k: int = 5, threshold: float = 0.35
|
||||
) -> List[tuple[Document, float]]:
|
||||
"""
|
||||
Search with similarity scores.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
k: Maximum number of results to return
|
||||
threshold: Minimum similarity threshold
|
||||
|
||||
Returns:
|
||||
List of (Document, score) tuples above threshold
|
||||
"""
|
||||
if not self._indexed:
|
||||
self.index()
|
||||
|
||||
if not self.documents or self.embeddings is None:
|
||||
return []
|
||||
|
||||
# Get query embedding
|
||||
if self.use_local_embeddings:
|
||||
from .embeddings import get_embeddings_local
|
||||
|
||||
query_embedding = get_embeddings_local([query])[0]
|
||||
else:
|
||||
query_embedding = get_embeddings([query], model=self.embedding_model)[0]
|
||||
|
||||
# Compute cosine similarities
|
||||
similarities = np.dot(self.embeddings, query_embedding) / (
|
||||
np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_embedding)
|
||||
+ 1e-10
|
||||
)
|
||||
|
||||
# Get top k above threshold
|
||||
indices_above_threshold = np.where(similarities >= threshold)[0]
|
||||
|
||||
if len(indices_above_threshold) > 0:
|
||||
sorted_indices = indices_above_threshold[
|
||||
np.argsort(similarities[indices_above_threshold])[::-1]
|
||||
]
|
||||
top_indices = sorted_indices[:k]
|
||||
else:
|
||||
# Fallback: return single best result even if below threshold
|
||||
top_indices = [np.argmax(similarities)]
|
||||
|
||||
return [(self.documents[i], float(similarities[i])) for i in top_indices]
|
||||
|
||||
def add_document(
|
||||
self, content: str, source: str = "user", metadata: Optional[dict] = None
|
||||
):
|
||||
"""
|
||||
Add a document to the knowledge base.
|
||||
|
||||
Args:
|
||||
content: The document content
|
||||
source: The source identifier
|
||||
metadata: Optional metadata
|
||||
"""
|
||||
doc = Document(content=content, source=source, metadata=metadata)
|
||||
|
||||
# Generate embedding
|
||||
if self.use_local_embeddings:
|
||||
from .embeddings import get_embeddings_local
|
||||
|
||||
new_embedding = get_embeddings_local([content])
|
||||
else:
|
||||
new_embedding = get_embeddings([content], model=self.embedding_model)
|
||||
|
||||
doc.embedding = new_embedding[0]
|
||||
self.documents.append(doc)
|
||||
|
||||
# Update embeddings array
|
||||
if self.embeddings is not None:
|
||||
self.embeddings = np.vstack([self.embeddings, new_embedding])
|
||||
else:
|
||||
self.embeddings = new_embedding
|
||||
|
||||
def add_documents(self, documents: List[Document]):
|
||||
"""
|
||||
Add multiple documents to the knowledge base.
|
||||
|
||||
Args:
|
||||
documents: List of Document objects to add
|
||||
"""
|
||||
if not documents:
|
||||
return
|
||||
|
||||
texts = [doc.content for doc in documents]
|
||||
|
||||
if self.use_local_embeddings:
|
||||
from .embeddings import get_embeddings_local
|
||||
|
||||
new_embeddings = get_embeddings_local(texts)
|
||||
else:
|
||||
new_embeddings = get_embeddings(texts, model=self.embedding_model)
|
||||
|
||||
for i, doc in enumerate(documents):
|
||||
doc.embedding = new_embeddings[i]
|
||||
self.documents.append(doc)
|
||||
|
||||
if self.embeddings is not None:
|
||||
self.embeddings = np.vstack([self.embeddings, new_embeddings])
|
||||
else:
|
||||
self.embeddings = new_embeddings
|
||||
|
||||
def remove_document(self, doc_id: str) -> bool:
|
||||
"""
|
||||
Remove a document by ID.
|
||||
|
||||
Args:
|
||||
doc_id: The document ID to remove
|
||||
|
||||
Returns:
|
||||
True if removed, False if not found
|
||||
"""
|
||||
for i, doc in enumerate(self.documents):
|
||||
if doc.doc_id == doc_id:
|
||||
self.documents.pop(i)
|
||||
if self.embeddings is not None:
|
||||
self.embeddings = np.delete(self.embeddings, i, axis=0)
|
||||
return True
|
||||
return False
|
||||
|
||||
def clear(self):
|
||||
"""Clear all documents and embeddings."""
|
||||
self.documents.clear()
|
||||
self.embeddings = None
|
||||
self._indexed = False
|
||||
self._source_files = set()
|
||||
|
||||
def get_document_count(self) -> int:
|
||||
"""Get the number of source files indexed."""
|
||||
return len(self._source_files)
|
||||
|
||||
def get_chunk_count(self) -> int:
|
||||
"""Get the number of indexed chunks (internal document segments)."""
|
||||
return len(self.documents)
|
||||
|
||||
def save_index(self, path: Path):
|
||||
"""
|
||||
Save the index to disk.
|
||||
|
||||
Args:
|
||||
path: Path to save the index
|
||||
"""
|
||||
import pickle
|
||||
|
||||
data = {
|
||||
"documents": [
|
||||
{
|
||||
"content": doc.content,
|
||||
"source": doc.source,
|
||||
"metadata": doc.metadata,
|
||||
"doc_id": doc.doc_id,
|
||||
}
|
||||
for doc in self.documents
|
||||
],
|
||||
"embeddings": self.embeddings,
|
||||
}
|
||||
|
||||
with open(path, "wb") as f:
|
||||
pickle.dump(data, f)
|
||||
|
||||
def load_index(self, path: Path):
|
||||
"""
|
||||
Load the index from disk.
|
||||
|
||||
Args:
|
||||
path: Path to load the index from
|
||||
"""
|
||||
import pickle
|
||||
|
||||
with open(path, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
|
||||
self.documents = [
|
||||
Document(
|
||||
content=d["content"],
|
||||
source=d["source"],
|
||||
metadata=d["metadata"],
|
||||
doc_id=d["doc_id"],
|
||||
)
|
||||
for d in data["documents"]
|
||||
]
|
||||
self.embeddings = data["embeddings"]
|
||||
|
||||
# Restore embeddings in documents
|
||||
if self.embeddings is not None:
|
||||
for i, doc in enumerate(self.documents):
|
||||
doc.embedding = self.embeddings[i]
|
||||
|
||||
self._indexed = True
|
||||
1
ghostcrew/knowledge/sources/cves.json
Normal file
1
ghostcrew/knowledge/sources/cves.json
Normal file
@@ -0,0 +1 @@
|
||||
[{"id":"CVE-2025-55182","description":"React4Shell - RCE via unsafe deserialization in React Server Components","cvss":10.0,"affected":"react-server-dom-* 19.0.0-19.2.0","patched":"19.0.1, 19.1.2, 19.2.1"},{"id":"CVE-2025-66478","description":"Next.js RCE via React Server Components vulnerability","cvss":10.0,"affected":"Next.js using vulnerable React RSC"}]
|
||||
5
ghostcrew/knowledge/sources/methodologies.md
Normal file
5
ghostcrew/knowledge/sources/methodologies.md
Normal file
@@ -0,0 +1,5 @@
|
||||
# Example Knowledge
|
||||
|
||||
Add .md, .txt, or .json files here for RAG context.
|
||||
|
||||
Delete these examples when adding real content.
|
||||
4
ghostcrew/knowledge/sources/wordlists.txt
Normal file
4
ghostcrew/knowledge/sources/wordlists.txt
Normal file
@@ -0,0 +1,4 @@
|
||||
admin
|
||||
root
|
||||
password
|
||||
test
|
||||
15
ghostcrew/llm/__init__.py
Normal file
15
ghostcrew/llm/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
"""LLM integration for GhostCrew."""
|
||||
|
||||
from .config import ModelConfig
|
||||
from .llm import LLM, LLMResponse
|
||||
from .memory import ConversationMemory
|
||||
from .utils import count_tokens, truncate_to_tokens
|
||||
|
||||
__all__ = [
|
||||
"LLM",
|
||||
"LLMResponse",
|
||||
"ModelConfig",
|
||||
"ConversationMemory",
|
||||
"count_tokens",
|
||||
"truncate_to_tokens",
|
||||
]
|
||||
54
ghostcrew/llm/config.py
Normal file
54
ghostcrew/llm/config.py
Normal file
@@ -0,0 +1,54 @@
|
||||
"""LLM configuration for GhostCrew."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""LLM model configuration."""
|
||||
|
||||
# Generation parameters
|
||||
temperature: float = 0.7
|
||||
max_tokens: int = 4096
|
||||
top_p: float = 1.0
|
||||
frequency_penalty: float = 0.0
|
||||
presence_penalty: float = 0.0
|
||||
|
||||
# Context management
|
||||
max_context_tokens: int = 128000
|
||||
|
||||
# Retry settings for rate limits
|
||||
max_retries: int = 5 # Retry up to 5 times for rate limits
|
||||
retry_delay: float = 2.0 # Base delay - will exponentially increase
|
||||
|
||||
# Timeout
|
||||
timeout: int = 120
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for LLM calls."""
|
||||
return {
|
||||
"temperature": self.temperature,
|
||||
"max_tokens": self.max_tokens,
|
||||
"top_p": self.top_p,
|
||||
"frequency_penalty": self.frequency_penalty,
|
||||
"presence_penalty": self.presence_penalty,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def for_model(cls, model: str) -> "ModelConfig":
|
||||
"""Get configuration for a model. Uses sensible defaults for modern LLMs."""
|
||||
return cls(temperature=0.7, max_tokens=4096, max_context_tokens=128000)
|
||||
|
||||
|
||||
# Preset configurations
|
||||
CREATIVE_CONFIG = ModelConfig(
|
||||
temperature=0.9, top_p=0.95, frequency_penalty=0.5, presence_penalty=0.5
|
||||
)
|
||||
|
||||
PRECISE_CONFIG = ModelConfig(
|
||||
temperature=0.1, top_p=1.0, frequency_penalty=0.0, presence_penalty=0.0
|
||||
)
|
||||
|
||||
BALANCED_CONFIG = ModelConfig(
|
||||
temperature=0.7, top_p=1.0, frequency_penalty=0.0, presence_penalty=0.0
|
||||
)
|
||||
325
ghostcrew/llm/llm.py
Normal file
325
ghostcrew/llm/llm.py
Normal file
@@ -0,0 +1,325 @@
|
||||
"""LiteLLM wrapper for GhostCrew."""
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, List, Optional
|
||||
|
||||
from ..config.constants import DEFAULT_MODEL
|
||||
from .config import ModelConfig
|
||||
from .memory import ConversationMemory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..knowledge import RAGEngine
|
||||
from ..tools import Tool
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponse:
|
||||
"""Response from LLM."""
|
||||
|
||||
content: Optional[str]
|
||||
tool_calls: Optional[List[Any]]
|
||||
usage: Optional[dict]
|
||||
model: str = ""
|
||||
finish_reason: str = ""
|
||||
|
||||
|
||||
class LLM:
|
||||
"""LiteLLM wrapper with tool calling support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = None,
|
||||
config: Optional[ModelConfig] = None,
|
||||
rag_engine: Optional["RAGEngine"] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the LLM wrapper.
|
||||
|
||||
Args:
|
||||
model: The model to use (supports LiteLLM model names)
|
||||
config: Model configuration
|
||||
rag_engine: Optional RAG engine for context injection
|
||||
"""
|
||||
self.model = model or DEFAULT_MODEL
|
||||
self.config = config or ModelConfig()
|
||||
self.rag_engine = rag_engine
|
||||
self.memory = ConversationMemory(max_tokens=self.config.max_context_tokens)
|
||||
|
||||
# Ensure litellm is available
|
||||
try:
|
||||
import litellm
|
||||
|
||||
# Drop unsupported params for models that don't support them
|
||||
litellm.drop_params = True
|
||||
self._litellm = litellm
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"litellm is required for LLM functionality. "
|
||||
"Install with: pip install litellm"
|
||||
) from e
|
||||
|
||||
def _is_rate_limit_error(self, error: Exception) -> bool:
|
||||
"""Check if an error is a rate limit error."""
|
||||
error_str = str(error).lower()
|
||||
error_type = type(error).__name__.lower()
|
||||
return (
|
||||
"rate" in error_str
|
||||
and "limit" in error_str
|
||||
or "ratelimit" in error_type
|
||||
or "429" in error_str
|
||||
or "too many requests" in error_str
|
||||
)
|
||||
|
||||
async def _retry_with_backoff(self, coro_factory, max_retries: int = None):
|
||||
"""
|
||||
Retry a coroutine with exponential backoff for rate limits.
|
||||
|
||||
Args:
|
||||
coro_factory: A callable that returns a new coroutine each call
|
||||
max_retries: Max retry attempts (uses config if not specified)
|
||||
"""
|
||||
retries = max_retries or self.config.max_retries
|
||||
base_delay = self.config.retry_delay
|
||||
|
||||
for attempt in range(retries + 1):
|
||||
try:
|
||||
return await coro_factory()
|
||||
except Exception as e:
|
||||
if not self._is_rate_limit_error(e) or attempt >= retries:
|
||||
raise
|
||||
|
||||
# Exponential backoff with jitter
|
||||
delay = base_delay * (2**attempt) + random.uniform(0, 1)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Should not reach here
|
||||
raise RuntimeError("Retry logic failed unexpectedly")
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
system_prompt: str,
|
||||
messages: List[dict],
|
||||
tools: Optional[List["Tool"]] = None,
|
||||
stream: bool = False,
|
||||
) -> LLMResponse:
|
||||
"""
|
||||
Generate a response from the LLM.
|
||||
|
||||
Args:
|
||||
system_prompt: The system prompt
|
||||
messages: Conversation messages
|
||||
tools: Available tools for function calling
|
||||
stream: Whether to stream the response
|
||||
|
||||
Returns:
|
||||
LLMResponse with the result
|
||||
"""
|
||||
# Build messages list
|
||||
llm_messages = [{"role": "system", "content": system_prompt}]
|
||||
|
||||
# Add conversation history with summarization if needed
|
||||
history = await self.memory.get_messages_with_summary(
|
||||
messages, llm_call=self._summarize_call
|
||||
)
|
||||
llm_messages.extend(history)
|
||||
|
||||
# Build tools list
|
||||
llm_tools = None
|
||||
if tools:
|
||||
llm_tools = [tool.to_llm_format() for tool in tools if tool.enabled]
|
||||
|
||||
try:
|
||||
# Build call kwargs - only pass non-default optional params
|
||||
# to avoid conflicts (e.g., Claude doesn't allow temperature + top_p together)
|
||||
call_kwargs = {
|
||||
"model": self.model,
|
||||
"messages": llm_messages,
|
||||
"tools": llm_tools,
|
||||
"temperature": self.config.temperature,
|
||||
"max_tokens": self.config.max_tokens,
|
||||
}
|
||||
# Only add optional params if explicitly changed from defaults
|
||||
if self.config.top_p != 1.0:
|
||||
call_kwargs["top_p"] = self.config.top_p
|
||||
if self.config.frequency_penalty != 0.0:
|
||||
call_kwargs["frequency_penalty"] = self.config.frequency_penalty
|
||||
if self.config.presence_penalty != 0.0:
|
||||
call_kwargs["presence_penalty"] = self.config.presence_penalty
|
||||
|
||||
# Call LLM with retry for rate limits
|
||||
async def _call():
|
||||
return await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
response = await self._retry_with_backoff(_call)
|
||||
|
||||
# Parse response
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
|
||||
# Handle usage - convert to dict safely
|
||||
usage_dict = None
|
||||
if response.usage:
|
||||
try:
|
||||
usage_dict = dict(response.usage)
|
||||
except (TypeError, ValueError):
|
||||
usage_dict = {
|
||||
"prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
|
||||
"completion_tokens": getattr(
|
||||
response.usage, "completion_tokens", 0
|
||||
),
|
||||
"total_tokens": getattr(response.usage, "total_tokens", 0),
|
||||
}
|
||||
|
||||
return LLMResponse(
|
||||
content=message.content,
|
||||
tool_calls=message.tool_calls,
|
||||
usage=usage_dict,
|
||||
model=response.model if hasattr(response, "model") else self.model,
|
||||
finish_reason=choice.finish_reason or "",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# Return error as response (after retries exhausted)
|
||||
return LLMResponse(
|
||||
content=f"LLM Error: {str(e)}",
|
||||
tool_calls=None,
|
||||
usage=None,
|
||||
model=self.model,
|
||||
finish_reason="error",
|
||||
)
|
||||
|
||||
async def generate_stream(
|
||||
self,
|
||||
system_prompt: str,
|
||||
messages: List[dict],
|
||||
tools: Optional[List["Tool"]] = None,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream a response from the LLM.
|
||||
|
||||
Args:
|
||||
system_prompt: The system prompt
|
||||
messages: Conversation messages
|
||||
tools: Available tools for function calling
|
||||
|
||||
Yields:
|
||||
Response content chunks
|
||||
"""
|
||||
llm_messages = [{"role": "system", "content": system_prompt}]
|
||||
history = await self.memory.get_messages_with_summary(
|
||||
messages, llm_call=self._summarize_call
|
||||
)
|
||||
llm_messages.extend(history)
|
||||
|
||||
llm_tools = None
|
||||
if tools:
|
||||
llm_tools = [tool.to_llm_format() for tool in tools if tool.enabled]
|
||||
|
||||
try:
|
||||
response = await self._litellm.acompletion(
|
||||
model=self.model,
|
||||
messages=llm_messages,
|
||||
tools=llm_tools,
|
||||
temperature=self.config.temperature,
|
||||
max_tokens=self.config.max_tokens,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
async for chunk in response:
|
||||
if chunk.choices[0].delta.content:
|
||||
yield chunk.choices[0].delta.content
|
||||
|
||||
except Exception as e:
|
||||
yield f"\nLLM Error: {str(e)}"
|
||||
|
||||
async def simple_completion(
|
||||
self, prompt: str, system: str = "You are a helpful assistant."
|
||||
) -> str:
|
||||
"""
|
||||
Simple completion without tools.
|
||||
|
||||
Args:
|
||||
prompt: The user prompt
|
||||
system: The system prompt
|
||||
|
||||
Returns:
|
||||
The response text
|
||||
"""
|
||||
response = await self.generate(
|
||||
system_prompt=system,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
tools=None,
|
||||
)
|
||||
return response.content or ""
|
||||
|
||||
def set_model(self, model: str):
|
||||
"""Change the model."""
|
||||
self.model = model
|
||||
|
||||
def update_config(self, **kwargs):
|
||||
"""Update configuration parameters."""
|
||||
for key, value in kwargs.items():
|
||||
if hasattr(self.config, key):
|
||||
setattr(self.config, key, value)
|
||||
|
||||
async def _summarize_call(self, prompt: str) -> str:
|
||||
"""
|
||||
Internal LLM call for summarization.
|
||||
|
||||
Args:
|
||||
prompt: The summarization prompt
|
||||
|
||||
Returns:
|
||||
Summary text
|
||||
"""
|
||||
try:
|
||||
response = await self._litellm.acompletion(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a terse summarizer for a pentesting agent.",
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
temperature=0.3, # Lower temperature for consistent summaries
|
||||
max_tokens=1000, # Summaries should be concise
|
||||
)
|
||||
return response.choices[0].message.content or ""
|
||||
except Exception as e:
|
||||
return f"[Summarization failed: {e}]"
|
||||
|
||||
def clear_memory(self):
|
||||
"""Clear conversation memory and summary cache."""
|
||||
self.memory.clear_summary_cache()
|
||||
|
||||
def get_memory_stats(self) -> dict:
|
||||
"""Get memory usage statistics."""
|
||||
return self.memory.get_stats()
|
||||
|
||||
def get_available_models(self) -> List[str]:
|
||||
"""
|
||||
Get list of commonly available models.
|
||||
|
||||
Returns:
|
||||
List of model names
|
||||
"""
|
||||
return [
|
||||
# OpenAI
|
||||
"gpt-5",
|
||||
"gpt-4.1",
|
||||
"gpt-4.1-mini",
|
||||
"gpt-4.1-nano",
|
||||
# Anthropic
|
||||
"claude-sonnet-4-20250514",
|
||||
"claude-opus-4-20250514",
|
||||
# Google
|
||||
"gemini-2.5-pro",
|
||||
"gemini-2.5-flash",
|
||||
# Others via LiteLLM
|
||||
"ollama/llama3",
|
||||
"ollama/mixtral",
|
||||
"groq/llama3-70b-8192",
|
||||
]
|
||||
253
ghostcrew/llm/memory.py
Normal file
253
ghostcrew/llm/memory.py
Normal file
@@ -0,0 +1,253 @@
|
||||
"""Conversation memory management for GhostCrew."""
|
||||
|
||||
from typing import Awaitable, Callable, List, Optional
|
||||
|
||||
SUMMARY_PROMPT = """Summarize this conversation history for a pentesting agent. Be terse.
|
||||
|
||||
Focus on:
|
||||
- Targets discovered (IPs, domains, hosts)
|
||||
- Open ports and services found
|
||||
- Credentials or secrets discovered
|
||||
- Vulnerabilities identified
|
||||
- What was attempted and failed (to avoid repeating)
|
||||
- Current objective/progress
|
||||
|
||||
Omit: verbose tool output, back-and-forth clarifications, redundant info.
|
||||
|
||||
Conversation to summarize:
|
||||
{conversation}
|
||||
|
||||
Summary:"""
|
||||
|
||||
|
||||
class ConversationMemory:
|
||||
"""Manages conversation history with token limits and summarization."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_tokens: int = 128000,
|
||||
reserve_ratio: float = 0.8,
|
||||
recent_to_keep: int = 10,
|
||||
summarize_threshold: float = 0.6,
|
||||
):
|
||||
"""
|
||||
Initialize conversation memory.
|
||||
|
||||
Args:
|
||||
max_tokens: Maximum context tokens
|
||||
reserve_ratio: Ratio of tokens to use (leave room for response)
|
||||
recent_to_keep: Number of recent messages to keep in full
|
||||
summarize_threshold: Summarize when history exceeds this ratio of budget
|
||||
"""
|
||||
self.max_tokens = max_tokens
|
||||
self.reserve_ratio = reserve_ratio
|
||||
self.recent_to_keep = recent_to_keep
|
||||
self.summarize_threshold = summarize_threshold
|
||||
self._encoder = None
|
||||
self._cached_summary: Optional[str] = None
|
||||
self._summarized_count: int = 0
|
||||
|
||||
@property
|
||||
def encoder(self):
|
||||
"""Lazy load the tokenizer."""
|
||||
if self._encoder is None:
|
||||
try:
|
||||
import tiktoken
|
||||
|
||||
self._encoder = tiktoken.get_encoding("cl100k_base")
|
||||
except ImportError:
|
||||
self._encoder = None
|
||||
return self._encoder
|
||||
|
||||
@property
|
||||
def token_budget(self) -> int:
|
||||
"""Available tokens for history."""
|
||||
return int(self.max_tokens * self.reserve_ratio)
|
||||
|
||||
def get_messages(self, messages: List[dict]) -> List[dict]:
|
||||
"""
|
||||
Get messages that fit within token limit (sync, no summarization).
|
||||
Falls back to truncation if over budget.
|
||||
|
||||
Args:
|
||||
messages: Full conversation history
|
||||
|
||||
Returns:
|
||||
Messages that fit within the token budget
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
# If we have a cached summary, prepend it
|
||||
if self._cached_summary and len(messages) > self._summarized_count:
|
||||
result = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"Previous conversation summary:\n{self._cached_summary}",
|
||||
}
|
||||
]
|
||||
# Add messages after the summarized portion
|
||||
recent = messages[self._summarized_count :]
|
||||
result.extend(
|
||||
self._truncate_to_fit(
|
||||
recent, self.token_budget - self._count_tokens(result[0])
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
return self._truncate_to_fit(messages, self.token_budget)
|
||||
|
||||
async def get_messages_with_summary(
|
||||
self, messages: List[dict], llm_call: Callable[[str], Awaitable[str]]
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Get messages, summarizing older ones if needed.
|
||||
|
||||
Args:
|
||||
messages: Full conversation history
|
||||
llm_call: Async function to call LLM for summarization
|
||||
|
||||
Returns:
|
||||
Messages with older ones summarized if over threshold
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
total_tokens = self.get_total_tokens(messages)
|
||||
threshold_tokens = int(self.token_budget * self.summarize_threshold)
|
||||
|
||||
# Check if we need to summarize
|
||||
if total_tokens <= threshold_tokens:
|
||||
return messages
|
||||
|
||||
# Don't summarize if we don't have enough messages
|
||||
if len(messages) <= self.recent_to_keep:
|
||||
return self._truncate_to_fit(messages, self.token_budget)
|
||||
|
||||
# Split messages: older to summarize, recent to keep
|
||||
split_point = len(messages) - self.recent_to_keep
|
||||
older = messages[:split_point]
|
||||
recent = messages[-self.recent_to_keep :]
|
||||
|
||||
# Check if we already summarized these messages
|
||||
if split_point <= self._summarized_count and self._cached_summary:
|
||||
result = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"Previous conversation summary:\n{self._cached_summary}",
|
||||
}
|
||||
]
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
# Summarize older messages
|
||||
summary = await self._summarize(older, llm_call)
|
||||
|
||||
# Cache the summary
|
||||
self._cached_summary = summary
|
||||
self._summarized_count = split_point
|
||||
|
||||
# Build result
|
||||
result = [
|
||||
{"role": "system", "content": f"Previous conversation summary:\n{summary}"}
|
||||
]
|
||||
result.extend(recent)
|
||||
|
||||
return result
|
||||
|
||||
async def _summarize(
|
||||
self, messages: List[dict], llm_call: Callable[[str], Awaitable[str]]
|
||||
) -> str:
|
||||
"""
|
||||
Summarize a list of messages.
|
||||
|
||||
Args:
|
||||
messages: Messages to summarize
|
||||
llm_call: Async function to call LLM
|
||||
|
||||
Returns:
|
||||
Summary string
|
||||
"""
|
||||
# Format messages for summarization
|
||||
conversation_text = self._format_for_summary(messages)
|
||||
|
||||
# Call LLM for summary
|
||||
prompt = SUMMARY_PROMPT.format(conversation=conversation_text)
|
||||
|
||||
try:
|
||||
summary = await llm_call(prompt)
|
||||
return summary.strip()
|
||||
except Exception as e:
|
||||
# Fallback: simple truncation indicator
|
||||
return f"[{len(messages)} earlier messages - summarization failed: {e}]"
|
||||
|
||||
def _format_for_summary(self, messages: List[dict]) -> str:
|
||||
"""Format messages as text for summarization."""
|
||||
lines = []
|
||||
for msg in messages:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
|
||||
# Truncate very long messages for summarization input
|
||||
if len(content) > 2000:
|
||||
content = content[:2000] + "...[truncated]"
|
||||
|
||||
if role == "user":
|
||||
lines.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
lines.append(f"Assistant: {content}")
|
||||
elif role == "tool":
|
||||
tool_name = msg.get("name", "tool")
|
||||
lines.append(f"Tool ({tool_name}): {content}")
|
||||
|
||||
return "\n\n".join(lines)
|
||||
|
||||
def _truncate_to_fit(self, messages: List[dict], budget: int) -> List[dict]:
|
||||
"""Truncate messages from the beginning to fit budget."""
|
||||
total_tokens = 0
|
||||
result = []
|
||||
|
||||
for msg in reversed(messages):
|
||||
msg_tokens = self._count_tokens(msg)
|
||||
if total_tokens + msg_tokens > budget:
|
||||
break
|
||||
result.insert(0, msg)
|
||||
total_tokens += msg_tokens
|
||||
|
||||
return result
|
||||
|
||||
def _count_tokens(self, message: dict) -> int:
|
||||
"""Count tokens in a message."""
|
||||
content = message.get("content", "")
|
||||
|
||||
if isinstance(content, str):
|
||||
if self.encoder:
|
||||
return len(self.encoder.encode(content))
|
||||
else:
|
||||
return int(len(content.split()) * 1.3)
|
||||
|
||||
return 0
|
||||
|
||||
def get_total_tokens(self, messages: List[dict]) -> int:
|
||||
"""Get total token count for messages."""
|
||||
return sum(self._count_tokens(msg) for msg in messages)
|
||||
|
||||
def fits_in_context(self, messages: List[dict]) -> bool:
|
||||
"""Check if messages fit in context window."""
|
||||
return self.get_total_tokens(messages) <= self.token_budget
|
||||
|
||||
def clear_summary_cache(self):
|
||||
"""Clear the cached summary (call when conversation is cleared)."""
|
||||
self._cached_summary = None
|
||||
self._summarized_count = 0
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
"""Get memory statistics."""
|
||||
return {
|
||||
"max_tokens": self.max_tokens,
|
||||
"token_budget": self.token_budget,
|
||||
"summarize_threshold": int(self.token_budget * self.summarize_threshold),
|
||||
"recent_to_keep": self.recent_to_keep,
|
||||
"has_summary": self._cached_summary is not None,
|
||||
"summarized_message_count": self._summarized_count,
|
||||
}
|
||||
201
ghostcrew/llm/utils.py
Normal file
201
ghostcrew/llm/utils.py
Normal file
@@ -0,0 +1,201 @@
|
||||
"""LLM utility functions for GhostCrew."""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
def count_tokens(text: str, model: str = "gpt-4") -> int:
|
||||
"""
|
||||
Count tokens in text.
|
||||
|
||||
Args:
|
||||
text: The text to count
|
||||
model: The model (used to select tokenizer)
|
||||
|
||||
Returns:
|
||||
Token count
|
||||
"""
|
||||
try:
|
||||
import tiktoken
|
||||
|
||||
# Select encoding based on model
|
||||
if "gpt-4" in model or "gpt-3.5" in model:
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
else:
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
return len(encoding.encode(text))
|
||||
|
||||
except ImportError:
|
||||
# Fallback approximation
|
||||
return int(len(text.split()) * 1.3)
|
||||
|
||||
|
||||
def truncate_to_tokens(text: str, max_tokens: int, model: str = "gpt-4") -> str:
|
||||
"""
|
||||
Truncate text to a maximum number of tokens.
|
||||
|
||||
Args:
|
||||
text: The text to truncate
|
||||
max_tokens: Maximum tokens
|
||||
model: The model (used to select tokenizer)
|
||||
|
||||
Returns:
|
||||
Truncated text
|
||||
"""
|
||||
try:
|
||||
import tiktoken
|
||||
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
tokens = encoding.encode(text)
|
||||
|
||||
if len(tokens) <= max_tokens:
|
||||
return text
|
||||
|
||||
truncated_tokens = tokens[:max_tokens]
|
||||
return encoding.decode(truncated_tokens)
|
||||
|
||||
except ImportError:
|
||||
# Fallback approximation
|
||||
words = text.split()
|
||||
target_words = int(max_tokens / 1.3)
|
||||
return " ".join(words[:target_words])
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""
|
||||
Quick estimation of token count without loading tokenizer.
|
||||
|
||||
Args:
|
||||
text: The text to estimate
|
||||
|
||||
Returns:
|
||||
Estimated token count
|
||||
"""
|
||||
# Average: ~4 characters per token for English
|
||||
return len(text) // 4
|
||||
|
||||
|
||||
def format_messages_for_display(messages: List[dict], max_length: int = 500) -> str:
|
||||
"""
|
||||
Format messages for display (e.g., in logs).
|
||||
|
||||
Args:
|
||||
messages: Messages to format
|
||||
max_length: Maximum length per message
|
||||
|
||||
Returns:
|
||||
Formatted string
|
||||
"""
|
||||
lines = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "unknown").upper()
|
||||
content = msg.get("content", "")
|
||||
|
||||
if len(content) > max_length:
|
||||
content = content[:max_length] + "..."
|
||||
|
||||
lines.append(f"[{role}] {content}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def extract_code_blocks(text: str) -> List[dict]:
|
||||
"""
|
||||
Extract code blocks from markdown text.
|
||||
|
||||
Args:
|
||||
text: Text containing code blocks
|
||||
|
||||
Returns:
|
||||
List of dicts with 'language' and 'code' keys
|
||||
"""
|
||||
import re
|
||||
|
||||
pattern = r"```(\w*)\n(.*?)```"
|
||||
matches = re.findall(pattern, text, re.DOTALL)
|
||||
|
||||
return [
|
||||
{"language": lang or "text", "code": code.strip()} for lang, code in matches
|
||||
]
|
||||
|
||||
|
||||
def extract_tool_calls_from_text(text: str) -> List[dict]:
|
||||
"""
|
||||
Extract tool call references from text (for display purposes).
|
||||
|
||||
Args:
|
||||
text: Text that may contain tool references
|
||||
|
||||
Returns:
|
||||
List of potential tool calls
|
||||
"""
|
||||
import re
|
||||
|
||||
# Look for patterns like "use tool_name" or "call tool_name"
|
||||
pattern = r"(?:use|call|execute|run)\s+(\w+)"
|
||||
matches = re.findall(pattern, text.lower())
|
||||
|
||||
return [{"tool": match} for match in matches]
|
||||
|
||||
|
||||
def sanitize_for_shell(text: str) -> str:
|
||||
"""
|
||||
Sanitize text for safe shell usage.
|
||||
|
||||
Args:
|
||||
text: Text to sanitize
|
||||
|
||||
Returns:
|
||||
Sanitized text
|
||||
"""
|
||||
# Escape dangerous characters
|
||||
dangerous = ["`", "$", "\\", '"', "'", ";", "&", "|", ">", "<", "\n", "\r"]
|
||||
|
||||
result = text
|
||||
for char in dangerous:
|
||||
result = result.replace(char, f"\\{char}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_llm_json(text: str) -> Optional[dict]:
|
||||
"""
|
||||
Attempt to parse JSON from LLM output.
|
||||
|
||||
Args:
|
||||
text: Text that may contain JSON
|
||||
|
||||
Returns:
|
||||
Parsed dict or None
|
||||
"""
|
||||
import json
|
||||
import re
|
||||
|
||||
# Try direct parse first
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Try to find JSON in code blocks
|
||||
pattern = r"```(?:json)?\n?(.*?)```"
|
||||
matches = re.findall(pattern, text, re.DOTALL)
|
||||
|
||||
for match in matches:
|
||||
try:
|
||||
return json.loads(match.strip())
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# Try to find JSON object in text
|
||||
pattern = r"\{[^{}]*\}"
|
||||
matches = re.findall(pattern, text)
|
||||
|
||||
for match in matches:
|
||||
try:
|
||||
return json.loads(match)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
return None
|
||||
17
ghostcrew/mcp/__init__.py
Normal file
17
ghostcrew/mcp/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""MCP (Model Context Protocol) integration for GhostCrew."""
|
||||
|
||||
from .discovery import MCPDiscovery
|
||||
from .manager import MCPManager, MCPServer, MCPServerConfig
|
||||
from .tools import create_mcp_tool
|
||||
from .transport import MCPTransport, SSETransport, StdioTransport
|
||||
|
||||
__all__ = [
|
||||
"MCPManager",
|
||||
"MCPServerConfig",
|
||||
"MCPServer",
|
||||
"MCPTransport",
|
||||
"StdioTransport",
|
||||
"SSETransport",
|
||||
"create_mcp_tool",
|
||||
"MCPDiscovery",
|
||||
]
|
||||
204
ghostcrew/mcp/discovery.py
Normal file
204
ghostcrew/mcp/discovery.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""MCP tool discovery for GhostCrew."""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiscoveredServer:
|
||||
"""A discovered MCP server."""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
type: str # "stdio" or "sse"
|
||||
command: Optional[str] = None
|
||||
args: Optional[List[str]] = None
|
||||
url: Optional[str] = None
|
||||
tools: List[dict] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.tools is None:
|
||||
self.tools = []
|
||||
|
||||
|
||||
class MCPDiscovery:
|
||||
"""Discovers available MCP servers and tools."""
|
||||
|
||||
# Known MCP servers for security tools
|
||||
KNOWN_SERVERS = [
|
||||
{
|
||||
"name": "nmap",
|
||||
"description": "Network scanning and host discovery",
|
||||
"type": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-nmap"],
|
||||
},
|
||||
{
|
||||
"name": "filesystem",
|
||||
"description": "File system operations",
|
||||
"type": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem"],
|
||||
},
|
||||
{
|
||||
"name": "fetch",
|
||||
"description": "HTTP requests and web fetching",
|
||||
"type": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-fetch"],
|
||||
},
|
||||
]
|
||||
|
||||
def __init__(self, config_path: Path = Path("mcp.json")):
|
||||
"""
|
||||
Initialize MCP discovery.
|
||||
|
||||
Args:
|
||||
config_path: Path to the MCP configuration file
|
||||
"""
|
||||
self.config_path = config_path
|
||||
|
||||
def discover_local(self) -> List[DiscoveredServer]:
|
||||
"""
|
||||
Discover locally installed MCP servers.
|
||||
|
||||
Returns:
|
||||
List of discovered servers
|
||||
"""
|
||||
discovered = []
|
||||
|
||||
# Check for npm global packages
|
||||
# Check for Python packages
|
||||
# This is a simplified implementation
|
||||
|
||||
for server_info in self.KNOWN_SERVERS:
|
||||
discovered.append(DiscoveredServer(**server_info))
|
||||
|
||||
return discovered
|
||||
|
||||
def load_from_config(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Load server configurations from file.
|
||||
|
||||
Returns:
|
||||
List of server configurations
|
||||
"""
|
||||
if not self.config_path.exists():
|
||||
return []
|
||||
|
||||
try:
|
||||
config = json.loads(self.config_path.read_text(encoding="utf-8"))
|
||||
return config.get("servers", [])
|
||||
except json.JSONDecodeError:
|
||||
return []
|
||||
|
||||
def add_server_to_config(
|
||||
self,
|
||||
name: str,
|
||||
server_type: str,
|
||||
command: Optional[str] = None,
|
||||
args: Optional[List[str]] = None,
|
||||
url: Optional[str] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Add a server to the configuration file.
|
||||
|
||||
Args:
|
||||
name: Server name
|
||||
server_type: "stdio" or "sse"
|
||||
command: Command for stdio servers
|
||||
args: Arguments for stdio servers
|
||||
url: URL for SSE servers
|
||||
env: Environment variables
|
||||
|
||||
Returns:
|
||||
True if added successfully
|
||||
"""
|
||||
# Load existing config
|
||||
if self.config_path.exists():
|
||||
try:
|
||||
config = json.loads(self.config_path.read_text(encoding="utf-8"))
|
||||
except json.JSONDecodeError:
|
||||
config = {"servers": []}
|
||||
else:
|
||||
config = {"servers": []}
|
||||
|
||||
# Check if server already exists
|
||||
for existing in config["servers"]:
|
||||
if existing.get("name") == name:
|
||||
return False
|
||||
|
||||
# Build server config
|
||||
server_config = {"name": name, "type": server_type, "enabled": True}
|
||||
|
||||
if server_type == "stdio":
|
||||
server_config["command"] = command
|
||||
server_config["args"] = args or []
|
||||
if env:
|
||||
server_config["env"] = env
|
||||
elif server_type == "sse":
|
||||
server_config["url"] = url
|
||||
|
||||
config["servers"].append(server_config)
|
||||
|
||||
# Save config
|
||||
self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
||||
|
||||
return True
|
||||
|
||||
def remove_server_from_config(self, name: str) -> bool:
|
||||
"""
|
||||
Remove a server from the configuration file.
|
||||
|
||||
Args:
|
||||
name: Server name to remove
|
||||
|
||||
Returns:
|
||||
True if removed successfully
|
||||
"""
|
||||
if not self.config_path.exists():
|
||||
return False
|
||||
|
||||
try:
|
||||
config = json.loads(self.config_path.read_text(encoding="utf-8"))
|
||||
except json.JSONDecodeError:
|
||||
return False
|
||||
|
||||
original_count = len(config.get("servers", []))
|
||||
config["servers"] = [
|
||||
s for s in config.get("servers", []) if s.get("name") != name
|
||||
]
|
||||
|
||||
if len(config["servers"]) == original_count:
|
||||
return False
|
||||
|
||||
self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
||||
|
||||
return True
|
||||
|
||||
def generate_default_config(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate a default MCP configuration.
|
||||
|
||||
Returns:
|
||||
Default configuration dictionary
|
||||
"""
|
||||
return {
|
||||
"servers": [
|
||||
{
|
||||
"name": "filesystem",
|
||||
"type": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
|
||||
"enabled": True,
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
def save_default_config(self):
|
||||
"""Save the default configuration to file."""
|
||||
config = self.generate_default_config()
|
||||
self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
||||
279
ghostcrew/mcp/manager.py
Normal file
279
ghostcrew/mcp/manager.py
Normal file
@@ -0,0 +1,279 @@
|
||||
"""MCP server connection manager for GhostCrew.
|
||||
|
||||
Uses standard MCP configuration format:
|
||||
{
|
||||
"mcpServers": {
|
||||
"server-name": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "package-name"],
|
||||
"env": {"VAR": "value"}
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .tools import create_mcp_tool
|
||||
from .transport import MCPTransport, StdioTransport
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPServerConfig:
|
||||
"""Configuration for an MCP server."""
|
||||
|
||||
name: str
|
||||
command: str
|
||||
args: List[str] = field(default_factory=list)
|
||||
env: Dict[str, str] = field(default_factory=dict)
|
||||
enabled: bool = True
|
||||
description: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPServer:
|
||||
"""Represents a connected MCP server."""
|
||||
|
||||
name: str
|
||||
config: MCPServerConfig
|
||||
transport: MCPTransport
|
||||
tools: List[dict] = field(default_factory=list)
|
||||
connected: bool = False
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect from the server."""
|
||||
if self.connected:
|
||||
await self.transport.disconnect()
|
||||
self.connected = False
|
||||
|
||||
|
||||
class MCPManager:
|
||||
"""Manages MCP server connections and exposes tools to agents."""
|
||||
|
||||
DEFAULT_CONFIG_PATHS = [
|
||||
Path.cwd() / "mcp_servers.json",
|
||||
Path.cwd() / "mcp.json",
|
||||
Path(__file__).parent / "mcp_servers.json",
|
||||
Path.home() / ".ghostcrew" / "mcp_servers.json",
|
||||
]
|
||||
|
||||
def __init__(self, config_path: Optional[Path] = None):
|
||||
self.config_path = config_path or self._find_config()
|
||||
self.servers: Dict[str, MCPServer] = {}
|
||||
self._message_id = 0
|
||||
|
||||
def _find_config(self) -> Path:
|
||||
for path in self.DEFAULT_CONFIG_PATHS:
|
||||
if path.exists():
|
||||
return path
|
||||
return self.DEFAULT_CONFIG_PATHS[0]
|
||||
|
||||
def _get_next_id(self) -> int:
|
||||
self._message_id += 1
|
||||
return self._message_id
|
||||
|
||||
def _load_config(self) -> Dict[str, MCPServerConfig]:
|
||||
if not self.config_path.exists():
|
||||
return {}
|
||||
try:
|
||||
raw = json.loads(self.config_path.read_text(encoding="utf-8"))
|
||||
servers = {}
|
||||
mcp_servers = raw.get("mcpServers", {})
|
||||
for name, config in mcp_servers.items():
|
||||
if not config.get("command"):
|
||||
continue
|
||||
servers[name] = MCPServerConfig(
|
||||
name=name,
|
||||
command=config["command"],
|
||||
args=config.get("args", []),
|
||||
env=config.get("env", {}),
|
||||
enabled=config.get("enabled", True),
|
||||
description=config.get("description", ""),
|
||||
)
|
||||
return servers
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"[MCP] Error loading config: {e}")
|
||||
return {}
|
||||
|
||||
def _save_config(self, servers: Dict[str, MCPServerConfig]):
|
||||
config = {"mcpServers": {}}
|
||||
for name, server in servers.items():
|
||||
server_config = {"command": server.command, "args": server.args}
|
||||
if server.env:
|
||||
server_config["env"] = server.env
|
||||
if server.description:
|
||||
server_config["description"] = server.description
|
||||
if not server.enabled:
|
||||
server_config["enabled"] = False
|
||||
config["mcpServers"][name] = server_config
|
||||
self.config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
||||
|
||||
def add_server(
|
||||
self,
|
||||
name: str,
|
||||
command: str,
|
||||
args: List[str] = None,
|
||||
env: Dict[str, str] = None,
|
||||
description: str = "",
|
||||
):
|
||||
servers = self._load_config()
|
||||
servers[name] = MCPServerConfig(
|
||||
name=name,
|
||||
command=command,
|
||||
args=args or [],
|
||||
env=env or {},
|
||||
description=description,
|
||||
)
|
||||
self._save_config(servers)
|
||||
print(f"[MCP] Added server: {name}")
|
||||
|
||||
def remove_server(self, name: str) -> bool:
|
||||
servers = self._load_config()
|
||||
if name in servers:
|
||||
del servers[name]
|
||||
self._save_config(servers)
|
||||
return True
|
||||
return False
|
||||
|
||||
def list_configured_servers(self) -> List[dict]:
|
||||
servers = self._load_config()
|
||||
return [
|
||||
{
|
||||
"name": n,
|
||||
"command": s.command,
|
||||
"args": s.args,
|
||||
"env": s.env,
|
||||
"enabled": s.enabled,
|
||||
"description": s.description,
|
||||
"connected": n in self.servers and self.servers[n].connected,
|
||||
}
|
||||
for n, s in servers.items()
|
||||
]
|
||||
|
||||
async def connect_all(self) -> List[Any]:
|
||||
servers_config = self._load_config()
|
||||
all_tools = []
|
||||
for name, config in servers_config.items():
|
||||
if not config.enabled:
|
||||
continue
|
||||
server = await self._connect_server(config)
|
||||
if server:
|
||||
self.servers[name] = server
|
||||
for tool_def in server.tools:
|
||||
tool = create_mcp_tool(tool_def, server, self)
|
||||
all_tools.append(tool)
|
||||
print(f"[MCP] Connected to {name} with {len(server.tools)} tools")
|
||||
return all_tools
|
||||
|
||||
async def connect_server(self, name: str) -> Optional[MCPServer]:
|
||||
servers_config = self._load_config()
|
||||
if name not in servers_config:
|
||||
return None
|
||||
config = servers_config[name]
|
||||
server = await self._connect_server(config)
|
||||
if server:
|
||||
self.servers[name] = server
|
||||
return server
|
||||
|
||||
async def _connect_server(self, config: MCPServerConfig) -> Optional[MCPServer]:
|
||||
transport = None
|
||||
try:
|
||||
env = {**os.environ, **config.env}
|
||||
transport = StdioTransport(
|
||||
command=config.command, args=config.args, env=env
|
||||
)
|
||||
await transport.connect()
|
||||
|
||||
await transport.send(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": "ghostcrew", "version": "0.2.0"},
|
||||
},
|
||||
"id": self._get_next_id(),
|
||||
}
|
||||
)
|
||||
await transport.send(
|
||||
{"jsonrpc": "2.0", "method": "notifications/initialized"}
|
||||
)
|
||||
|
||||
tools_response = await transport.send(
|
||||
{"jsonrpc": "2.0", "method": "tools/list", "id": self._get_next_id()}
|
||||
)
|
||||
tools = tools_response.get("result", {}).get("tools", [])
|
||||
|
||||
return MCPServer(
|
||||
name=config.name,
|
||||
config=config,
|
||||
transport=transport,
|
||||
tools=tools,
|
||||
connected=True,
|
||||
)
|
||||
except Exception as e:
|
||||
# Clean up transport on failure
|
||||
if transport:
|
||||
try:
|
||||
await transport.disconnect()
|
||||
except Exception:
|
||||
pass
|
||||
print(f"[MCP] Failed to connect to {config.name}: {e}")
|
||||
return None
|
||||
|
||||
async def call_tool(self, server_name: str, tool_name: str, arguments: dict) -> Any:
|
||||
server = self.servers.get(server_name)
|
||||
if not server or not server.connected:
|
||||
raise ValueError(f"Server '{server_name}' not connected")
|
||||
|
||||
# Use 5 minute timeout for tool calls (scans can take a while)
|
||||
response = await server.transport.send(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "tools/call",
|
||||
"params": {"name": tool_name, "arguments": arguments},
|
||||
"id": self._get_next_id(),
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
if "error" in response:
|
||||
raise RuntimeError(f"MCP error: {response['error'].get('message')}")
|
||||
return response.get("result", {}).get("content", [])
|
||||
|
||||
async def disconnect_server(self, name: str):
|
||||
server = self.servers.get(name)
|
||||
if server:
|
||||
await server.disconnect()
|
||||
del self.servers[name]
|
||||
|
||||
async def disconnect_all(self):
|
||||
for server in list(self.servers.values()):
|
||||
await server.disconnect()
|
||||
self.servers.clear()
|
||||
|
||||
async def reconnect_all(self) -> List[Any]:
|
||||
"""Disconnect all servers and reconnect them.
|
||||
|
||||
Useful after cancellation leaves servers in a bad state.
|
||||
"""
|
||||
# Disconnect all
|
||||
await self.disconnect_all()
|
||||
|
||||
# Reconnect all configured servers
|
||||
return await self.connect_all()
|
||||
|
||||
def get_server(self, name: str) -> Optional[MCPServer]:
|
||||
return self.servers.get(name)
|
||||
|
||||
def get_all_servers(self) -> List[MCPServer]:
|
||||
return list(self.servers.values())
|
||||
|
||||
def is_connected(self, name: str) -> bool:
|
||||
server = self.servers.get(name)
|
||||
return server is not None and server.connected
|
||||
3
ghostcrew/mcp/mcp_servers.json
Normal file
3
ghostcrew/mcp/mcp_servers.json
Normal file
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"mcpServers": {}
|
||||
}
|
||||
123
ghostcrew/mcp/tools.py
Normal file
123
ghostcrew/mcp/tools.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""MCP tool wrapper for GhostCrew."""
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..runtime import Runtime
|
||||
from ..tools import Tool
|
||||
from .manager import MCPManager, MCPServer
|
||||
|
||||
|
||||
def create_mcp_tool(
|
||||
tool_def: dict, server: "MCPServer", manager: "MCPManager"
|
||||
) -> "Tool":
|
||||
"""
|
||||
Create a Tool instance from an MCP tool definition.
|
||||
|
||||
Args:
|
||||
tool_def: The MCP tool definition
|
||||
server: The MCP server that provides this tool
|
||||
manager: The MCP manager for making calls
|
||||
|
||||
Returns:
|
||||
A Tool instance that wraps the MCP tool
|
||||
"""
|
||||
from ..tools import Tool, ToolSchema
|
||||
|
||||
async def execute_mcp(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""Execute this MCP tool."""
|
||||
# Get the tool name (without mcp_ prefix)
|
||||
original_name = tool_def["name"]
|
||||
|
||||
try:
|
||||
result = await manager.call_tool(server.name, original_name, arguments)
|
||||
|
||||
# Format result
|
||||
if isinstance(result, list):
|
||||
formatted_parts = []
|
||||
for item in result:
|
||||
if isinstance(item, dict):
|
||||
if item.get("type") == "text":
|
||||
formatted_parts.append(item.get("text", ""))
|
||||
elif item.get("type") == "image":
|
||||
formatted_parts.append(
|
||||
f"[Image: {item.get('mimeType', 'unknown')}]"
|
||||
)
|
||||
elif item.get("type") == "resource":
|
||||
formatted_parts.append(
|
||||
f"[Resource: {item.get('uri', 'unknown')}]"
|
||||
)
|
||||
else:
|
||||
formatted_parts.append(str(item))
|
||||
else:
|
||||
formatted_parts.append(str(item))
|
||||
return "\n".join(formatted_parts)
|
||||
|
||||
return str(result)
|
||||
|
||||
except Exception as e:
|
||||
return f"MCP tool error: {str(e)}"
|
||||
|
||||
# Convert MCP schema to our schema format
|
||||
input_schema = tool_def.get("inputSchema", {})
|
||||
schema = ToolSchema(
|
||||
type=input_schema.get("type", "object"),
|
||||
properties=input_schema.get("properties", {}),
|
||||
required=input_schema.get("required", []),
|
||||
)
|
||||
|
||||
# Create unique name with server prefix
|
||||
tool_name = f"mcp_{server.name}_{tool_def['name']}"
|
||||
|
||||
return Tool(
|
||||
name=tool_name,
|
||||
description=tool_def.get("description", f"MCP tool from {server.name}"),
|
||||
schema=schema,
|
||||
execute_fn=execute_mcp,
|
||||
category=f"mcp:{server.name}",
|
||||
metadata={
|
||||
"mcp_server": server.name,
|
||||
"mcp_tool": tool_def["name"],
|
||||
"original_schema": input_schema,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def format_mcp_result(result: Any) -> str:
|
||||
"""
|
||||
Format an MCP tool result for display.
|
||||
|
||||
Args:
|
||||
result: The raw MCP result
|
||||
|
||||
Returns:
|
||||
Formatted string
|
||||
"""
|
||||
if isinstance(result, list):
|
||||
parts = []
|
||||
for item in result:
|
||||
if isinstance(item, dict):
|
||||
content_type = item.get("type", "unknown")
|
||||
|
||||
if content_type == "text":
|
||||
parts.append(item.get("text", ""))
|
||||
elif content_type == "image":
|
||||
mime = item.get("mimeType", "unknown")
|
||||
data_preview = item.get("data", "")[:50]
|
||||
parts.append(f"[Image ({mime}): {data_preview}...]")
|
||||
elif content_type == "resource":
|
||||
uri = item.get("uri", "unknown")
|
||||
parts.append(f"[Resource: {uri}]")
|
||||
else:
|
||||
parts.append(str(item))
|
||||
else:
|
||||
parts.append(str(item))
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
elif isinstance(result, dict):
|
||||
if "content" in result:
|
||||
return format_mcp_result(result["content"])
|
||||
return str(result)
|
||||
|
||||
return str(result)
|
||||
287
ghostcrew/mcp/transport.py
Normal file
287
ghostcrew/mcp/transport.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""MCP transport implementations for GhostCrew."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
class MCPTransport(ABC):
|
||||
"""Abstract base class for MCP transports."""
|
||||
|
||||
@abstractmethod
|
||||
async def connect(self):
|
||||
"""Establish the connection."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def send(self, message: dict, timeout: float = 15.0) -> dict:
|
||||
"""Send a message and receive a response."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def disconnect(self):
|
||||
"""Close the connection."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_connected(self) -> bool:
|
||||
"""Check if the transport is connected."""
|
||||
pass
|
||||
|
||||
|
||||
class StdioTransport(MCPTransport):
|
||||
"""MCP transport over stdio (for npx/uvx commands)."""
|
||||
|
||||
def __init__(self, command: str, args: list[str], env: Dict[str, str]):
|
||||
"""
|
||||
Initialize stdio transport.
|
||||
|
||||
Args:
|
||||
command: The command to run (e.g., 'npx', 'uvx')
|
||||
args: Command arguments
|
||||
env: Additional environment variables
|
||||
"""
|
||||
self.command = command
|
||||
self.args = args
|
||||
self.env = env
|
||||
self.process: Optional[asyncio.subprocess.Process] = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Check if the process is running."""
|
||||
return self.process is not None and self.process.returncode is None
|
||||
|
||||
async def connect(self):
|
||||
"""Start the MCP server process."""
|
||||
import shutil
|
||||
|
||||
# Merge environment variables
|
||||
full_env = {**os.environ, **self.env}
|
||||
|
||||
# On Windows, resolve commands like npx, uvx that may be .cmd/.ps1 wrappers
|
||||
if os.name == "nt":
|
||||
# Check for .cmd version first (more compatible)
|
||||
cmd_path = shutil.which(f"{self.command}.cmd")
|
||||
if cmd_path:
|
||||
resolved_command = cmd_path
|
||||
else:
|
||||
# Fall back to regular which
|
||||
resolved_command = shutil.which(self.command) or self.command
|
||||
else:
|
||||
resolved_command = self.command
|
||||
|
||||
self.process = await asyncio.create_subprocess_exec(
|
||||
resolved_command,
|
||||
*self.args,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=full_env,
|
||||
)
|
||||
|
||||
async def send(self, message: dict, timeout: float = 15.0) -> dict:
|
||||
"""
|
||||
Send a JSON-RPC message and wait for response.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
timeout: Timeout in seconds for response (default 15s)
|
||||
|
||||
Returns:
|
||||
The JSON-RPC response
|
||||
"""
|
||||
if not self.process or not self.process.stdin or not self.process.stdout:
|
||||
raise RuntimeError("Transport not connected")
|
||||
|
||||
async with self._lock:
|
||||
# Send JSON-RPC message with newline
|
||||
msg_bytes = (json.dumps(message) + "\n").encode()
|
||||
self.process.stdin.write(msg_bytes)
|
||||
await self.process.stdin.drain()
|
||||
|
||||
# Notifications don't have responses
|
||||
if "id" not in message:
|
||||
return {}
|
||||
|
||||
# Read response line
|
||||
try:
|
||||
response_line = await asyncio.wait_for(
|
||||
self.process.stdout.readline(), timeout=timeout
|
||||
)
|
||||
|
||||
if not response_line:
|
||||
raise RuntimeError("Server closed connection")
|
||||
|
||||
return json.loads(response_line.decode())
|
||||
|
||||
except asyncio.TimeoutError as e:
|
||||
raise RuntimeError("Timeout waiting for MCP response") from e
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Invalid JSON response: {e}") from e
|
||||
|
||||
async def disconnect(self):
|
||||
"""Terminate the MCP server process cleanly."""
|
||||
if not self.process:
|
||||
return
|
||||
|
||||
proc = self.process
|
||||
self.process = None
|
||||
|
||||
# Close all pipes first to prevent "unclosed transport" warnings
|
||||
for pipe in (proc.stdin, proc.stdout, proc.stderr):
|
||||
if pipe:
|
||||
try:
|
||||
pipe.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Wait for pipe transports to close
|
||||
if proc.stdin:
|
||||
try:
|
||||
await proc.stdin.wait_closed()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Terminate the process
|
||||
try:
|
||||
proc.terminate()
|
||||
await asyncio.wait_for(proc.wait(), timeout=2.0)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
try:
|
||||
await proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class SSETransport(MCPTransport):
|
||||
"""MCP transport over Server-Sent Events (HTTP)."""
|
||||
|
||||
def __init__(self, url: str):
|
||||
"""
|
||||
Initialize SSE transport.
|
||||
|
||||
Args:
|
||||
url: The HTTP endpoint URL
|
||||
"""
|
||||
self.url = url
|
||||
self.session: Optional[Any] = None # aiohttp.ClientSession
|
||||
self._connected = False
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Check if the session is active."""
|
||||
return self._connected and self.session is not None
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to the SSE endpoint."""
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
self.session = aiohttp.ClientSession()
|
||||
self._connected = True
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
"aiohttp is required for SSE transport. Install with: pip install aiohttp"
|
||||
) from e
|
||||
|
||||
async def send(self, message: dict) -> dict:
|
||||
"""
|
||||
Send a message via HTTP POST.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
|
||||
Returns:
|
||||
The JSON-RPC response
|
||||
"""
|
||||
if not self.session:
|
||||
raise RuntimeError("Transport not connected")
|
||||
|
||||
try:
|
||||
async with self.session.post(
|
||||
self.url, json=message, headers={"Content-Type": "application/json"}
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
raise RuntimeError(f"HTTP error: {response.status}")
|
||||
|
||||
return await response.json()
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"SSE request failed: {e}") from e
|
||||
|
||||
async def disconnect(self):
|
||||
"""Close the HTTP session."""
|
||||
if self.session:
|
||||
await self.session.close()
|
||||
self.session = None
|
||||
self._connected = False
|
||||
|
||||
|
||||
class WebSocketTransport(MCPTransport):
|
||||
"""MCP transport over WebSocket."""
|
||||
|
||||
def __init__(self, url: str):
|
||||
"""
|
||||
Initialize WebSocket transport.
|
||||
|
||||
Args:
|
||||
url: The WebSocket endpoint URL
|
||||
"""
|
||||
self.url = url
|
||||
self.ws: Optional[Any] = None
|
||||
self._connected = False
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Check if the WebSocket is connected."""
|
||||
return self._connected and self.ws is not None
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to the WebSocket endpoint."""
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
self._session = aiohttp.ClientSession()
|
||||
self.ws = await self._session.ws_connect(self.url)
|
||||
self._connected = True
|
||||
except ImportError as e:
|
||||
raise RuntimeError("aiohttp is required for WebSocket transport") from e
|
||||
|
||||
async def send(self, message: dict) -> dict:
|
||||
"""
|
||||
Send a message via WebSocket.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
|
||||
Returns:
|
||||
The JSON-RPC response
|
||||
"""
|
||||
if not self.ws:
|
||||
raise RuntimeError("Transport not connected")
|
||||
|
||||
await self.ws.send_json(message)
|
||||
|
||||
# Notifications don't have responses
|
||||
if "id" not in message:
|
||||
return {}
|
||||
|
||||
response = await self.ws.receive_json()
|
||||
return response
|
||||
|
||||
async def disconnect(self):
|
||||
"""Close the WebSocket connection."""
|
||||
if self.ws:
|
||||
await self.ws.close()
|
||||
self.ws = None
|
||||
if hasattr(self, "_session") and self._session:
|
||||
await self._session.close()
|
||||
self._session = None
|
||||
self._connected = False
|
||||
14
ghostcrew/runtime/__init__.py
Normal file
14
ghostcrew/runtime/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
"""Runtime environment for GhostCrew."""
|
||||
|
||||
from .docker_runtime import DockerRuntime
|
||||
from .runtime import CommandResult, EnvironmentInfo, LocalRuntime, Runtime
|
||||
from .tool_server import ToolServer
|
||||
|
||||
__all__ = [
|
||||
"Runtime",
|
||||
"CommandResult",
|
||||
"LocalRuntime",
|
||||
"DockerRuntime",
|
||||
"ToolServer",
|
||||
"EnvironmentInfo",
|
||||
]
|
||||
373
ghostcrew/runtime/docker_runtime.py
Normal file
373
ghostcrew/runtime/docker_runtime.py
Normal file
@@ -0,0 +1,373 @@
|
||||
"""Docker runtime for GhostCrew."""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import tarfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from .runtime import CommandResult, Runtime
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..mcp import MCPManager
|
||||
|
||||
|
||||
@dataclass
|
||||
class DockerConfig:
|
||||
"""Docker runtime configuration."""
|
||||
|
||||
image: str = "ghostcrew-kali:latest" # Built from Dockerfile.kali
|
||||
container_name: str = "ghostcrew-sandbox"
|
||||
network_mode: str = "bridge"
|
||||
cap_add: list = None
|
||||
volumes: dict = None
|
||||
environment: dict = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.cap_add is None:
|
||||
self.cap_add = ["NET_ADMIN"] # For VPN
|
||||
if self.volumes is None:
|
||||
self.volumes = {}
|
||||
if self.environment is None:
|
||||
self.environment = {}
|
||||
|
||||
|
||||
class DockerRuntime(Runtime):
|
||||
"""Manages Docker sandbox for tool execution."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[DockerConfig] = None,
|
||||
vpn_config: Optional[Path] = None,
|
||||
mcp_manager: Optional["MCPManager"] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the Docker runtime.
|
||||
|
||||
Args:
|
||||
config: Docker configuration
|
||||
vpn_config: Path to OpenVPN config file
|
||||
mcp_manager: MCP manager for tool calls
|
||||
"""
|
||||
super().__init__(mcp_manager)
|
||||
self.config = config or DockerConfig()
|
||||
self.vpn_config = vpn_config
|
||||
self.client = None
|
||||
self.container = None
|
||||
self._browser_context = None
|
||||
self._proxy_running = False
|
||||
self._proxy_port = 8080
|
||||
|
||||
async def start(self):
|
||||
"""Start the sandbox container."""
|
||||
try:
|
||||
import docker
|
||||
|
||||
self.client = docker.from_env()
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"docker is required for Docker runtime. "
|
||||
"Install with: pip install docker"
|
||||
) from e
|
||||
|
||||
# Check if container already exists
|
||||
try:
|
||||
self.container = self.client.containers.get(self.config.container_name)
|
||||
if self.container.status != "running":
|
||||
self.container.start()
|
||||
await asyncio.sleep(2) # Wait for container to fully start
|
||||
except Exception:
|
||||
# Create new container
|
||||
volumes = {
|
||||
str(Path.home() / ".ghostcrew"): {
|
||||
"bind": "/root/.ghostcrew",
|
||||
"mode": "rw",
|
||||
},
|
||||
**self.config.volumes,
|
||||
}
|
||||
|
||||
self.container = self.client.containers.run(
|
||||
self.config.image,
|
||||
name=self.config.container_name,
|
||||
detach=True,
|
||||
tty=True,
|
||||
cap_add=self.config.cap_add,
|
||||
volumes=volumes,
|
||||
network_mode=self.config.network_mode,
|
||||
environment=self.config.environment,
|
||||
)
|
||||
await asyncio.sleep(2)
|
||||
|
||||
# Setup VPN if configured
|
||||
if self.vpn_config:
|
||||
await self._setup_vpn()
|
||||
|
||||
async def stop(self):
|
||||
"""Stop and remove the sandbox container."""
|
||||
if self.container:
|
||||
try:
|
||||
self.container.stop(timeout=10)
|
||||
self.container.remove()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
self.container = None
|
||||
|
||||
async def execute_command(self, command: str, timeout: int = 300) -> CommandResult:
|
||||
"""
|
||||
Execute a command in the sandbox.
|
||||
|
||||
Args:
|
||||
command: The shell command to execute
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
CommandResult with output
|
||||
"""
|
||||
if not self.container:
|
||||
raise RuntimeError("Sandbox not started")
|
||||
|
||||
try:
|
||||
# Execute command
|
||||
exec_result = self.container.exec_run(
|
||||
cmd=["bash", "-c", command], demux=True
|
||||
)
|
||||
|
||||
stdout = (
|
||||
exec_result.output[0].decode(errors="replace")
|
||||
if exec_result.output[0]
|
||||
else ""
|
||||
)
|
||||
stderr = (
|
||||
exec_result.output[1].decode(errors="replace")
|
||||
if exec_result.output[1]
|
||||
else ""
|
||||
)
|
||||
|
||||
return CommandResult(
|
||||
exit_code=exec_result.exit_code, stdout=stdout, stderr=stderr
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return CommandResult(
|
||||
exit_code=-1, stdout="", stderr=f"Execution error: {str(e)}"
|
||||
)
|
||||
|
||||
async def browser_action(self, action: str, **kwargs) -> dict:
|
||||
"""
|
||||
Perform browser automation in the sandbox.
|
||||
|
||||
Args:
|
||||
action: The browser action to perform
|
||||
**kwargs: Action-specific arguments
|
||||
|
||||
Returns:
|
||||
Action result dictionary
|
||||
"""
|
||||
# This would communicate with a browser automation service in the container
|
||||
# For now, we'll use a simple implementation via terminal commands
|
||||
|
||||
if action == "navigate":
|
||||
url = kwargs.get("url", "")
|
||||
result = await self.execute_command(f"curl -s -L -o /tmp/page.html '{url}'")
|
||||
if result.success:
|
||||
content_result = await self.execute_command(
|
||||
"head -c 5000 /tmp/page.html"
|
||||
)
|
||||
return {
|
||||
"url": url,
|
||||
"title": "Retrieved",
|
||||
"content": content_result.stdout,
|
||||
}
|
||||
return {"error": result.stderr}
|
||||
|
||||
elif action == "get_content":
|
||||
result = await self.execute_command(
|
||||
"cat /tmp/page.html 2>/dev/null || echo 'No page loaded'"
|
||||
)
|
||||
return {"content": result.stdout}
|
||||
|
||||
elif action == "get_links":
|
||||
result = await self.execute_command(
|
||||
"grep -oP 'href=\"\\K[^\"]+' /tmp/page.html 2>/dev/null | head -50"
|
||||
)
|
||||
links = [
|
||||
{"href": link, "text": ""}
|
||||
for link in result.stdout.strip().split("\n")
|
||||
if link
|
||||
]
|
||||
return {"links": links}
|
||||
|
||||
elif action == "screenshot":
|
||||
return {
|
||||
"error": "Screenshot requires graphical browser - not available in sandbox"
|
||||
}
|
||||
|
||||
return {"error": f"Unknown browser action: {action}"}
|
||||
|
||||
async def proxy_action(self, action: str, **kwargs) -> dict:
|
||||
"""
|
||||
Control the HTTP proxy in the sandbox.
|
||||
|
||||
Args:
|
||||
action: The proxy action to perform
|
||||
**kwargs: Action-specific arguments
|
||||
|
||||
Returns:
|
||||
Action result dictionary
|
||||
"""
|
||||
port = kwargs.get("port", self._proxy_port)
|
||||
|
||||
if action == "start":
|
||||
# Start mitmproxy in the background
|
||||
result = await self.execute_command(
|
||||
f"mitmdump -p {port} --set block_global=false -w /tmp/proxy.flow &"
|
||||
)
|
||||
if result.exit_code == 0:
|
||||
self._proxy_running = True
|
||||
self._proxy_port = port
|
||||
return {"status": "started", "port": port}
|
||||
return {"error": result.stderr}
|
||||
|
||||
elif action == "stop":
|
||||
await self.execute_command("pkill -f mitmdump")
|
||||
self._proxy_running = False
|
||||
return {"status": "stopped"}
|
||||
|
||||
elif action == "status":
|
||||
result = await self.execute_command("pgrep -f mitmdump")
|
||||
return {
|
||||
"status": "running" if result.exit_code == 0 else "stopped",
|
||||
"port": self._proxy_port,
|
||||
"request_count": 0, # Would need to parse proxy logs
|
||||
}
|
||||
|
||||
elif action == "get_history":
|
||||
# Would parse /tmp/proxy.flow
|
||||
return {"requests": []}
|
||||
|
||||
elif action == "clear_history":
|
||||
await self.execute_command("rm -f /tmp/proxy.flow")
|
||||
return {"status": "cleared"}
|
||||
|
||||
return {"error": f"Unknown proxy action: {action}"}
|
||||
|
||||
async def is_running(self) -> bool:
|
||||
"""Check if the container is running."""
|
||||
if not self.container:
|
||||
return False
|
||||
|
||||
try:
|
||||
self.container.reload()
|
||||
return self.container.status == "running"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_status(self) -> dict:
|
||||
"""Get runtime status information."""
|
||||
running = await self.is_running()
|
||||
|
||||
status = {
|
||||
"type": "docker",
|
||||
"running": running,
|
||||
"container_name": self.config.container_name,
|
||||
"image": self.config.image,
|
||||
"proxy_running": self._proxy_running,
|
||||
}
|
||||
|
||||
if running:
|
||||
# Get container info
|
||||
self.container.reload()
|
||||
status["container_id"] = self.container.short_id
|
||||
status["network_mode"] = self.config.network_mode
|
||||
|
||||
return status
|
||||
|
||||
async def _setup_vpn(self):
|
||||
"""Configure VPN in the sandbox."""
|
||||
if not self.vpn_config or not self.vpn_config.exists():
|
||||
return
|
||||
|
||||
# Copy VPN config to container
|
||||
config_content = self.vpn_config.read_bytes()
|
||||
tar_data = self._create_tar(config_content, "client.ovpn")
|
||||
|
||||
self.container.put_archive("/etc/openvpn", tar_data)
|
||||
|
||||
# Start OpenVPN
|
||||
await self.execute_command(
|
||||
"openvpn --config /etc/openvpn/client.ovpn --daemon --log /var/log/openvpn.log"
|
||||
)
|
||||
|
||||
# Wait for connection
|
||||
await asyncio.sleep(5)
|
||||
|
||||
# Verify connection
|
||||
result = await self.execute_command("curl -s --max-time 10 ifconfig.me")
|
||||
if result.success:
|
||||
print(f"[VPN] Connected. External IP: {result.stdout.strip()}")
|
||||
else:
|
||||
print("[VPN] Connection may have failed")
|
||||
|
||||
def _create_tar(self, content: bytes, filename: str) -> bytes:
|
||||
"""Create a tar archive for container upload."""
|
||||
tar_stream = io.BytesIO()
|
||||
tar = tarfile.open(fileobj=tar_stream, mode="w")
|
||||
|
||||
file_data = io.BytesIO(content)
|
||||
info = tarfile.TarInfo(name=filename)
|
||||
info.size = len(content)
|
||||
tar.addfile(info, file_data)
|
||||
tar.close()
|
||||
|
||||
tar_stream.seek(0)
|
||||
return tar_stream.read()
|
||||
|
||||
async def copy_to_container(self, local_path: Path, container_path: str):
|
||||
"""
|
||||
Copy a file to the container.
|
||||
|
||||
Args:
|
||||
local_path: Local file path
|
||||
container_path: Destination path in container
|
||||
"""
|
||||
if not self.container:
|
||||
raise RuntimeError("Container not running")
|
||||
|
||||
content = local_path.read_bytes()
|
||||
filename = local_path.name
|
||||
tar_data = self._create_tar(content, filename)
|
||||
|
||||
# Ensure directory exists
|
||||
dir_path = str(Path(container_path).parent)
|
||||
await self.execute_command(f"mkdir -p {dir_path}")
|
||||
|
||||
self.container.put_archive(dir_path, tar_data)
|
||||
|
||||
async def copy_from_container(self, container_path: str, local_path: Path):
|
||||
"""
|
||||
Copy a file from the container.
|
||||
|
||||
Args:
|
||||
container_path: Source path in container
|
||||
local_path: Local destination path
|
||||
"""
|
||||
if not self.container:
|
||||
raise RuntimeError("Container not running")
|
||||
|
||||
bits, stat = self.container.get_archive(container_path)
|
||||
|
||||
# Extract from tar
|
||||
tar_stream = io.BytesIO()
|
||||
for chunk in bits:
|
||||
tar_stream.write(chunk)
|
||||
tar_stream.seek(0)
|
||||
|
||||
tar = tarfile.open(fileobj=tar_stream)
|
||||
for member in tar.getmembers():
|
||||
f = tar.extractfile(member)
|
||||
if f:
|
||||
local_path.write_bytes(f.read())
|
||||
break
|
||||
tar.close()
|
||||
496
ghostcrew/runtime/runtime.py
Normal file
496
ghostcrew/runtime/runtime.py
Normal file
@@ -0,0 +1,496 @@
|
||||
"""Runtime abstraction for GhostCrew."""
|
||||
|
||||
import platform
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..mcp import MCPManager
|
||||
|
||||
|
||||
@dataclass
|
||||
class EnvironmentInfo:
|
||||
"""System environment information."""
|
||||
|
||||
os: str # "Windows", "Linux", "Darwin"
|
||||
os_version: str
|
||||
shell: str # "powershell", "bash", "zsh", etc.
|
||||
architecture: str # "x86_64", "arm64", etc.
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Concise string representation for prompts."""
|
||||
return f"{self.os} ({self.architecture}), shell: {self.shell}"
|
||||
|
||||
|
||||
def detect_environment() -> EnvironmentInfo:
|
||||
"""Detect the current system environment."""
|
||||
os_name = platform.system()
|
||||
os_version = platform.release()
|
||||
arch = platform.machine()
|
||||
|
||||
# Detect shell
|
||||
if os_name == "Windows":
|
||||
# Check for PowerShell vs CMD
|
||||
shell = "powershell"
|
||||
else:
|
||||
# Unix-like: check common shells
|
||||
import os
|
||||
|
||||
shell_path = os.environ.get("SHELL", "/bin/sh")
|
||||
shell = shell_path.split("/")[-1] # Extract shell name
|
||||
|
||||
return EnvironmentInfo(
|
||||
os=os_name, os_version=os_version, shell=shell, architecture=arch
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommandResult:
|
||||
"""Result of a command execution."""
|
||||
|
||||
exit_code: int
|
||||
stdout: str
|
||||
stderr: str
|
||||
|
||||
@property
|
||||
def success(self) -> bool:
|
||||
"""Check if the command succeeded."""
|
||||
return self.exit_code == 0
|
||||
|
||||
@property
|
||||
def output(self) -> str:
|
||||
"""Get combined output."""
|
||||
parts = []
|
||||
if self.stdout:
|
||||
parts.append(self.stdout)
|
||||
if self.stderr:
|
||||
parts.append(self.stderr)
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
class Runtime(ABC):
|
||||
"""Abstract base class for runtime environments."""
|
||||
|
||||
_environment: Optional[EnvironmentInfo] = None
|
||||
|
||||
def __init__(self, mcp_manager: Optional["MCPManager"] = None):
|
||||
"""
|
||||
Initialize the runtime.
|
||||
|
||||
Args:
|
||||
mcp_manager: Optional MCP manager for tool calls
|
||||
"""
|
||||
self.mcp_manager = mcp_manager
|
||||
|
||||
@property
|
||||
def environment(self) -> EnvironmentInfo:
|
||||
"""Get environment info (cached)."""
|
||||
if Runtime._environment is None:
|
||||
Runtime._environment = detect_environment()
|
||||
return Runtime._environment
|
||||
|
||||
@abstractmethod
|
||||
async def start(self):
|
||||
"""Start the runtime environment."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def stop(self):
|
||||
"""Stop the runtime environment."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def execute_command(self, command: str, timeout: int = 300) -> CommandResult:
|
||||
"""
|
||||
Execute a shell command.
|
||||
|
||||
Args:
|
||||
command: The command to execute
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
CommandResult with output
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def browser_action(self, action: str, **kwargs) -> dict:
|
||||
"""
|
||||
Perform a browser automation action.
|
||||
|
||||
Args:
|
||||
action: The action to perform
|
||||
**kwargs: Action-specific arguments
|
||||
|
||||
Returns:
|
||||
Action result
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def proxy_action(self, action: str, **kwargs) -> dict:
|
||||
"""
|
||||
Perform an HTTP proxy action.
|
||||
|
||||
Args:
|
||||
action: The action to perform
|
||||
**kwargs: Action-specific arguments
|
||||
|
||||
Returns:
|
||||
Action result
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def is_running(self) -> bool:
|
||||
"""Check if the runtime is running."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_status(self) -> dict:
|
||||
"""
|
||||
Get runtime status information.
|
||||
|
||||
Returns:
|
||||
Status dictionary
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class LocalRuntime(Runtime):
|
||||
"""Local runtime that executes commands directly on the host."""
|
||||
|
||||
def __init__(self, mcp_manager: Optional["MCPManager"] = None):
|
||||
super().__init__(mcp_manager)
|
||||
self._running = False
|
||||
self._browser = None
|
||||
self._browser_context = None
|
||||
self._page = None
|
||||
self._playwright = None
|
||||
self._active_processes: list = []
|
||||
|
||||
async def start(self):
|
||||
"""Start the local runtime."""
|
||||
self._running = True
|
||||
# Create loot directory for scan output
|
||||
Path("loot").mkdir(exist_ok=True)
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the local runtime gracefully."""
|
||||
# Clean up any active subprocesses
|
||||
for proc in self._active_processes:
|
||||
try:
|
||||
if proc.returncode is None:
|
||||
proc.terminate()
|
||||
await proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
self._active_processes.clear()
|
||||
|
||||
# Clean up browser
|
||||
await self._cleanup_browser()
|
||||
self._running = False
|
||||
|
||||
async def _cleanup_browser(self):
|
||||
"""Clean up browser resources properly."""
|
||||
# Close in reverse order of creation
|
||||
if self._page:
|
||||
try:
|
||||
await self._page.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._page = None
|
||||
|
||||
if self._browser_context:
|
||||
try:
|
||||
await self._browser_context.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._browser_context = None
|
||||
|
||||
if self._browser:
|
||||
try:
|
||||
await self._browser.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._browser = None
|
||||
|
||||
if self._playwright:
|
||||
try:
|
||||
await self._playwright.stop()
|
||||
except Exception:
|
||||
pass
|
||||
self._playwright = None
|
||||
|
||||
async def _ensure_browser(self):
|
||||
"""Ensure browser is initialized."""
|
||||
if self._page is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
from playwright.async_api import async_playwright
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
"Playwright not installed. Install with:\n"
|
||||
" pip install playwright\n"
|
||||
" playwright install chromium"
|
||||
) from e
|
||||
|
||||
self._playwright = await async_playwright().start()
|
||||
self._browser = await self._playwright.chromium.launch(headless=True)
|
||||
self._browser_context = await self._browser.new_context(
|
||||
user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
|
||||
)
|
||||
self._page = await self._browser_context.new_page()
|
||||
|
||||
async def execute_command(self, command: str, timeout: int = 300) -> CommandResult:
|
||||
"""Execute a command locally."""
|
||||
import asyncio
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
stdin=subprocess.DEVNULL, # Prevent interactive prompts
|
||||
)
|
||||
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(), timeout=timeout
|
||||
)
|
||||
|
||||
return CommandResult(
|
||||
exit_code=process.returncode or 0,
|
||||
stdout=stdout.decode(errors="replace"),
|
||||
stderr=stderr.decode(errors="replace"),
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
return CommandResult(
|
||||
exit_code=-1,
|
||||
stdout="",
|
||||
stderr=f"Command timed out after {timeout} seconds",
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
# Handle Ctrl+C gracefully
|
||||
return CommandResult(exit_code=-1, stdout="", stderr="Command cancelled")
|
||||
except Exception as e:
|
||||
return CommandResult(exit_code=-1, stdout="", stderr=str(e))
|
||||
|
||||
async def browser_action(self, action: str, **kwargs) -> dict:
|
||||
"""Perform browser automation actions using Playwright."""
|
||||
try:
|
||||
await self._ensure_browser()
|
||||
except RuntimeError as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
timeout = kwargs.get("timeout", 30) * 1000 # Convert to ms
|
||||
|
||||
try:
|
||||
if action == "navigate":
|
||||
url = kwargs.get("url")
|
||||
if not url:
|
||||
return {"error": "URL is required for navigate action"}
|
||||
|
||||
await self._page.goto(
|
||||
url, timeout=timeout, wait_until="domcontentloaded"
|
||||
)
|
||||
|
||||
if kwargs.get("wait_for"):
|
||||
await self._page.wait_for_selector(
|
||||
kwargs["wait_for"], timeout=timeout
|
||||
)
|
||||
|
||||
return {"url": self._page.url, "title": await self._page.title()}
|
||||
|
||||
elif action == "screenshot":
|
||||
from pathlib import Path
|
||||
|
||||
# Navigate first if URL provided
|
||||
if kwargs.get("url"):
|
||||
await self._page.goto(
|
||||
kwargs["url"], timeout=timeout, wait_until="domcontentloaded"
|
||||
)
|
||||
|
||||
# Save screenshot to loot directory
|
||||
output_dir = Path("loot/screenshots")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
filename = f"screenshot_{int(__import__('time').time())}.png"
|
||||
filepath = output_dir / filename
|
||||
|
||||
await self._page.screenshot(path=str(filepath), full_page=True)
|
||||
|
||||
return {"path": str(filepath)}
|
||||
|
||||
elif action == "get_content":
|
||||
if kwargs.get("url"):
|
||||
await self._page.goto(
|
||||
kwargs["url"], timeout=timeout, wait_until="domcontentloaded"
|
||||
)
|
||||
|
||||
content = await self._page.content()
|
||||
|
||||
# Also get text content for easier reading
|
||||
text_content = await self._page.evaluate(
|
||||
"() => document.body.innerText"
|
||||
)
|
||||
|
||||
return {
|
||||
"content": text_content,
|
||||
"html": content[:10000] if len(content) > 10000 else content,
|
||||
}
|
||||
|
||||
elif action == "get_links":
|
||||
if kwargs.get("url"):
|
||||
await self._page.goto(
|
||||
kwargs["url"], timeout=timeout, wait_until="domcontentloaded"
|
||||
)
|
||||
|
||||
links = await self._page.evaluate(
|
||||
"""() => {
|
||||
return Array.from(document.querySelectorAll('a[href]')).map(a => ({
|
||||
href: a.href,
|
||||
text: a.innerText.trim()
|
||||
}));
|
||||
}"""
|
||||
)
|
||||
|
||||
return {"links": links}
|
||||
|
||||
elif action == "get_forms":
|
||||
if kwargs.get("url"):
|
||||
await self._page.goto(
|
||||
kwargs["url"], timeout=timeout, wait_until="domcontentloaded"
|
||||
)
|
||||
|
||||
forms = await self._page.evaluate(
|
||||
"""() => {
|
||||
return Array.from(document.querySelectorAll('form')).map(form => ({
|
||||
action: form.action,
|
||||
method: form.method || 'GET',
|
||||
inputs: Array.from(form.querySelectorAll('input, textarea, select')).map(input => ({
|
||||
name: input.name,
|
||||
type: input.type || 'text',
|
||||
value: input.value
|
||||
}))
|
||||
}));
|
||||
}"""
|
||||
)
|
||||
|
||||
return {"forms": forms}
|
||||
|
||||
elif action == "click":
|
||||
selector = kwargs.get("selector")
|
||||
if not selector:
|
||||
return {"error": "Selector is required for click action"}
|
||||
|
||||
await self._page.click(selector, timeout=timeout)
|
||||
return {"selector": selector, "clicked": True}
|
||||
|
||||
elif action == "type":
|
||||
selector = kwargs.get("selector")
|
||||
text = kwargs.get("text", "")
|
||||
if not selector:
|
||||
return {"error": "Selector is required for type action"}
|
||||
|
||||
await self._page.fill(selector, text, timeout=timeout)
|
||||
return {"selector": selector, "typed": True}
|
||||
|
||||
elif action == "execute_js":
|
||||
javascript = kwargs.get("javascript")
|
||||
if not javascript:
|
||||
return {"error": "JavaScript code is required"}
|
||||
|
||||
result = await self._page.evaluate(javascript)
|
||||
return {"result": str(result) if result else ""}
|
||||
|
||||
else:
|
||||
return {"error": f"Unknown browser action: {action}"}
|
||||
|
||||
except Exception as e:
|
||||
return {"error": f"Browser action failed: {str(e)}"}
|
||||
|
||||
async def proxy_action(self, action: str, **kwargs) -> dict:
|
||||
"""HTTP proxy actions using httpx."""
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
return {"error": "httpx not installed. Install with: pip install httpx"}
|
||||
|
||||
timeout = kwargs.get("timeout", 30)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=timeout, follow_redirects=True
|
||||
) as client:
|
||||
if action == "request":
|
||||
method = kwargs.get("method", "GET").upper()
|
||||
url = kwargs.get("url")
|
||||
headers = kwargs.get("headers", {})
|
||||
data = kwargs.get("data")
|
||||
|
||||
if not url:
|
||||
return {"error": "URL is required"}
|
||||
|
||||
response = await client.request(
|
||||
method, url, headers=headers, data=data
|
||||
)
|
||||
|
||||
return {
|
||||
"status_code": response.status_code,
|
||||
"headers": dict(response.headers),
|
||||
"body": (
|
||||
response.text[:10000]
|
||||
if len(response.text) > 10000
|
||||
else response.text
|
||||
),
|
||||
}
|
||||
|
||||
elif action == "get":
|
||||
url = kwargs.get("url")
|
||||
if not url:
|
||||
return {"error": "URL is required"}
|
||||
|
||||
response = await client.get(url, headers=kwargs.get("headers", {}))
|
||||
return {
|
||||
"status_code": response.status_code,
|
||||
"headers": dict(response.headers),
|
||||
"body": response.text[:10000],
|
||||
}
|
||||
|
||||
elif action == "post":
|
||||
url = kwargs.get("url")
|
||||
if not url:
|
||||
return {"error": "URL is required"}
|
||||
|
||||
response = await client.post(
|
||||
url,
|
||||
headers=kwargs.get("headers", {}),
|
||||
data=kwargs.get("data"),
|
||||
json=kwargs.get("json"),
|
||||
)
|
||||
return {
|
||||
"status_code": response.status_code,
|
||||
"headers": dict(response.headers),
|
||||
"body": response.text[:10000],
|
||||
}
|
||||
|
||||
else:
|
||||
return {"error": f"Unknown proxy action: {action}"}
|
||||
|
||||
except Exception as e:
|
||||
return {"error": f"Proxy action failed: {str(e)}"}
|
||||
|
||||
async def is_running(self) -> bool:
|
||||
return self._running
|
||||
|
||||
async def get_status(self) -> dict:
|
||||
return {
|
||||
"type": "local",
|
||||
"running": self._running,
|
||||
"browser_active": self._page is not None,
|
||||
}
|
||||
202
ghostcrew/runtime/tool_server.py
Normal file
202
ghostcrew/runtime/tool_server.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""Tool server for running tools in the sandbox."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolRequest:
|
||||
"""A tool execution request."""
|
||||
|
||||
tool_name: str
|
||||
arguments: dict
|
||||
request_id: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResponse:
|
||||
"""A tool execution response."""
|
||||
|
||||
request_id: str
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
success: bool = True
|
||||
|
||||
|
||||
class ToolServer:
|
||||
"""
|
||||
Server that runs inside the sandbox to handle tool requests.
|
||||
|
||||
This is used for more complex tool orchestration where
|
||||
tools need to run inside the container.
|
||||
"""
|
||||
|
||||
def __init__(self, host: str = "0.0.0.0", port: int = 9999):
|
||||
"""
|
||||
Initialize the tool server.
|
||||
|
||||
Args:
|
||||
host: Host to bind to
|
||||
port: Port to listen on
|
||||
"""
|
||||
self.host = host
|
||||
self.port = port
|
||||
self._tools: Dict[str, Callable] = {}
|
||||
self._server = None
|
||||
self._running = False
|
||||
|
||||
def register_tool(self, name: str, handler: Callable):
|
||||
"""
|
||||
Register a tool handler.
|
||||
|
||||
Args:
|
||||
name: Tool name
|
||||
handler: Async function to handle the tool
|
||||
"""
|
||||
self._tools[name] = handler
|
||||
|
||||
async def start(self):
|
||||
"""Start the tool server."""
|
||||
self._server = await asyncio.start_server(
|
||||
self._handle_connection, self.host, self.port
|
||||
)
|
||||
self._running = True
|
||||
|
||||
async with self._server:
|
||||
await self._server.serve_forever()
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the tool server."""
|
||||
self._running = False
|
||||
if self._server:
|
||||
self._server.close()
|
||||
await self._server.wait_closed()
|
||||
|
||||
async def _handle_connection(
|
||||
self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
||||
):
|
||||
"""Handle an incoming connection."""
|
||||
try:
|
||||
while self._running:
|
||||
# Read request
|
||||
data = await reader.readline()
|
||||
if not data:
|
||||
break
|
||||
|
||||
try:
|
||||
request_data = json.loads(data.decode())
|
||||
request = ToolRequest(
|
||||
tool_name=request_data["tool"],
|
||||
arguments=request_data.get("arguments", {}),
|
||||
request_id=request_data.get("id", "unknown"),
|
||||
)
|
||||
|
||||
# Execute tool
|
||||
response = await self._execute_tool(request)
|
||||
|
||||
# Send response
|
||||
response_data = {
|
||||
"id": response.request_id,
|
||||
"result": response.result,
|
||||
"error": response.error,
|
||||
"success": response.success,
|
||||
}
|
||||
writer.write((json.dumps(response_data) + "\n").encode())
|
||||
await writer.drain()
|
||||
|
||||
except json.JSONDecodeError:
|
||||
error_response = {"error": "Invalid JSON", "success": False}
|
||||
writer.write((json.dumps(error_response) + "\n").encode())
|
||||
await writer.drain()
|
||||
|
||||
except Exception as e:
|
||||
print(f"Connection error: {e}")
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
|
||||
async def _execute_tool(self, request: ToolRequest) -> ToolResponse:
|
||||
"""Execute a tool request."""
|
||||
handler = self._tools.get(request.tool_name)
|
||||
|
||||
if not handler:
|
||||
return ToolResponse(
|
||||
request_id=request.request_id,
|
||||
error=f"Tool '{request.tool_name}' not found",
|
||||
success=False,
|
||||
)
|
||||
|
||||
try:
|
||||
result = await handler(request.arguments)
|
||||
return ToolResponse(
|
||||
request_id=request.request_id, result=str(result), success=True
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResponse(
|
||||
request_id=request.request_id, error=str(e), success=False
|
||||
)
|
||||
|
||||
|
||||
class ToolClient:
|
||||
"""Client for communicating with the tool server in the sandbox."""
|
||||
|
||||
def __init__(self, host: str = "localhost", port: int = 9999):
|
||||
"""
|
||||
Initialize the tool client.
|
||||
|
||||
Args:
|
||||
host: Tool server host
|
||||
port: Tool server port
|
||||
"""
|
||||
self.host = host
|
||||
self.port = port
|
||||
self._reader: Optional[asyncio.StreamReader] = None
|
||||
self._writer: Optional[asyncio.StreamWriter] = None
|
||||
self._request_id = 0
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to the tool server."""
|
||||
self._reader, self._writer = await asyncio.open_connection(self.host, self.port)
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect from the tool server."""
|
||||
if self._writer:
|
||||
self._writer.close()
|
||||
await self._writer.wait_closed()
|
||||
self._writer = None
|
||||
self._reader = None
|
||||
|
||||
async def call_tool(self, tool_name: str, arguments: dict) -> ToolResponse:
|
||||
"""
|
||||
Call a tool on the server.
|
||||
|
||||
Args:
|
||||
tool_name: The tool to call
|
||||
arguments: Tool arguments
|
||||
|
||||
Returns:
|
||||
ToolResponse with result
|
||||
"""
|
||||
if not self._writer or not self._reader:
|
||||
raise RuntimeError("Not connected to tool server")
|
||||
|
||||
self._request_id += 1
|
||||
request_id = str(self._request_id)
|
||||
|
||||
# Send request
|
||||
request = {"id": request_id, "tool": tool_name, "arguments": arguments}
|
||||
self._writer.write((json.dumps(request) + "\n").encode())
|
||||
await self._writer.drain()
|
||||
|
||||
# Read response
|
||||
response_data = await self._reader.readline()
|
||||
response = json.loads(response_data.decode())
|
||||
|
||||
return ToolResponse(
|
||||
request_id=response.get("id", request_id),
|
||||
result=response.get("result"),
|
||||
error=response.get("error"),
|
||||
success=response.get("success", False),
|
||||
)
|
||||
42
ghostcrew/tools/__init__.py
Normal file
42
ghostcrew/tools/__init__.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""Tool system for GhostCrew."""
|
||||
|
||||
from .executor import ToolExecutor
|
||||
from .loader import discover_tools, get_tool_info, load_all_tools, reload_tools
|
||||
from .registry import (
|
||||
Tool,
|
||||
ToolSchema,
|
||||
clear_tools,
|
||||
disable_tool,
|
||||
enable_tool,
|
||||
get_all_tools,
|
||||
get_tool,
|
||||
get_tool_names,
|
||||
get_tools_by_category,
|
||||
register_tool,
|
||||
register_tool_instance,
|
||||
)
|
||||
|
||||
# Auto-load all built-in tools on import
|
||||
_loaded = load_all_tools()
|
||||
|
||||
__all__ = [
|
||||
# Registry
|
||||
"Tool",
|
||||
"ToolSchema",
|
||||
"register_tool",
|
||||
"get_all_tools",
|
||||
"get_tool",
|
||||
"register_tool_instance",
|
||||
"get_tools_by_category",
|
||||
"enable_tool",
|
||||
"disable_tool",
|
||||
"get_tool_names",
|
||||
"clear_tools",
|
||||
# Executor
|
||||
"ToolExecutor",
|
||||
# Loader
|
||||
"load_all_tools",
|
||||
"get_tool_info",
|
||||
"reload_tools",
|
||||
"discover_tools",
|
||||
]
|
||||
156
ghostcrew/tools/browser/__init__.py
Normal file
156
ghostcrew/tools/browser/__init__.py
Normal file
@@ -0,0 +1,156 @@
|
||||
"""Browser automation tool for GhostCrew."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..registry import ToolSchema, register_tool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...runtime import Runtime
|
||||
|
||||
|
||||
@register_tool(
|
||||
name="browser",
|
||||
description="Automate a headless browser. Actions: navigate, click, type, screenshot, get_content, get_links, get_forms, execute_js.",
|
||||
schema=ToolSchema(
|
||||
properties={
|
||||
"action": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"navigate",
|
||||
"click",
|
||||
"type",
|
||||
"screenshot",
|
||||
"get_content",
|
||||
"get_links",
|
||||
"get_forms",
|
||||
"execute_js",
|
||||
],
|
||||
"description": "The browser action to perform",
|
||||
},
|
||||
"url": {
|
||||
"type": "string",
|
||||
"description": "URL to navigate to (for 'navigate' action)",
|
||||
},
|
||||
"selector": {
|
||||
"type": "string",
|
||||
"description": "CSS selector for element (for 'click', 'type' actions)",
|
||||
},
|
||||
"text": {
|
||||
"type": "string",
|
||||
"description": "Text to type (for 'type' action)",
|
||||
},
|
||||
"javascript": {
|
||||
"type": "string",
|
||||
"description": "JavaScript code to execute (for 'execute_js' action)",
|
||||
},
|
||||
"wait_for": {
|
||||
"type": "string",
|
||||
"description": "CSS selector to wait for before continuing",
|
||||
},
|
||||
"timeout": {
|
||||
"type": "integer",
|
||||
"description": "Timeout in seconds (default: 30)",
|
||||
"default": 30,
|
||||
},
|
||||
},
|
||||
required=["action"],
|
||||
),
|
||||
category="web",
|
||||
)
|
||||
async def browser(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""
|
||||
Perform browser automation actions.
|
||||
|
||||
Args:
|
||||
arguments: Dictionary with action and related parameters
|
||||
runtime: The runtime environment
|
||||
|
||||
Returns:
|
||||
Result of the browser action
|
||||
"""
|
||||
action = arguments["action"]
|
||||
timeout = arguments.get("timeout", 30)
|
||||
|
||||
try:
|
||||
result = await runtime.browser_action(
|
||||
action=action,
|
||||
url=arguments.get("url"),
|
||||
selector=arguments.get("selector"),
|
||||
text=arguments.get("text"),
|
||||
javascript=arguments.get("javascript"),
|
||||
wait_for=arguments.get("wait_for"),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
return _format_browser_result(action, result)
|
||||
|
||||
except Exception as e:
|
||||
return f"Browser action '{action}' failed: {str(e)}"
|
||||
|
||||
|
||||
def _format_browser_result(action: str, result: dict) -> str:
|
||||
"""Format browser action result for display."""
|
||||
# Check for errors first
|
||||
if "error" in result:
|
||||
return f"Browser error: {result['error']}"
|
||||
|
||||
if action == "navigate":
|
||||
return f"Navigated to: {result.get('url', 'unknown')}\nTitle: {result.get('title', 'N/A')}"
|
||||
|
||||
elif action == "screenshot":
|
||||
return f"Screenshot saved to: {result.get('path', 'unknown')}"
|
||||
|
||||
elif action == "get_content":
|
||||
content = result.get("content", "")
|
||||
if len(content) > 5000:
|
||||
content = content[:5000] + "\n... (truncated)"
|
||||
return f"Page content:\n{content}"
|
||||
|
||||
elif action == "get_links":
|
||||
links = result.get("links", [])
|
||||
if not links:
|
||||
return "No links found on page"
|
||||
|
||||
formatted = ["Found links:"]
|
||||
for link in links[:50]: # Limit to 50 links
|
||||
text = link.get("text", "").strip()[:50]
|
||||
href = link.get("href", "")
|
||||
formatted.append(f" - [{text}] {href}")
|
||||
|
||||
if len(links) > 50:
|
||||
formatted.append(f" ... and {len(links) - 50} more links")
|
||||
|
||||
return "\n".join(formatted)
|
||||
|
||||
elif action == "get_forms":
|
||||
forms = result.get("forms", [])
|
||||
if not forms:
|
||||
return "No forms found on page"
|
||||
|
||||
formatted = ["Found forms:"]
|
||||
for i, form in enumerate(forms):
|
||||
formatted.append(f"\nForm {i + 1}:")
|
||||
formatted.append(f" Action: {form.get('action', 'N/A')}")
|
||||
formatted.append(f" Method: {form.get('method', 'GET')}")
|
||||
inputs = form.get("inputs", [])
|
||||
if inputs:
|
||||
formatted.append(" Inputs:")
|
||||
for inp in inputs:
|
||||
formatted.append(
|
||||
f" - {inp.get('name', 'unnamed')} ({inp.get('type', 'text')})"
|
||||
)
|
||||
|
||||
return "\n".join(formatted)
|
||||
|
||||
elif action == "click":
|
||||
return f"Clicked element: {result.get('selector', 'unknown')}"
|
||||
|
||||
elif action == "type":
|
||||
return f"Typed text into: {result.get('selector', 'unknown')}"
|
||||
|
||||
elif action == "execute_js":
|
||||
output = result.get("result", "")
|
||||
return f"JavaScript result:\n{output}"
|
||||
|
||||
else:
|
||||
return str(result)
|
||||
5
ghostcrew/tools/browser/browser.py
Normal file
5
ghostcrew/tools/browser/browser.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Browser tool implementation."""
|
||||
|
||||
from . import browser
|
||||
|
||||
__all__ = ["browser"]
|
||||
52
ghostcrew/tools/completion/__init__.py
Normal file
52
ghostcrew/tools/completion/__init__.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""Task completion tool for GhostCrew agent loop control."""
|
||||
|
||||
from ..registry import ToolSchema, register_tool
|
||||
|
||||
# Sentinel value to signal task completion
|
||||
TASK_COMPLETE_SIGNAL = "__TASK_COMPLETE__"
|
||||
|
||||
|
||||
@register_tool(
|
||||
name="finish",
|
||||
description="Signal that the current task is finished. Call this when you have completed ALL steps of the user's request. Include a concise summary of what was accomplished.",
|
||||
schema=ToolSchema(
|
||||
properties={
|
||||
"summary": {
|
||||
"type": "string",
|
||||
"description": "Brief summary of what was accomplished and any key findings",
|
||||
},
|
||||
},
|
||||
required=["summary"],
|
||||
),
|
||||
category="control",
|
||||
)
|
||||
async def finish(arguments: dict, runtime) -> str:
|
||||
"""
|
||||
Signal task completion to the agent framework.
|
||||
|
||||
This tool is called by the agent when it has finished all steps
|
||||
of the user's task. The framework uses this as an explicit
|
||||
termination signal rather than relying on LLM text output.
|
||||
|
||||
Args:
|
||||
arguments: Dictionary with 'summary' key
|
||||
runtime: The runtime environment (unused)
|
||||
|
||||
Returns:
|
||||
The completion signal with summary
|
||||
"""
|
||||
summary = arguments.get("summary", "Task completed.")
|
||||
# Return special signal that the framework recognizes
|
||||
return f"{TASK_COMPLETE_SIGNAL}:{summary}"
|
||||
|
||||
|
||||
def is_task_complete(result: str) -> bool:
|
||||
"""Check if a tool result signals task completion."""
|
||||
return result.startswith(TASK_COMPLETE_SIGNAL)
|
||||
|
||||
|
||||
def extract_completion_summary(result: str) -> str:
|
||||
"""Extract the summary from a task_complete result."""
|
||||
if is_task_complete(result):
|
||||
return result[len(TASK_COMPLETE_SIGNAL) + 1:] # +1 for the colon
|
||||
return result
|
||||
213
ghostcrew/tools/executor.py
Normal file
213
ghostcrew/tools/executor.py
Normal file
@@ -0,0 +1,213 @@
|
||||
"""Tool executor for GhostCrew."""
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..runtime import Runtime
|
||||
from .registry import Tool
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionResult:
|
||||
"""Result of a tool execution."""
|
||||
|
||||
tool_name: str
|
||||
arguments: dict
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
success: bool = True
|
||||
start_time: Optional[datetime] = None
|
||||
end_time: Optional[datetime] = None
|
||||
duration_ms: float = 0.0
|
||||
|
||||
@property
|
||||
def duration(self) -> float:
|
||||
"""Get execution duration in seconds."""
|
||||
return self.duration_ms / 1000.0
|
||||
|
||||
|
||||
class ToolExecutor:
|
||||
"""Handles tool execution with error handling and logging."""
|
||||
|
||||
def __init__(self, runtime: "Runtime", timeout: int = 300, max_retries: int = 0):
|
||||
"""
|
||||
Initialize the tool executor.
|
||||
|
||||
Args:
|
||||
runtime: The runtime environment
|
||||
timeout: Default timeout for tool execution in seconds
|
||||
max_retries: Number of retries on failure
|
||||
"""
|
||||
self.runtime = runtime
|
||||
self.timeout = timeout
|
||||
self.max_retries = max_retries
|
||||
self.execution_history: List[ExecutionResult] = []
|
||||
|
||||
async def execute(
|
||||
self, tool: "Tool", arguments: dict, timeout: Optional[int] = None
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute a tool with the given arguments.
|
||||
|
||||
Args:
|
||||
tool: The tool to execute
|
||||
arguments: The arguments to pass to the tool
|
||||
timeout: Optional timeout override
|
||||
|
||||
Returns:
|
||||
ExecutionResult with the outcome
|
||||
"""
|
||||
execution_timeout = timeout or self.timeout
|
||||
start_time = datetime.now()
|
||||
|
||||
# Validate arguments
|
||||
is_valid, error_msg = tool.validate_arguments(arguments)
|
||||
if not is_valid:
|
||||
result = ExecutionResult(
|
||||
tool_name=tool.name,
|
||||
arguments=arguments,
|
||||
error=error_msg,
|
||||
success=False,
|
||||
start_time=start_time,
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
result.duration_ms = (result.end_time - start_time).total_seconds() * 1000
|
||||
self.execution_history.append(result)
|
||||
return result
|
||||
|
||||
# Execute with retries
|
||||
last_error = None
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
# Execute with timeout
|
||||
output = await asyncio.wait_for(
|
||||
tool.execute(arguments, self.runtime), timeout=execution_timeout
|
||||
)
|
||||
|
||||
end_time = datetime.now()
|
||||
result = ExecutionResult(
|
||||
tool_name=tool.name,
|
||||
arguments=arguments,
|
||||
result=output,
|
||||
success=True,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
duration_ms=(end_time - start_time).total_seconds() * 1000,
|
||||
)
|
||||
self.execution_history.append(result)
|
||||
return result
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
last_error = f"Execution timed out after {execution_timeout} seconds"
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
|
||||
# Wait before retry
|
||||
if attempt < self.max_retries:
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# All attempts failed
|
||||
end_time = datetime.now()
|
||||
result = ExecutionResult(
|
||||
tool_name=tool.name,
|
||||
arguments=arguments,
|
||||
error=last_error,
|
||||
success=False,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
duration_ms=(end_time - start_time).total_seconds() * 1000,
|
||||
)
|
||||
self.execution_history.append(result)
|
||||
return result
|
||||
|
||||
async def execute_batch(
|
||||
self, executions: List[tuple["Tool", dict]], parallel: bool = False
|
||||
) -> List[ExecutionResult]:
|
||||
"""
|
||||
Execute multiple tools.
|
||||
|
||||
Args:
|
||||
executions: List of (tool, arguments) tuples
|
||||
parallel: Whether to execute in parallel
|
||||
|
||||
Returns:
|
||||
List of ExecutionResults
|
||||
"""
|
||||
if parallel:
|
||||
tasks = [self.execute(tool, args) for tool, args in executions]
|
||||
return await asyncio.gather(*tasks)
|
||||
else:
|
||||
results = []
|
||||
for tool, args in executions:
|
||||
result = await self.execute(tool, args)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def get_execution_stats(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get statistics about tool executions.
|
||||
|
||||
Returns:
|
||||
Dictionary with execution statistics
|
||||
"""
|
||||
if not self.execution_history:
|
||||
return {
|
||||
"total_executions": 0,
|
||||
"successful": 0,
|
||||
"failed": 0,
|
||||
"success_rate": 0.0,
|
||||
"avg_duration_ms": 0.0,
|
||||
"tools_used": {},
|
||||
}
|
||||
|
||||
total = len(self.execution_history)
|
||||
successful = sum(1 for r in self.execution_history if r.success)
|
||||
failed = total - successful
|
||||
|
||||
# Calculate average duration
|
||||
durations = [r.duration_ms for r in self.execution_history]
|
||||
avg_duration = sum(durations) / len(durations) if durations else 0.0
|
||||
|
||||
# Count tool usage
|
||||
tools_used: Dict[str, int] = {}
|
||||
for result in self.execution_history:
|
||||
tools_used[result.tool_name] = tools_used.get(result.tool_name, 0) + 1
|
||||
|
||||
return {
|
||||
"total_executions": total,
|
||||
"successful": successful,
|
||||
"failed": failed,
|
||||
"success_rate": successful / total if total > 0 else 0.0,
|
||||
"avg_duration_ms": avg_duration,
|
||||
"tools_used": tools_used,
|
||||
}
|
||||
|
||||
def clear_history(self):
|
||||
"""Clear the execution history."""
|
||||
self.execution_history.clear()
|
||||
|
||||
def get_last_result(
|
||||
self, tool_name: Optional[str] = None
|
||||
) -> Optional[ExecutionResult]:
|
||||
"""
|
||||
Get the last execution result.
|
||||
|
||||
Args:
|
||||
tool_name: Optional filter by tool name
|
||||
|
||||
Returns:
|
||||
The last ExecutionResult or None
|
||||
"""
|
||||
if not self.execution_history:
|
||||
return None
|
||||
|
||||
if tool_name:
|
||||
for result in reversed(self.execution_history):
|
||||
if result.tool_name == tool_name:
|
||||
return result
|
||||
return None
|
||||
|
||||
return self.execution_history[-1]
|
||||
143
ghostcrew/tools/loader.py
Normal file
143
ghostcrew/tools/loader.py
Normal file
@@ -0,0 +1,143 @@
|
||||
"""Dynamic tool loader for GhostCrew."""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from .registry import get_all_tools
|
||||
|
||||
|
||||
def discover_tools(tools_dir: Optional[Path] = None) -> List[str]:
|
||||
"""
|
||||
Discover all tool modules in the tools directory.
|
||||
|
||||
Args:
|
||||
tools_dir: Path to tools directory. Defaults to package tools dir.
|
||||
|
||||
Returns:
|
||||
List of discovered tool module names
|
||||
"""
|
||||
if tools_dir is None:
|
||||
tools_dir = Path(__file__).parent
|
||||
|
||||
discovered = []
|
||||
|
||||
for item in tools_dir.iterdir():
|
||||
# Skip non-tool items
|
||||
if item.name.startswith("_"):
|
||||
continue
|
||||
if item.name in ("registry.py", "executor.py", "loader.py"):
|
||||
continue
|
||||
|
||||
# Check if it's a tool module
|
||||
if item.is_dir() and (item / "__init__.py").exists():
|
||||
discovered.append(item.name)
|
||||
elif item.suffix == ".py":
|
||||
discovered.append(item.stem)
|
||||
|
||||
return discovered
|
||||
|
||||
|
||||
def load_tool_module(module_name: str, tools_dir: Optional[Path] = None) -> bool:
|
||||
"""
|
||||
Load a tool module by name.
|
||||
|
||||
Args:
|
||||
module_name: Name of the tool module to load
|
||||
tools_dir: Path to tools directory
|
||||
|
||||
Returns:
|
||||
True if loaded successfully, False otherwise
|
||||
"""
|
||||
if tools_dir is None:
|
||||
tools_dir = Path(__file__).parent
|
||||
|
||||
try:
|
||||
# Build the full module path
|
||||
full_module_name = f"ghostcrew.tools.{module_name}"
|
||||
|
||||
# Check if already loaded
|
||||
if full_module_name in sys.modules:
|
||||
return True
|
||||
|
||||
# Try to import the module
|
||||
importlib.import_module(full_module_name)
|
||||
return True
|
||||
|
||||
except ImportError as e:
|
||||
print(f"Warning: Failed to load tool module '{module_name}': {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Warning: Error loading tool module '{module_name}': {e}")
|
||||
return False
|
||||
|
||||
|
||||
def load_all_tools(tools_dir: Optional[Path] = None) -> List[str]:
|
||||
"""
|
||||
Discover and load all tool modules.
|
||||
|
||||
Args:
|
||||
tools_dir: Path to tools directory
|
||||
|
||||
Returns:
|
||||
List of successfully loaded tool module names
|
||||
"""
|
||||
discovered = discover_tools(tools_dir)
|
||||
loaded = []
|
||||
|
||||
for module_name in discovered:
|
||||
if load_tool_module(module_name, tools_dir):
|
||||
loaded.append(module_name)
|
||||
|
||||
return loaded
|
||||
|
||||
|
||||
def get_tool_info() -> List[dict]:
|
||||
"""
|
||||
Get information about all registered tools.
|
||||
|
||||
Returns:
|
||||
List of tool info dictionaries
|
||||
"""
|
||||
tools = get_all_tools()
|
||||
return [
|
||||
{
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"category": tool.category,
|
||||
"enabled": tool.enabled,
|
||||
"parameters": (
|
||||
list(tool.schema.properties.keys()) if tool.schema.properties else []
|
||||
),
|
||||
}
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
|
||||
def reload_tools():
|
||||
"""Reload all tool modules."""
|
||||
from .registry import clear_tools
|
||||
|
||||
# Clear existing tools
|
||||
clear_tools()
|
||||
|
||||
# Unload tool modules from sys.modules
|
||||
to_remove = [
|
||||
name
|
||||
for name in sys.modules
|
||||
if name.startswith("ghostcrew.tools.")
|
||||
and name
|
||||
not in (
|
||||
"ghostcrew.tools",
|
||||
"ghostcrew.tools.registry",
|
||||
"ghostcrew.tools.executor",
|
||||
"ghostcrew.tools.loader",
|
||||
)
|
||||
]
|
||||
|
||||
for name in to_remove:
|
||||
del sys.modules[name]
|
||||
|
||||
# Reload all tools
|
||||
return load_all_tools()
|
||||
140
ghostcrew/tools/notes/__init__.py
Normal file
140
ghostcrew/tools/notes/__init__.py
Normal file
@@ -0,0 +1,140 @@
|
||||
"""Notes tool for GhostCrew - persistent key findings storage."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
from ..registry import ToolSchema, register_tool
|
||||
|
||||
# Notes storage
|
||||
_notes: Dict[str, str] = {}
|
||||
_notes_file: Path = Path("loot/notes.json")
|
||||
|
||||
|
||||
def _load_notes() -> None:
|
||||
"""Load notes from file."""
|
||||
global _notes
|
||||
if _notes_file.exists():
|
||||
try:
|
||||
_notes = json.loads(_notes_file.read_text())
|
||||
except (json.JSONDecodeError, IOError):
|
||||
_notes = {}
|
||||
|
||||
|
||||
def _save_notes() -> None:
|
||||
"""Save notes to file."""
|
||||
_notes_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
_notes_file.write_text(json.dumps(_notes, indent=2))
|
||||
|
||||
|
||||
def get_all_notes() -> Dict[str, str]:
|
||||
"""Get all notes (for TUI /notes command)."""
|
||||
if not _notes:
|
||||
_load_notes()
|
||||
return _notes.copy()
|
||||
|
||||
|
||||
def set_notes_file(path: Path) -> None:
|
||||
"""Set custom notes file path."""
|
||||
global _notes_file
|
||||
_notes_file = path
|
||||
_load_notes()
|
||||
|
||||
|
||||
# Load notes on module import
|
||||
_load_notes()
|
||||
|
||||
|
||||
@register_tool(
|
||||
name="notes",
|
||||
description="Manage persistent notes for key findings. Actions: create, read, update, delete, list.",
|
||||
schema=ToolSchema(
|
||||
properties={
|
||||
"action": {
|
||||
"type": "string",
|
||||
"enum": ["create", "read", "update", "delete", "list"],
|
||||
"description": "The action to perform",
|
||||
},
|
||||
"key": {
|
||||
"type": "string",
|
||||
"description": "Note identifier (e.g., 'creds_ssh', 'open_ports', 'vuln_sqli')",
|
||||
},
|
||||
"value": {
|
||||
"type": "string",
|
||||
"description": "Note content (for create/update)",
|
||||
},
|
||||
},
|
||||
required=["action"],
|
||||
),
|
||||
category="utility",
|
||||
)
|
||||
async def notes(arguments: dict, runtime) -> str:
|
||||
"""
|
||||
Manage persistent notes.
|
||||
|
||||
Args:
|
||||
arguments: Dictionary with action, key, value
|
||||
runtime: The runtime environment (unused)
|
||||
|
||||
Returns:
|
||||
Result message
|
||||
"""
|
||||
action = arguments["action"]
|
||||
key = arguments.get("key", "").strip()
|
||||
value = arguments.get("value", "")
|
||||
|
||||
if action == "create":
|
||||
if not key:
|
||||
return "Error: key is required for create"
|
||||
if not value:
|
||||
return "Error: value is required for create"
|
||||
if key in _notes:
|
||||
return f"Error: note '{key}' already exists. Use 'update' to modify."
|
||||
|
||||
_notes[key] = value
|
||||
_save_notes()
|
||||
return f"Created note '{key}'"
|
||||
|
||||
elif action == "read":
|
||||
if not key:
|
||||
return "Error: key is required for read"
|
||||
if key not in _notes:
|
||||
return f"Note '{key}' not found"
|
||||
|
||||
return f"[{key}] {_notes[key]}"
|
||||
|
||||
elif action == "update":
|
||||
if not key:
|
||||
return "Error: key is required for update"
|
||||
if not value:
|
||||
return "Error: value is required for update"
|
||||
|
||||
existed = key in _notes
|
||||
_notes[key] = value
|
||||
_save_notes()
|
||||
return f"{'Updated' if existed else 'Created'} note '{key}'"
|
||||
|
||||
elif action == "delete":
|
||||
if not key:
|
||||
return "Error: key is required for delete"
|
||||
if key not in _notes:
|
||||
return f"Note '{key}' not found"
|
||||
|
||||
del _notes[key]
|
||||
_save_notes()
|
||||
return f"Deleted note '{key}'"
|
||||
|
||||
elif action == "list":
|
||||
if not _notes:
|
||||
return "No notes saved"
|
||||
|
||||
lines = [f"Notes ({len(_notes)} entries):"]
|
||||
for k, v in _notes.items():
|
||||
# Truncate long values for display
|
||||
display_val = v if len(v) <= 60 else v[:57] + "..."
|
||||
lines.append(f" [{k}] {display_val}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
else:
|
||||
return f"Unknown action: {action}"
|
||||
235
ghostcrew/tools/registry.py
Normal file
235
ghostcrew/tools/registry.py
Normal file
@@ -0,0 +1,235 @@
|
||||
"""Tool registry for GhostCrew."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..runtime import Runtime
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolSchema:
|
||||
"""JSON Schema for tool parameters."""
|
||||
|
||||
type: str = "object"
|
||||
properties: Optional[Dict[str, Any]] = None
|
||||
required: Optional[List[str]] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.properties is None:
|
||||
self.properties = {}
|
||||
if self.required is None:
|
||||
self.required = []
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary format."""
|
||||
return {
|
||||
"type": self.type,
|
||||
"properties": self.properties,
|
||||
"required": self.required,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tool:
|
||||
"""Represents a tool available to agents."""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
schema: ToolSchema
|
||||
execute_fn: Callable
|
||||
category: str = "general"
|
||||
enabled: bool = True
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
async def execute(self, arguments: dict, runtime: "Runtime") -> str:
|
||||
"""
|
||||
Execute the tool with given arguments.
|
||||
|
||||
Args:
|
||||
arguments: The arguments for the tool
|
||||
runtime: The runtime environment
|
||||
|
||||
Returns:
|
||||
The tool execution result as a string
|
||||
"""
|
||||
if not self.enabled:
|
||||
return f"Tool '{self.name}' is currently disabled."
|
||||
|
||||
return await self.execute_fn(arguments, runtime)
|
||||
|
||||
def to_llm_format(self) -> dict:
|
||||
"""Convert to LLM function calling format."""
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"parameters": {
|
||||
"type": self.schema.type,
|
||||
"properties": self.schema.properties or {},
|
||||
"required": self.schema.required or [],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def validate_arguments(self, arguments: dict) -> tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Validate arguments against the schema.
|
||||
|
||||
Args:
|
||||
arguments: The arguments to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message)
|
||||
"""
|
||||
# Check required fields
|
||||
for required_field in self.schema.required or []:
|
||||
if required_field not in arguments:
|
||||
return False, f"Missing required field: {required_field}"
|
||||
|
||||
# Type checking (basic)
|
||||
for key, value in arguments.items():
|
||||
if key in (self.schema.properties or {}):
|
||||
expected_type = self.schema.properties[key].get("type")
|
||||
if expected_type:
|
||||
if not self._check_type(value, expected_type):
|
||||
return (
|
||||
False,
|
||||
f"Invalid type for {key}: expected {expected_type}",
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
def _check_type(self, value: Any, expected_type: str) -> bool:
|
||||
"""Check if a value matches the expected type."""
|
||||
type_map = {
|
||||
"string": str,
|
||||
"integer": int,
|
||||
"number": (int, float),
|
||||
"boolean": bool,
|
||||
"array": list,
|
||||
"object": dict,
|
||||
}
|
||||
|
||||
expected = type_map.get(expected_type)
|
||||
if expected is None:
|
||||
return True # Unknown type, allow
|
||||
|
||||
return isinstance(value, expected)
|
||||
|
||||
|
||||
# Global tool registry
|
||||
_tools: Dict[str, Tool] = {}
|
||||
|
||||
|
||||
def register_tool(
|
||||
name: str, description: str, schema: ToolSchema, category: str = "general"
|
||||
) -> Callable:
|
||||
"""
|
||||
Decorator to register a tool.
|
||||
|
||||
Args:
|
||||
name: The tool name
|
||||
description: Description of what the tool does
|
||||
schema: The parameter schema
|
||||
category: Tool category for organization
|
||||
|
||||
Returns:
|
||||
Decorator function
|
||||
"""
|
||||
|
||||
def decorator(fn: Callable) -> Callable:
|
||||
@wraps(fn)
|
||||
async def wrapper(arguments: dict, runtime: "Runtime") -> str:
|
||||
return await fn(arguments, runtime)
|
||||
|
||||
tool = Tool(
|
||||
name=name,
|
||||
description=description,
|
||||
schema=schema,
|
||||
execute_fn=wrapper,
|
||||
category=category,
|
||||
)
|
||||
_tools[name] = tool
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def get_all_tools() -> List[Tool]:
|
||||
"""Get all registered tools."""
|
||||
return list(_tools.values())
|
||||
|
||||
|
||||
def get_tool(name: str) -> Optional[Tool]:
|
||||
"""Get a tool by name."""
|
||||
return _tools.get(name)
|
||||
|
||||
|
||||
def register_tool_instance(tool: Tool) -> None:
|
||||
"""
|
||||
Register a pre-built tool instance (used for MCP tools).
|
||||
|
||||
Args:
|
||||
tool: The Tool instance to register
|
||||
"""
|
||||
_tools[tool.name] = tool
|
||||
|
||||
|
||||
def unregister_tool(name: str) -> bool:
|
||||
"""
|
||||
Unregister a tool by name.
|
||||
|
||||
Args:
|
||||
name: The tool name to unregister
|
||||
|
||||
Returns:
|
||||
True if tool was unregistered, False if not found
|
||||
"""
|
||||
if name in _tools:
|
||||
del _tools[name]
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_tools_by_category(category: str) -> List[Tool]:
|
||||
"""
|
||||
Get all tools in a specific category.
|
||||
|
||||
Args:
|
||||
category: The category to filter by
|
||||
|
||||
Returns:
|
||||
List of tools in that category
|
||||
"""
|
||||
return [tool for tool in _tools.values() if tool.category == category]
|
||||
|
||||
|
||||
def clear_tools() -> None:
|
||||
"""Clear all registered tools."""
|
||||
_tools.clear()
|
||||
|
||||
|
||||
def get_tool_names() -> List[str]:
|
||||
"""Get list of all registered tool names."""
|
||||
return list(_tools.keys())
|
||||
|
||||
|
||||
def enable_tool(name: str) -> bool:
|
||||
"""Enable a tool by name."""
|
||||
tool = _tools.get(name)
|
||||
if tool:
|
||||
tool.enabled = True
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def disable_tool(name: str) -> bool:
|
||||
"""Disable a tool by name."""
|
||||
tool = _tools.get(name)
|
||||
if tool:
|
||||
tool.enabled = False
|
||||
return True
|
||||
return False
|
||||
74
ghostcrew/tools/terminal/__init__.py
Normal file
74
ghostcrew/tools/terminal/__init__.py
Normal file
@@ -0,0 +1,74 @@
|
||||
"""Terminal tool for GhostCrew."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..registry import ToolSchema, register_tool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...runtime import Runtime
|
||||
|
||||
|
||||
@register_tool(
|
||||
name="terminal",
|
||||
description="Execute shell commands.",
|
||||
schema=ToolSchema(
|
||||
properties={
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "The shell command to execute",
|
||||
},
|
||||
"timeout": {
|
||||
"type": "integer",
|
||||
"description": "Timeout in seconds (default: 300)",
|
||||
"default": 300,
|
||||
},
|
||||
"working_dir": {
|
||||
"type": "string",
|
||||
"description": "Working directory for the command (optional)",
|
||||
},
|
||||
},
|
||||
required=["command"],
|
||||
),
|
||||
category="execution",
|
||||
)
|
||||
async def terminal(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""
|
||||
Execute a terminal command in the sandbox.
|
||||
|
||||
Args:
|
||||
arguments: Dictionary with 'command', optional 'timeout' and 'working_dir'
|
||||
runtime: The runtime environment
|
||||
|
||||
Returns:
|
||||
Formatted output string with command results
|
||||
"""
|
||||
command = arguments["command"]
|
||||
timeout = arguments.get("timeout", 300)
|
||||
working_dir = arguments.get("working_dir")
|
||||
|
||||
# Build the full command with working directory if specified
|
||||
if working_dir:
|
||||
full_command = f"cd {working_dir} && {command}"
|
||||
else:
|
||||
full_command = command
|
||||
|
||||
result = await runtime.execute_command(full_command, timeout=timeout)
|
||||
|
||||
# Format the output
|
||||
output_parts = [f"Command: {command}"]
|
||||
|
||||
if working_dir:
|
||||
output_parts.append(f"Working Directory: {working_dir}")
|
||||
|
||||
output_parts.append(f"Exit Code: {result.exit_code}")
|
||||
|
||||
if result.stdout:
|
||||
output_parts.append(f"\n--- stdout ---\n{result.stdout}")
|
||||
|
||||
if result.stderr:
|
||||
output_parts.append(f"\n--- stderr ---\n{result.stderr}")
|
||||
|
||||
if not result.stdout and not result.stderr:
|
||||
output_parts.append("\n(No output)")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
5
ghostcrew/tools/terminal/terminal.py
Normal file
5
ghostcrew/tools/terminal/terminal.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Terminal tool implementation."""
|
||||
|
||||
from . import terminal
|
||||
|
||||
__all__ = ["terminal"]
|
||||
104
ghostcrew/tools/web_search/__init__.py
Normal file
104
ghostcrew/tools/web_search/__init__.py
Normal file
@@ -0,0 +1,104 @@
|
||||
"""Web search tool for GhostCrew."""
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
|
||||
from ..registry import ToolSchema, register_tool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...runtime import Runtime
|
||||
|
||||
|
||||
@register_tool(
|
||||
name="web_search",
|
||||
description="Search the web for new security research, CVEs, exploits, bypass techniques, and documentation.",
|
||||
schema=ToolSchema(
|
||||
properties={
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Search query (be specific - include CVE numbers, tool names, versions)",
|
||||
}
|
||||
},
|
||||
required=["query"],
|
||||
),
|
||||
category="research",
|
||||
)
|
||||
async def web_search(arguments: dict, runtime: "Runtime") -> str:
|
||||
"""
|
||||
Search the web using Tavily API.
|
||||
|
||||
Args:
|
||||
arguments: Dictionary with 'query'
|
||||
runtime: The runtime environment
|
||||
|
||||
Returns:
|
||||
Search results formatted for the LLM
|
||||
"""
|
||||
query = arguments.get("query", "").strip()
|
||||
if not query:
|
||||
return "Error: No search query provided"
|
||||
|
||||
api_key = os.getenv("TAVILY_API_KEY")
|
||||
if not api_key:
|
||||
return (
|
||||
"Error: TAVILY_API_KEY environment variable not set.\n"
|
||||
"Get a free API key at https://tavily.com (1000 searches/month free)"
|
||||
)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
"https://api.tavily.com/search",
|
||||
json={
|
||||
"api_key": api_key,
|
||||
"query": query,
|
||||
"search_depth": "advanced",
|
||||
"include_answer": True,
|
||||
"include_raw_content": False,
|
||||
"max_results": 5,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
return _format_results(query, data)
|
||||
|
||||
except httpx.TimeoutException:
|
||||
return "Error: Search request timed out"
|
||||
except httpx.HTTPStatusError as e:
|
||||
return f"Error: Search API returned {e.response.status_code}"
|
||||
except Exception as e:
|
||||
return f"Error: Search failed - {str(e)}"
|
||||
|
||||
|
||||
def _format_results(query: str, data: dict) -> str:
|
||||
"""Format Tavily results for LLM consumption."""
|
||||
parts = [f"Search: {query}\n"]
|
||||
|
||||
# Include synthesized answer if available
|
||||
if answer := data.get("answer"):
|
||||
parts.append(f"Summary:\n{answer}\n")
|
||||
|
||||
# Include individual results
|
||||
results = data.get("results", [])
|
||||
if results:
|
||||
parts.append("Sources:")
|
||||
for i, result in enumerate(results, 1):
|
||||
title = result.get("title", "Untitled")
|
||||
url = result.get("url", "")
|
||||
content = result.get("content", "")
|
||||
|
||||
# Truncate content if too long
|
||||
if len(content) > 500:
|
||||
content = content[:500] + "..."
|
||||
|
||||
parts.append(f"\n[{i}] {title}")
|
||||
parts.append(f" URL: {url}")
|
||||
if content:
|
||||
parts.append(f" {content}")
|
||||
else:
|
||||
parts.append("No results found")
|
||||
|
||||
return "\n".join(parts)
|
||||
@@ -1,34 +0,0 @@
|
||||
admin
|
||||
login
|
||||
dashboard
|
||||
test
|
||||
dev
|
||||
backup
|
||||
config
|
||||
old
|
||||
index.html
|
||||
index.php
|
||||
robots.txt
|
||||
.htaccess
|
||||
uploads
|
||||
testpass
|
||||
images
|
||||
js
|
||||
css
|
||||
api
|
||||
server-status
|
||||
cgi-bin
|
||||
cgi-local
|
||||
cgi-shl
|
||||
cgi
|
||||
scripts
|
||||
bin
|
||||
test-cgi
|
||||
nph-test
|
||||
nph-ping
|
||||
nph-trace
|
||||
formmail.cgi
|
||||
test.cgi
|
||||
guestbook.cgi
|
||||
count.cgi
|
||||
submit.cgi
|
||||
42
main.py
42
main.py
@@ -1,42 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
GHOSTCREW - AI-driven penetration testing assistant
|
||||
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from colorama import init
|
||||
|
||||
init(autoreset=True)
|
||||
|
||||
from agents import set_tracing_disabled
|
||||
set_tracing_disabled(True)
|
||||
|
||||
from agents.mcp import MCPServerStdio, MCPServerSse
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main application entry point."""
|
||||
try:
|
||||
from core.pentest_agent import PentestAgent
|
||||
|
||||
agent = PentestAgent(MCPServerStdio, MCPServerSse)
|
||||
await agent.run()
|
||||
|
||||
except ImportError as e:
|
||||
print(f"Error importing required modules: {e}")
|
||||
print("Please ensure all dependencies are installed: pip install -r requirements.txt")
|
||||
sys.exit(1)
|
||||
except KeyboardInterrupt:
|
||||
print("\nApplication interrupted by user.")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Unexpected error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
131
pyproject.toml
Normal file
131
pyproject.toml
Normal file
@@ -0,0 +1,131 @@
|
||||
[project]
|
||||
name = "ghostcrew"
|
||||
version = "0.2.0"
|
||||
description = "AI penetration testing"
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
requires-python = ">=3.10"
|
||||
authors = [
|
||||
{name = "Masic"}
|
||||
]
|
||||
keywords = [
|
||||
"penetration-testing",
|
||||
"security",
|
||||
"ai",
|
||||
"llm",
|
||||
"mcp",
|
||||
"automation"
|
||||
]
|
||||
classifiers = [
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Information Technology",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Operating System :: OS Independent",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Topic :: Security",
|
||||
"Topic :: Software Development :: Libraries :: Python Modules",
|
||||
]
|
||||
dependencies = [
|
||||
"litellm>=1.40.0",
|
||||
"openai>=1.30.0",
|
||||
"anthropic>=0.25.0",
|
||||
"tiktoken>=0.7.0",
|
||||
"aiohttp>=3.9.0",
|
||||
"aiofiles>=23.2.0",
|
||||
"playwright>=1.44.0",
|
||||
"beautifulsoup4>=4.12.0",
|
||||
"httpx>=0.27.0",
|
||||
"numpy>=1.26.0",
|
||||
"docker>=7.0.0",
|
||||
"rich>=13.7.0",
|
||||
"textual>=0.63.0",
|
||||
"typer>=0.12.0",
|
||||
"pydantic>=2.7.0",
|
||||
"pydantic-settings>=2.2.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
"pyyaml>=6.0.0",
|
||||
"jinja2>=3.1.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.2.0",
|
||||
"pytest-asyncio>=0.23.0",
|
||||
"pytest-cov>=5.0.0",
|
||||
"pytest-mock>=3.14.0",
|
||||
"black>=24.4.0",
|
||||
"isort>=5.13.0",
|
||||
"mypy>=1.10.0",
|
||||
"ruff>=0.4.0",
|
||||
]
|
||||
rag = [
|
||||
"sentence-transformers>=2.7.0",
|
||||
"faiss-cpu>=1.8.0",
|
||||
]
|
||||
all = [
|
||||
"ghostcrew[dev,rag]",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/GH05TCREW/ghostcrew"
|
||||
|
||||
[project.scripts]
|
||||
ghostcrew = "ghostcrew.interface.main:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["ghostcrew"]
|
||||
|
||||
[tool.hatch.build.targets.sdist]
|
||||
include = [
|
||||
"ghostcrew/**",
|
||||
"*.md",
|
||||
"*.txt"
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
python_files = ["test_*.py"]
|
||||
python_functions = ["test_*"]
|
||||
addopts = "-v --tb=short"
|
||||
|
||||
[tool.black]
|
||||
line-length = 88
|
||||
target-version = ["py310", "py311", "py312"]
|
||||
include = '\.pyi?$'
|
||||
|
||||
[tool.isort]
|
||||
profile = "black"
|
||||
line_length = 88
|
||||
known_first_party = ["ghostcrew"]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 88
|
||||
target-version = "py310"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
"F", # pyflakes
|
||||
"I", # isort
|
||||
"B", # flake8-bugbear
|
||||
"C4", # flake8-comprehensions
|
||||
]
|
||||
ignore = [
|
||||
"E501", # line too long (handled by black)
|
||||
]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.10"
|
||||
warn_return_any = true
|
||||
warn_unused_configs = true
|
||||
ignore_missing_imports = true
|
||||
@@ -1 +0,0 @@
|
||||
"""RAG (Retrieval-Augmented Generation) system for GHOSTCREW."""
|
||||
@@ -1,25 +0,0 @@
|
||||
import os
|
||||
import json
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env file
|
||||
load_dotenv()
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"), # Use the standard OpenAI API key env variable
|
||||
base_url=os.getenv("OPENAI_BASE_URL") # Read base_url from environment variable
|
||||
)
|
||||
|
||||
completion = client.embeddings.create(
|
||||
model="text-embedding-ada-002",
|
||||
input='This is a sample text for embedding generation to test the functionality.',
|
||||
encoding_format="float"
|
||||
)
|
||||
|
||||
response_json = completion.model_dump_json()
|
||||
embedding_data = json.loads(response_json)
|
||||
embedding_array = embedding_data['data'][0]['embedding']
|
||||
print(len(embedding_array))
|
||||
print(type(embedding_array))
|
||||
print("Extracted embedding array:", embedding_array)
|
||||
@@ -1,136 +0,0 @@
|
||||
#from curses import color_content
|
||||
from ollama import chat,Message
|
||||
from ollama import embeddings
|
||||
import os
|
||||
import json
|
||||
from openai import OpenAI
|
||||
import numpy as np
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env file
|
||||
load_dotenv()
|
||||
|
||||
# Set numpy print options to display full arrays
|
||||
np.set_printoptions(threshold=np.inf)
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"), # Use the standard OpenAI API key env variable
|
||||
base_url=os.getenv("OPENAI_BASE_URL") # Read base_url from environment variable
|
||||
)
|
||||
|
||||
import os # Added for directory operations
|
||||
|
||||
class Kb:
|
||||
def __init__(self, dirpath): # Read all documents in the directory
|
||||
all_content = ""
|
||||
if not os.path.isdir(dirpath):
|
||||
print(f"Error: {dirpath} is not a valid directory.")
|
||||
self.docs = []
|
||||
self.embedss = np.array([])
|
||||
return
|
||||
|
||||
# Define binary file extensions to skip
|
||||
binary_extensions = {
|
||||
'.exe', '.dll', '.so', '.dylib', '.elf', '.bin', '.dat',
|
||||
'.zip', '.tar', '.gz', '.7z', '.rar', '.pdf', '.doc', '.docx',
|
||||
'.xls', '.xlsx', '.ppt', '.pptx', '.jpg', '.jpeg', '.png', '.gif',
|
||||
'.bmp', '.ico', '.mp3', '.mp4', '.avi', '.mov', '.wav', '.flv',
|
||||
'.iso', '.img', '.vmdk', '.vdi'
|
||||
}
|
||||
|
||||
for filename in os.listdir(dirpath):
|
||||
filepath = os.path.join(dirpath, filename)
|
||||
if os.path.isfile(filepath):
|
||||
# Get file extension in lowercase
|
||||
file_ext = os.path.splitext(filename)[1].lower()
|
||||
|
||||
# Skip binary files
|
||||
if file_ext in binary_extensions:
|
||||
continue
|
||||
|
||||
try:
|
||||
with open(filepath, 'r', encoding="utf-8") as f:
|
||||
all_content += f.read() + "\n" # Add a newline to separate file contents
|
||||
except Exception as e:
|
||||
print(f"Error reading file {filepath}: {e}")
|
||||
|
||||
if not all_content.strip():
|
||||
print(f"Warning: No content found in directory {dirpath}.")
|
||||
self.docs = []
|
||||
self.embedss = np.array([])
|
||||
return
|
||||
|
||||
self.docs = self.split_content(all_content) # Split all document content after merging
|
||||
if self.docs:
|
||||
self.embedss = self.encode(self.docs)
|
||||
else:
|
||||
self.embedss = np.array([])
|
||||
|
||||
@staticmethod
|
||||
def split_content(content,max_length=5000):
|
||||
chuncks=[]
|
||||
for i in range(0,len(content),max_length):
|
||||
chuncks.append(content[i:i+max_length])
|
||||
return chuncks
|
||||
|
||||
|
||||
def encode(self,texts):
|
||||
embeds=[]
|
||||
for text in texts:
|
||||
completion = client.embeddings.create(
|
||||
model="text-embedding-ada-002",
|
||||
input=text,
|
||||
encoding_format="float"
|
||||
)
|
||||
response_json = completion.model_dump_json()
|
||||
embedding_data = json.loads(response_json)
|
||||
embedding_array = embedding_data['data'][0]['embedding']
|
||||
embeds.append(embedding_array)
|
||||
return np.array(embeds)
|
||||
|
||||
|
||||
@staticmethod #similarity
|
||||
def similarity(A,B):
|
||||
dot_product=np.dot(A,B)
|
||||
norm_A=np.linalg.norm(A)
|
||||
norm_B=np.linalg.norm(B)
|
||||
similarity=dot_product/(norm_A*norm_B)
|
||||
return similarity
|
||||
|
||||
def search(self,query):
|
||||
max_similarity=0
|
||||
max_similarity_index=0
|
||||
query_embedding=self.encode([query])[0]
|
||||
for idx,te in enumerate(self.embedss):
|
||||
similarity=self.similarity(query_embedding,te)
|
||||
if similarity>max_similarity:
|
||||
max_similarity=similarity
|
||||
max_similarity_index=idx
|
||||
return self.docs[max_similarity_index]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage: Create a dummy directory and file for testing
|
||||
test_kb_dir = "knowledge_test"
|
||||
if not os.path.exists(test_kb_dir):
|
||||
os.makedirs(test_kb_dir)
|
||||
with open(os.path.join(test_kb_dir, "test_doc.txt"), 'w', encoding='utf-8') as f:
|
||||
f.write("This is a test document for security audit information.")
|
||||
|
||||
kb = Kb(test_kb_dir)
|
||||
if kb.docs: # Check if docs were loaded
|
||||
#for doc in kb.docs:
|
||||
# print("========================================================")
|
||||
# print(doc)
|
||||
|
||||
#for e in kb.embedss:
|
||||
# print(e)
|
||||
result = kb.search("security audit")
|
||||
print(f"Search result: {result}")
|
||||
else:
|
||||
print("Knowledge base is empty or failed to load.")
|
||||
|
||||
# Clean up dummy directory and file
|
||||
# import shutil
|
||||
# if os.path.exists(test_kb_dir):
|
||||
# shutil.rmtree(test_kb_dir)
|
||||
@@ -1 +0,0 @@
|
||||
"""Report generation system for GHOSTCREW."""
|
||||
@@ -1,817 +0,0 @@
|
||||
"""
|
||||
Professional Markdown Report Generator for GHOSTCREW
|
||||
Processes workflow conversation history and generates beautiful reports
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Any
|
||||
import re
|
||||
from colorama import Fore, Style
|
||||
|
||||
|
||||
class PentestReportGenerator:
|
||||
"""Generate professional penetration testing reports from workflow data"""
|
||||
|
||||
def __init__(self, report_data: Dict[str, Any]):
|
||||
self.workflow_name = report_data['workflow_name']
|
||||
self.workflow_key = report_data['workflow_key']
|
||||
self.target = report_data['target']
|
||||
self.timestamp = report_data['timestamp']
|
||||
self.conversation_history = report_data['conversation_history']
|
||||
self.tools_used = report_data.get('tools_used', [])
|
||||
|
||||
# Will be populated by AI analysis
|
||||
self.structured_findings = {}
|
||||
|
||||
def format_conversation_history(self) -> str:
|
||||
"""Format conversation history for AI analysis"""
|
||||
formatted = []
|
||||
|
||||
for i, entry in enumerate(self.conversation_history, 1):
|
||||
formatted.append(f"\n--- STEP {i} ---")
|
||||
formatted.append(f"QUERY: {entry.get('user_query', '')}")
|
||||
formatted.append(f"RESPONSE: {entry.get('ai_response', '')}")
|
||||
formatted.append("=" * 50)
|
||||
|
||||
return "\n".join(formatted)
|
||||
|
||||
def create_analysis_prompt(self) -> str:
|
||||
"""Create comprehensive analysis prompt for AI"""
|
||||
|
||||
prompt = f"""
|
||||
You are analyzing a complete penetration testing workflow to create a professional security assessment report.
|
||||
|
||||
ASSESSMENT DETAILS:
|
||||
- Workflow: {self.workflow_name}
|
||||
- Target: {self.target}
|
||||
- Date: {self.timestamp.strftime('%Y-%m-%d')}
|
||||
- Tools Available: {', '.join(self.tools_used) if self.tools_used else 'Various security tools'}
|
||||
|
||||
COMPLETE WORKFLOW CONVERSATION LOG:
|
||||
{self.format_conversation_history()}
|
||||
|
||||
Please analyze this entire workflow and extract structured information. When extracting evidence, include the actual commands used and their outputs. Respond with a JSON object containing:
|
||||
|
||||
{{
|
||||
"executive_summary": "2-3 paragraph executive summary focusing on business impact and key risks",
|
||||
"key_statistics": {{
|
||||
"total_vulnerabilities": 0,
|
||||
"critical_count": 0,
|
||||
"high_count": 0,
|
||||
"systems_tested": 0,
|
||||
"systems_compromised": 0
|
||||
}},
|
||||
"vulnerabilities": [
|
||||
{{
|
||||
"severity": "Critical|High|Medium|Low|Informational",
|
||||
"title": "Descriptive vulnerability title",
|
||||
"description": "Technical description of the vulnerability",
|
||||
"impact": "Business impact if exploited",
|
||||
"affected_systems": ["IP/hostname"],
|
||||
"remediation": "Specific remediation steps",
|
||||
"evidence": "Include actual commands used and key outputs that demonstrate this finding. Format as: 'Command: [command] Output: [relevant output]'",
|
||||
"cvss_score": "If applicable",
|
||||
"references": ["CVE numbers, links, etc."]
|
||||
}}
|
||||
],
|
||||
"compromised_systems": [
|
||||
{{
|
||||
"system": "IP/hostname",
|
||||
"access_level": "user|admin|root|system",
|
||||
"method": "How access was gained",
|
||||
"evidence": "Actual commands and outputs showing compromise"
|
||||
}}
|
||||
],
|
||||
"credentials_found": [
|
||||
{{
|
||||
"username": "username",
|
||||
"credential_type": "password|hash|token|key",
|
||||
"system": "where found",
|
||||
"strength": "weak|moderate|strong"
|
||||
}}
|
||||
],
|
||||
"tools_used": ["List of tools actually used in testing"],
|
||||
"attack_paths": [
|
||||
{{
|
||||
"path_description": "Description of attack chain",
|
||||
"steps": ["Step 1", "Step 2", "etc"],
|
||||
"impact": "What this path achieves"
|
||||
}}
|
||||
],
|
||||
"recommendations": [
|
||||
{{
|
||||
"priority": "Immediate|Short-term|Medium-term|Long-term",
|
||||
"category": "Network|Application|System|Process",
|
||||
"recommendation": "Specific actionable recommendation",
|
||||
"business_justification": "Why this is important for business"
|
||||
}}
|
||||
],
|
||||
"methodology": "Brief description of testing methodology used",
|
||||
"scope_limitations": "Any scope limitations or areas not tested",
|
||||
"conclusion": "Overall security posture assessment and key takeaways"
|
||||
}}
|
||||
|
||||
Focus on extracting real findings from the conversation. Include actual command examples and outputs in the evidence fields. If no vulnerabilities were found, that's also valuable information. Be accurate and professional.
|
||||
"""
|
||||
return prompt
|
||||
|
||||
async def analyze_with_ai(self, prompt: str, run_agent_func, connected_servers, kb_instance=None):
|
||||
"""Run AI analysis using the main agent function"""
|
||||
try:
|
||||
# Use streaming=True since run_agent doesn't properly handle streaming=False
|
||||
result = await run_agent_func(
|
||||
prompt,
|
||||
connected_servers,
|
||||
history=[],
|
||||
streaming=True,
|
||||
kb_instance=kb_instance
|
||||
)
|
||||
|
||||
if result and hasattr(result, "final_output"):
|
||||
return result.final_output
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in AI analysis: {e}")
|
||||
return None
|
||||
|
||||
def parse_ai_response(self, ai_response: str) -> Dict[str, Any]:
|
||||
"""Parse AI response and extract JSON data"""
|
||||
try:
|
||||
# Try to find JSON in the response
|
||||
json_start = ai_response.find('{')
|
||||
json_end = ai_response.rfind('}') + 1
|
||||
|
||||
if json_start != -1 and json_end != -1:
|
||||
json_str = ai_response[json_start:json_end]
|
||||
return json.loads(json_str)
|
||||
else:
|
||||
# Fallback - create basic structure
|
||||
return {
|
||||
"executive_summary": ai_response[:500] + "...",
|
||||
"key_statistics": {"total_vulnerabilities": 0},
|
||||
"vulnerabilities": [],
|
||||
"compromised_systems": [],
|
||||
"recommendations": [],
|
||||
"methodology": "Standard penetration testing methodology",
|
||||
"conclusion": "Assessment completed successfully."
|
||||
}
|
||||
|
||||
except json.JSONDecodeError:
|
||||
# Fallback structure
|
||||
return {
|
||||
"executive_summary": "Assessment completed. See technical findings for details.",
|
||||
"key_statistics": {"total_vulnerabilities": 0},
|
||||
"vulnerabilities": [],
|
||||
"compromised_systems": [],
|
||||
"recommendations": [],
|
||||
"methodology": "Standard penetration testing methodology",
|
||||
"conclusion": "Unable to parse detailed findings."
|
||||
}
|
||||
|
||||
def generate_markdown_report(self) -> str:
|
||||
"""Generate the final markdown report"""
|
||||
findings = self.structured_findings
|
||||
|
||||
report = []
|
||||
|
||||
# Title Page
|
||||
report.append(f"# Penetration Testing Report")
|
||||
report.append(f"\n## {self.workflow_name}")
|
||||
report.append(f"\n**Target:** {self.target} ")
|
||||
report.append(f"**Assessment Date:** {self.timestamp.strftime('%Y-%m-%d')} ")
|
||||
report.append(f"**Report Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ")
|
||||
report.append(f"**Report ID:** GHOSTCREW-{self.workflow_key}-{int(self.timestamp.timestamp())} ")
|
||||
report.append(f"\n---\n")
|
||||
|
||||
# Table of Contents
|
||||
report.append("## Table of Contents\n")
|
||||
report.append("1. [Executive Summary](#1-executive-summary)")
|
||||
report.append("2. [Assessment Overview](#2-assessment-overview)")
|
||||
report.append("3. [Key Findings](#3-key-findings)")
|
||||
report.append("4. [Vulnerability Details](#4-vulnerability-details)")
|
||||
report.append("5. [Compromised Systems](#5-compromised-systems)")
|
||||
report.append("6. [Attack Paths](#6-attack-paths)")
|
||||
report.append("7. [Recommendations](#7-recommendations)")
|
||||
report.append("8. [Technical Methodology](#8-technical-methodology)")
|
||||
report.append("9. [Conclusion](#9-conclusion)")
|
||||
report.append("\n---\n")
|
||||
|
||||
# Executive Summary
|
||||
report.append("## 1. Executive Summary\n")
|
||||
report.append(findings.get('executive_summary', 'Assessment completed successfully.'))
|
||||
report.append("\n---\n")
|
||||
|
||||
# Assessment Overview
|
||||
report.append("## 2. Assessment Overview\n")
|
||||
report.append(f"### Scope")
|
||||
report.append(f"- **Primary Target:** {self.target}")
|
||||
report.append(f"- **Assessment Type:** {self.workflow_name}")
|
||||
report.append(f"- **Testing Window:** {self.timestamp.strftime('%Y-%m-%d')}")
|
||||
|
||||
if self.tools_used:
|
||||
report.append(f"\n### Tools Used")
|
||||
for tool in self.tools_used:
|
||||
report.append(f"- {tool}")
|
||||
|
||||
stats = findings.get('key_statistics', {})
|
||||
if stats:
|
||||
report.append(f"\n### Key Statistics")
|
||||
report.append(f"- **Total Vulnerabilities:** {stats.get('total_vulnerabilities', 0)}")
|
||||
report.append(f"- **Critical Severity:** {stats.get('critical_count', 0)}")
|
||||
report.append(f"- **High Severity:** {stats.get('high_count', 0)}")
|
||||
report.append(f"- **Systems Compromised:** {stats.get('systems_compromised', 0)}")
|
||||
|
||||
report.append("\n---\n")
|
||||
|
||||
# Key Findings Summary
|
||||
report.append("## 3. Key Findings\n")
|
||||
|
||||
vulnerabilities = findings.get('vulnerabilities', [])
|
||||
if vulnerabilities:
|
||||
# Group by severity
|
||||
severity_groups = {'Critical': [], 'High': [], 'Medium': [], 'Low': [], 'Informational': []}
|
||||
for vuln in vulnerabilities:
|
||||
severity = vuln.get('severity', 'Low')
|
||||
if severity in severity_groups:
|
||||
severity_groups[severity].append(vuln)
|
||||
|
||||
report.append("### Vulnerability Summary\n")
|
||||
report.append("| Severity | Count | Description |")
|
||||
report.append("|----------|-------|-------------|")
|
||||
|
||||
for severity, vulns in severity_groups.items():
|
||||
if vulns:
|
||||
count = len(vulns)
|
||||
titles = [v.get('title', 'Unknown') for v in vulns[:3]]
|
||||
desc = ', '.join(titles)
|
||||
if len(vulns) > 3:
|
||||
desc += f' (and {len(vulns) - 3} more)'
|
||||
report.append(f"| {severity} | {count} | {desc} |")
|
||||
else:
|
||||
report.append("No significant vulnerabilities were identified during the assessment.")
|
||||
|
||||
report.append("\n---\n")
|
||||
|
||||
# Vulnerability Details
|
||||
report.append("## 4. Vulnerability Details\n")
|
||||
|
||||
if vulnerabilities:
|
||||
# Group by severity for detailed listing
|
||||
for severity in ['Critical', 'High', 'Medium', 'Low', 'Informational']:
|
||||
severity_vulns = [v for v in vulnerabilities if v.get('severity') == severity]
|
||||
|
||||
if severity_vulns:
|
||||
report.append(f"### {severity} Severity Vulnerabilities\n")
|
||||
|
||||
for i, vuln in enumerate(severity_vulns, 1):
|
||||
report.append(f"#### {severity.upper()}-{i:03d}: {vuln.get('title', 'Unknown Vulnerability')}\n")
|
||||
report.append(f"**Description:** {vuln.get('description', 'No description provided')}\n")
|
||||
report.append(f"**Impact:** {vuln.get('impact', 'Impact assessment pending')}\n")
|
||||
|
||||
if vuln.get('affected_systems'):
|
||||
report.append(f"**Affected Systems:** {', '.join(vuln['affected_systems'])}\n")
|
||||
|
||||
report.append(f"**Remediation:** {vuln.get('remediation', 'Remediation steps pending')}\n")
|
||||
|
||||
if vuln.get('evidence'):
|
||||
report.append(f"**Evidence:**")
|
||||
report.append("```")
|
||||
report.append(vuln['evidence'])
|
||||
report.append("```")
|
||||
|
||||
if vuln.get('references'):
|
||||
report.append(f"**References:** {', '.join(vuln['references'])}\n")
|
||||
|
||||
report.append("\n")
|
||||
else:
|
||||
report.append("No vulnerabilities were identified during this assessment.")
|
||||
|
||||
report.append("---\n")
|
||||
|
||||
# Compromised Systems
|
||||
report.append("## 5. Compromised Systems\n")
|
||||
|
||||
compromised = findings.get('compromised_systems', [])
|
||||
if compromised:
|
||||
report.append("| System | Access Level | Method | Evidence |")
|
||||
report.append("|--------|--------------|--------|----------|")
|
||||
|
||||
for system in compromised:
|
||||
report.append(f"| {system.get('system', 'Unknown')} | {system.get('access_level', 'Unknown')} | {system.get('method', 'Unknown')} | {system.get('evidence', 'See technical details')[:50]}{'...' if len(system.get('evidence', '')) > 50 else ''} |")
|
||||
else:
|
||||
report.append("No systems were successfully compromised during the assessment.")
|
||||
|
||||
report.append("\n---\n")
|
||||
|
||||
# Attack Paths
|
||||
report.append("## 6. Attack Paths\n")
|
||||
|
||||
attack_paths = findings.get('attack_paths', [])
|
||||
if attack_paths:
|
||||
for i, path in enumerate(attack_paths, 1):
|
||||
report.append(f"### Attack Path {i}: {path.get('path_description', 'Unknown Path')}\n")
|
||||
report.append(f"**Impact:** {path.get('impact', 'Unknown impact')}\n")
|
||||
|
||||
steps = path.get('steps', [])
|
||||
if steps:
|
||||
report.append("**Steps:**")
|
||||
for step_num, step in enumerate(steps, 1):
|
||||
report.append(f"{step_num}. {step}")
|
||||
report.append("\n")
|
||||
else:
|
||||
report.append("No specific attack paths were identified or documented.")
|
||||
|
||||
report.append("---\n")
|
||||
|
||||
# Recommendations
|
||||
report.append("## 7. Recommendations\n")
|
||||
|
||||
recommendations = findings.get('recommendations', [])
|
||||
if recommendations:
|
||||
# Group by priority
|
||||
priority_groups = {'Immediate': [], 'Short-term': [], 'Medium-term': [], 'Long-term': []}
|
||||
for rec in recommendations:
|
||||
priority = rec.get('priority', 'Medium-term')
|
||||
if priority in priority_groups:
|
||||
priority_groups[priority].append(rec)
|
||||
|
||||
for priority, recs in priority_groups.items():
|
||||
if recs:
|
||||
report.append(f"### {priority} Priority\n")
|
||||
for rec in recs:
|
||||
report.append(f"**{rec.get('category', 'General')}:** {rec.get('recommendation', 'No recommendation provided')}")
|
||||
if rec.get('business_justification'):
|
||||
report.append(f" \n*Business Justification:* {rec['business_justification']}")
|
||||
report.append("\n")
|
||||
else:
|
||||
report.append("Continue following security best practices and conduct regular assessments.")
|
||||
|
||||
report.append("---\n")
|
||||
|
||||
# Technical Methodology
|
||||
report.append("## 8. Technical Methodology\n")
|
||||
report.append(findings.get('methodology', 'Standard penetration testing methodology was followed.'))
|
||||
|
||||
if findings.get('scope_limitations'):
|
||||
report.append(f"\n### Scope Limitations\n")
|
||||
report.append(findings['scope_limitations'])
|
||||
|
||||
report.append("\n---\n")
|
||||
|
||||
# Conclusion
|
||||
report.append("## 9. Conclusion\n")
|
||||
report.append(findings.get('conclusion', 'Assessment completed successfully.'))
|
||||
|
||||
report.append(f"\n\n---\n")
|
||||
report.append(f"*Report generated by GHOSTCREW v0.1.0* ")
|
||||
report.append(f"*{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*")
|
||||
|
||||
return "\n".join(report)
|
||||
|
||||
async def generate_report(self, run_agent_func, connected_servers, kb_instance=None, save_raw_history=False) -> str:
|
||||
"""Main method to generate the complete report"""
|
||||
|
||||
print(f"Analyzing workflow findings...")
|
||||
|
||||
# Step 1: Create analysis prompt
|
||||
analysis_prompt = self.create_analysis_prompt()
|
||||
|
||||
# Step 2: Get AI analysis
|
||||
ai_response = await self.analyze_with_ai(
|
||||
analysis_prompt,
|
||||
run_agent_func,
|
||||
connected_servers,
|
||||
kb_instance
|
||||
)
|
||||
|
||||
if not ai_response:
|
||||
raise Exception("Failed to get AI analysis")
|
||||
|
||||
print(f"Extracting structured findings...")
|
||||
|
||||
# Step 3: Parse AI response
|
||||
self.structured_findings = self.parse_ai_response(ai_response)
|
||||
|
||||
print(f"Generating markdown report...")
|
||||
|
||||
# Step 4: Generate markdown report
|
||||
markdown_report = self.generate_markdown_report()
|
||||
|
||||
# Step 5: Save report with options
|
||||
report_filename = self.save_report(markdown_report, save_raw_history)
|
||||
|
||||
return report_filename
|
||||
|
||||
def save_report(self, markdown_content: str, save_raw_history: bool = False) -> str:
|
||||
"""Save the report to file with optional raw history"""
|
||||
|
||||
# Create reports directory if it doesn't exist
|
||||
reports_dir = "reports"
|
||||
if not os.path.exists(reports_dir):
|
||||
os.makedirs(reports_dir)
|
||||
|
||||
# Generate filename
|
||||
timestamp_str = str(int(self.timestamp.timestamp()))
|
||||
filename = f"{reports_dir}/ghostcrew_{self.workflow_key}_{timestamp_str}.md"
|
||||
|
||||
# Save markdown file
|
||||
with open(filename, 'w', encoding='utf-8') as f:
|
||||
f.write(markdown_content)
|
||||
|
||||
# Optionally save raw conversation history
|
||||
if save_raw_history:
|
||||
raw_history_content = []
|
||||
raw_history_content.append(f"GHOSTCREW Raw Workflow History")
|
||||
raw_history_content.append(f"Workflow: {self.workflow_name}")
|
||||
raw_history_content.append(f"Target: {self.target}")
|
||||
raw_history_content.append(f"Date: {self.timestamp.strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
raw_history_content.append(f"=" * 60)
|
||||
raw_history_content.append("")
|
||||
|
||||
for i, entry in enumerate(self.conversation_history, 1):
|
||||
raw_history_content.append(f"STEP {i} - QUERY:")
|
||||
raw_history_content.append("-" * 40)
|
||||
raw_history_content.append(entry.get('user_query', 'No query recorded'))
|
||||
raw_history_content.append("")
|
||||
raw_history_content.append(f"STEP {i} - AI RESPONSE:")
|
||||
raw_history_content.append("-" * 40)
|
||||
raw_history_content.append(entry.get('ai_response', 'No response recorded'))
|
||||
raw_history_content.append("")
|
||||
raw_history_content.append("=" * 60)
|
||||
raw_history_content.append("")
|
||||
|
||||
raw_filename = f"{reports_dir}/ghostcrew_{self.workflow_key}_{timestamp_str}_raw_history.txt"
|
||||
with open(raw_filename, 'w', encoding='utf-8') as f:
|
||||
f.write('\n'.join(raw_history_content))
|
||||
print(f"Raw conversation history saved: {raw_filename}")
|
||||
|
||||
return filename
|
||||
|
||||
|
||||
async def generate_report_from_workflow(report_data: Dict[str, Any], run_agent_func, connected_servers, kb_instance=None, save_raw_history=False) -> str:
|
||||
"""
|
||||
Main function to generate a professional report from workflow data
|
||||
|
||||
Args:
|
||||
report_data: Dictionary containing workflow information
|
||||
run_agent_func: The main agent function for AI analysis
|
||||
connected_servers: Connected MCP servers
|
||||
kb_instance: Knowledge base instance
|
||||
save_raw_history: Whether to save raw conversation history
|
||||
|
||||
Returns:
|
||||
str: Path to generated report file
|
||||
"""
|
||||
|
||||
generator = PentestReportGenerator(report_data)
|
||||
return await generator.generate_report(run_agent_func, connected_servers, kb_instance, save_raw_history)
|
||||
|
||||
|
||||
async def generate_report_from_ptt(ptt_manager, conversation_history: List[Dict[str, Any]], run_agent_func=None, connected_servers=None, kb_instance=None, save_raw_history=False) -> str:
|
||||
"""
|
||||
Generate a professional report from PTT (Pentesting Task Tree) data
|
||||
|
||||
Args:
|
||||
ptt_manager: TaskTreeManager instance containing the PTT
|
||||
conversation_history: List of conversation history entries
|
||||
run_agent_func: The main agent function for AI analysis
|
||||
connected_servers: Connected MCP servers
|
||||
kb_instance: Knowledge base instance
|
||||
save_raw_history: Whether to save raw conversation history
|
||||
|
||||
Returns:
|
||||
str: Path to generated report file
|
||||
"""
|
||||
|
||||
# Convert PTT data to report-compatible format
|
||||
report_data = {
|
||||
'workflow_name': f"Agent Mode: {ptt_manager.goal}",
|
||||
'workflow_key': 'agent_mode',
|
||||
'target': ptt_manager.target,
|
||||
'timestamp': ptt_manager.creation_time,
|
||||
'conversation_history': conversation_history,
|
||||
'tools_used': [server.name for server in connected_servers] if connected_servers else [],
|
||||
'ptt_data': {
|
||||
'goal': ptt_manager.goal,
|
||||
'target': ptt_manager.target,
|
||||
'constraints': ptt_manager.constraints,
|
||||
'statistics': ptt_manager.get_statistics(),
|
||||
'tree_structure': ptt_manager.to_natural_language(),
|
||||
'nodes': {node_id: node.to_dict() for node_id, node in ptt_manager.nodes.items()}
|
||||
}
|
||||
}
|
||||
|
||||
# Create a specialized PTT report generator
|
||||
generator = PTTReportGenerator(report_data)
|
||||
|
||||
if run_agent_func and connected_servers:
|
||||
return await generator.generate_report(run_agent_func, connected_servers, kb_instance, save_raw_history)
|
||||
else:
|
||||
# Generate a basic report without AI analysis if no agent function available
|
||||
return generator.generate_basic_report(save_raw_history)
|
||||
|
||||
|
||||
class PTTReportGenerator:
|
||||
"""Generate professional penetration testing reports from PTT data"""
|
||||
|
||||
def __init__(self, report_data: Dict[str, Any]):
|
||||
self.workflow_name = report_data['workflow_name']
|
||||
self.workflow_key = report_data['workflow_key']
|
||||
self.target = report_data['target']
|
||||
self.timestamp = report_data['timestamp']
|
||||
self.conversation_history = report_data['conversation_history']
|
||||
self.tools_used = report_data.get('tools_used', [])
|
||||
self.ptt_data = report_data.get('ptt_data', {})
|
||||
|
||||
# Will be populated by AI analysis
|
||||
self.structured_findings = {}
|
||||
|
||||
def generate_basic_report(self, save_raw_history: bool = False) -> str:
|
||||
"""Generate a basic report without AI analysis"""
|
||||
# Extract findings from PTT nodes
|
||||
vulnerabilities = []
|
||||
completed_tasks = []
|
||||
failed_tasks = []
|
||||
|
||||
for node_data in self.ptt_data.get('nodes', {}).values():
|
||||
if node_data.get('status') == 'completed' and node_data.get('findings'):
|
||||
completed_tasks.append({
|
||||
'description': node_data.get('description', ''),
|
||||
'findings': node_data.get('findings', ''),
|
||||
'tool_used': node_data.get('tool_used', ''),
|
||||
'output_summary': node_data.get('output_summary', '')
|
||||
})
|
||||
elif node_data.get('status') == 'vulnerable':
|
||||
vulnerabilities.append({
|
||||
'title': node_data.get('description', 'Unknown Vulnerability'),
|
||||
'description': node_data.get('findings', 'No description available'),
|
||||
'severity': 'Medium', # Default severity
|
||||
'affected_systems': [self.target],
|
||||
'evidence': node_data.get('output_summary', ''),
|
||||
'remediation': 'Review and patch identified vulnerabilities'
|
||||
})
|
||||
elif node_data.get('status') == 'failed':
|
||||
failed_tasks.append({
|
||||
'description': node_data.get('description', ''),
|
||||
'tool_used': node_data.get('tool_used', ''),
|
||||
'error': node_data.get('output_summary', '')
|
||||
})
|
||||
|
||||
# Create structured findings
|
||||
self.structured_findings = {
|
||||
'executive_summary': f"Autonomous penetration testing completed against {self.target}. Goal: {self.ptt_data.get('goal', 'Unknown')}. {len(completed_tasks)} tasks completed successfully, {len(vulnerabilities)} vulnerabilities identified.",
|
||||
'vulnerabilities': vulnerabilities,
|
||||
'key_statistics': {
|
||||
'total_vulnerabilities': len(vulnerabilities),
|
||||
'critical_count': 0,
|
||||
'high_count': 0,
|
||||
'medium_count': len(vulnerabilities),
|
||||
'systems_compromised': 1 if vulnerabilities else 0
|
||||
},
|
||||
'methodology': f"Autonomous penetration testing using Pentesting Task Tree (PTT) methodology with intelligent task prioritization and execution.",
|
||||
'conclusion': f"Assessment {'successfully identified security weaknesses' if vulnerabilities else 'completed without identifying critical vulnerabilities'}. {'Immediate remediation recommended' if vulnerabilities else 'Continue monitoring and regular assessments'}.",
|
||||
'recommendations': [
|
||||
{
|
||||
'category': 'Patch Management',
|
||||
'recommendation': 'Apply security patches to all identified vulnerable services',
|
||||
'priority': 'Immediate',
|
||||
'business_justification': 'Prevents exploitation of known vulnerabilities'
|
||||
},
|
||||
{
|
||||
'category': 'Monitoring',
|
||||
'recommendation': 'Implement monitoring for the services and ports identified during reconnaissance',
|
||||
'priority': 'Short-term',
|
||||
'business_justification': 'Early detection of potential security incidents'
|
||||
}
|
||||
] if vulnerabilities else [
|
||||
{
|
||||
'category': 'Continued Security',
|
||||
'recommendation': 'Maintain current security posture and conduct regular assessments',
|
||||
'priority': 'Medium-term',
|
||||
'business_justification': 'Proactive security maintenance'
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# Generate the markdown report
|
||||
markdown_report = self.generate_markdown_report()
|
||||
|
||||
# Save report
|
||||
return self.save_report(markdown_report, save_raw_history)
|
||||
|
||||
async def generate_report(self, run_agent_func, connected_servers, kb_instance=None, save_raw_history=False) -> str:
|
||||
"""Generate a comprehensive report with AI analysis"""
|
||||
try:
|
||||
# Create analysis prompt specifically for PTT data
|
||||
analysis_prompt = self.create_ptt_analysis_prompt()
|
||||
|
||||
# Get AI analysis
|
||||
ai_response = await self.analyze_with_ai(
|
||||
analysis_prompt,
|
||||
run_agent_func,
|
||||
connected_servers,
|
||||
kb_instance
|
||||
)
|
||||
|
||||
if ai_response:
|
||||
# Parse AI response
|
||||
self.structured_findings = self.parse_ai_response(ai_response)
|
||||
else:
|
||||
print(f"{Fore.YELLOW}AI analysis failed, generating basic report...{Style.RESET_ALL}")
|
||||
return self.generate_basic_report(save_raw_history)
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Error in AI analysis: {e}. Generating basic report...{Style.RESET_ALL}")
|
||||
return self.generate_basic_report(save_raw_history)
|
||||
|
||||
# Generate markdown report
|
||||
markdown_report = self.generate_markdown_report()
|
||||
|
||||
# Save report
|
||||
return self.save_report(markdown_report, save_raw_history)
|
||||
|
||||
def create_ptt_analysis_prompt(self) -> str:
|
||||
"""Create analysis prompt for PTT data"""
|
||||
ptt_structure = self.ptt_data.get('tree_structure', 'No PTT structure available')
|
||||
goal = self.ptt_data.get('goal', 'Unknown goal')
|
||||
target = self.target
|
||||
statistics = self.ptt_data.get('statistics', {})
|
||||
|
||||
# Extract key findings from completed tasks
|
||||
key_findings = []
|
||||
for node_data in self.ptt_data.get('nodes', {}).values():
|
||||
if node_data.get('status') in ['completed', 'vulnerable'] and node_data.get('findings'):
|
||||
key_findings.append(f"- {node_data.get('description', '')}: {node_data.get('findings', '')}")
|
||||
|
||||
findings_text = '\n'.join(key_findings) if key_findings else 'No significant findings recorded'
|
||||
|
||||
prompt = f"""You are analyzing the results of an autonomous penetration test conducted using a Pentesting Task Tree (PTT) methodology.
|
||||
|
||||
ASSESSMENT DETAILS:
|
||||
Goal: {goal}
|
||||
Target: {target}
|
||||
Statistics: {statistics}
|
||||
|
||||
PTT STRUCTURE:
|
||||
{ptt_structure}
|
||||
|
||||
KEY FINDINGS:
|
||||
{findings_text}
|
||||
|
||||
Based on this PTT analysis, provide a comprehensive security assessment in the following JSON format:
|
||||
|
||||
{{
|
||||
"executive_summary": "Professional executive summary of the assessment",
|
||||
"key_statistics": {{
|
||||
"total_vulnerabilities": 0,
|
||||
"critical_count": 0,
|
||||
"high_count": 0,
|
||||
"medium_count": 0,
|
||||
"low_count": 0,
|
||||
"systems_compromised": 0
|
||||
}},
|
||||
"vulnerabilities": [
|
||||
{{
|
||||
"title": "Vulnerability name",
|
||||
"description": "Detailed description",
|
||||
"severity": "Critical/High/Medium/Low",
|
||||
"impact": "Business impact description",
|
||||
"affected_systems": ["system1", "system2"],
|
||||
"evidence": "Technical evidence",
|
||||
"remediation": "Specific remediation steps",
|
||||
"references": ["CVE-XXXX", "reference links"]
|
||||
}}
|
||||
],
|
||||
"compromised_systems": [
|
||||
{{
|
||||
"system": "system identifier",
|
||||
"access_level": "user/admin/root",
|
||||
"method": "exploitation method",
|
||||
"evidence": "proof of compromise"
|
||||
}}
|
||||
],
|
||||
"attack_paths": [
|
||||
{{
|
||||
"path_description": "Attack path name",
|
||||
"impact": "potential impact",
|
||||
"steps": ["step1", "step2", "step3"]
|
||||
}}
|
||||
],
|
||||
"recommendations": [
|
||||
{{
|
||||
"category": "category name",
|
||||
"recommendation": "specific recommendation",
|
||||
"priority": "Immediate/Short-term/Medium-term/Long-term",
|
||||
"business_justification": "why this matters to business"
|
||||
}}
|
||||
],
|
||||
"methodology": "Description of the PTT methodology used",
|
||||
"conclusion": "Professional conclusion of the assessment"
|
||||
}}
|
||||
|
||||
Focus on:
|
||||
1. Extracting real security findings from the PTT execution
|
||||
2. Proper risk classification
|
||||
3. Actionable recommendations
|
||||
4. Business-relevant impact assessment"""
|
||||
|
||||
return prompt
|
||||
|
||||
async def analyze_with_ai(self, prompt: str, run_agent_func, connected_servers, kb_instance) -> str:
|
||||
"""Analyze the assessment with AI"""
|
||||
try:
|
||||
result = await run_agent_func(
|
||||
prompt,
|
||||
connected_servers,
|
||||
history=[],
|
||||
streaming=True,
|
||||
kb_instance=kb_instance
|
||||
)
|
||||
|
||||
if hasattr(result, "final_output"):
|
||||
return result.final_output
|
||||
elif hasattr(result, "output"):
|
||||
return result.output
|
||||
elif isinstance(result, str):
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error in AI analysis: {e}{Style.RESET_ALL}")
|
||||
|
||||
return None
|
||||
|
||||
def parse_ai_response(self, response: str) -> Dict[str, Any]:
|
||||
"""Parse AI response for structured findings"""
|
||||
try:
|
||||
# Try to extract JSON from the response
|
||||
import re
|
||||
|
||||
# Look for JSON in the response
|
||||
json_match = re.search(r'\{.*\}', response, re.DOTALL)
|
||||
if json_match:
|
||||
json_str = json_match.group()
|
||||
return json.loads(json_str)
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Fore.YELLOW}Failed to parse AI response: {e}{Style.RESET_ALL}")
|
||||
|
||||
# Fallback to basic findings
|
||||
return {
|
||||
'executive_summary': 'Assessment completed successfully.',
|
||||
'vulnerabilities': [],
|
||||
'key_statistics': {'total_vulnerabilities': 0},
|
||||
'methodology': 'Autonomous penetration testing using PTT methodology.',
|
||||
'conclusion': 'Assessment completed.',
|
||||
'recommendations': []
|
||||
}
|
||||
|
||||
def generate_markdown_report(self) -> str:
|
||||
"""Generate the final markdown report using the same format as PentestReportGenerator"""
|
||||
# Use the same report generation logic as the workflow reporter
|
||||
temp_generator = PentestReportGenerator({
|
||||
'workflow_name': self.workflow_name,
|
||||
'workflow_key': self.workflow_key,
|
||||
'target': self.target,
|
||||
'timestamp': self.timestamp,
|
||||
'conversation_history': self.conversation_history,
|
||||
'tools_used': self.tools_used
|
||||
})
|
||||
temp_generator.structured_findings = self.structured_findings
|
||||
|
||||
return temp_generator.generate_markdown_report()
|
||||
|
||||
def save_report(self, markdown_content: str, save_raw_history: bool = False) -> str:
|
||||
"""Save the report to file"""
|
||||
# Create reports directory if it doesn't exist
|
||||
reports_dir = "reports"
|
||||
if not os.path.exists(reports_dir):
|
||||
os.makedirs(reports_dir)
|
||||
|
||||
# Generate filename
|
||||
timestamp_str = str(int(self.timestamp.timestamp()))
|
||||
safe_target = re.sub(r'[^\w\-_\.]', '_', self.target)
|
||||
filename = f"{reports_dir}/ghostcrew_agent_mode_{safe_target}_{timestamp_str}.md"
|
||||
|
||||
# Save markdown file
|
||||
with open(filename, 'w', encoding='utf-8') as f:
|
||||
f.write(markdown_content)
|
||||
|
||||
# Optionally save raw history and PTT data
|
||||
if save_raw_history:
|
||||
raw_filename = f"{reports_dir}/ghostcrew_agent_mode_{safe_target}_{timestamp_str}_raw.json"
|
||||
raw_data = {
|
||||
'ptt_data': self.ptt_data,
|
||||
'conversation_history': self.conversation_history,
|
||||
'timestamp': self.timestamp.isoformat()
|
||||
}
|
||||
|
||||
with open(raw_filename, 'w', encoding='utf-8') as f:
|
||||
json.dump(raw_data, f, indent=2, default=str)
|
||||
|
||||
print(f"{Fore.GREEN}Raw PTT data saved: {raw_filename}{Style.RESET_ALL}")
|
||||
|
||||
return filename
|
||||
@@ -1,14 +1,46 @@
|
||||
colorama==0.4.6
|
||||
python-dotenv==1.1.0
|
||||
openai==1.78.1
|
||||
uvicorn==0.34.0
|
||||
mcp==1.6.0
|
||||
langchain==0.3.25
|
||||
langchain-community==0.3.24
|
||||
numpy==2.2.5
|
||||
ollama==0.4.8
|
||||
openai-agents==0.0.14
|
||||
fastapi==0.115.9
|
||||
pymetasploit3==1.0.6
|
||||
tiktoken==0.9.0
|
||||
# Add other necessary dependencies
|
||||
# GhostCrew Dependencies
|
||||
|
||||
# Core LLM
|
||||
litellm>=1.40.0
|
||||
openai>=1.30.0
|
||||
anthropic>=0.25.0
|
||||
tiktoken>=0.7.0
|
||||
|
||||
# Async
|
||||
aiohttp>=3.9.0
|
||||
aiofiles>=23.2.0
|
||||
|
||||
# Web
|
||||
playwright>=1.44.0
|
||||
beautifulsoup4>=4.12.0
|
||||
httpx>=0.27.0
|
||||
|
||||
# RAG / Embeddings
|
||||
numpy>=1.26.0
|
||||
sentence-transformers>=2.7.0
|
||||
faiss-cpu>=1.8.0
|
||||
|
||||
# Docker
|
||||
docker>=7.0.0
|
||||
|
||||
# CLI / TUI
|
||||
rich>=13.7.0
|
||||
textual>=0.63.0
|
||||
typer>=0.12.0
|
||||
|
||||
# Config
|
||||
pydantic>=2.7.0
|
||||
pydantic-settings>=2.2.0
|
||||
python-dotenv>=1.0.0
|
||||
pyyaml>=6.0.0
|
||||
jinja2>=3.1.0
|
||||
|
||||
# Dev
|
||||
pytest>=8.2.0
|
||||
pytest-asyncio>=0.23.0
|
||||
pytest-cov>=5.0.0
|
||||
pytest-mock>=3.14.0
|
||||
black>=24.4.0
|
||||
isort>=5.13.0
|
||||
mypy>=1.10.0
|
||||
ruff>=0.4.0
|
||||
|
||||
74
scripts/run.sh
Normal file
74
scripts/run.sh
Normal file
@@ -0,0 +1,74 @@
|
||||
#!/bin/bash
|
||||
# GhostCrew Run Script
|
||||
|
||||
set -e
|
||||
|
||||
# Activate virtual environment if exists
|
||||
if [ -d "venv" ]; then
|
||||
source venv/bin/activate
|
||||
fi
|
||||
|
||||
# Load environment variables
|
||||
if [ -f ".env" ]; then
|
||||
export $(grep -v '^#' .env | xargs)
|
||||
fi
|
||||
|
||||
# Parse arguments
|
||||
MODE="cli"
|
||||
TARGET=""
|
||||
VPN_CONFIG=""
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--tui)
|
||||
MODE="tui"
|
||||
shift
|
||||
;;
|
||||
--target)
|
||||
TARGET="$2"
|
||||
shift 2
|
||||
;;
|
||||
--vpn)
|
||||
VPN_CONFIG="$2"
|
||||
shift 2
|
||||
;;
|
||||
--help)
|
||||
echo "GhostCrew - AI Penetration Testing"
|
||||
echo ""
|
||||
echo "Usage: run.sh [options]"
|
||||
echo ""
|
||||
echo "Options:"
|
||||
echo " --tui Run in TUI mode"
|
||||
echo " --target <url> Set initial target"
|
||||
echo " --vpn <config> Connect to VPN before starting"
|
||||
echo " --help Show this help message"
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Connect to VPN if specified
|
||||
if [ -n "$VPN_CONFIG" ]; then
|
||||
echo "Connecting to VPN..."
|
||||
sudo openvpn --config "$VPN_CONFIG" --daemon
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Build command
|
||||
CMD="python -m ghostcrew"
|
||||
|
||||
if [ "$MODE" = "tui" ]; then
|
||||
CMD="$CMD --tui"
|
||||
fi
|
||||
|
||||
if [ -n "$TARGET" ]; then
|
||||
CMD="$CMD --target $TARGET"
|
||||
fi
|
||||
|
||||
# Run GhostCrew
|
||||
echo "Starting GhostCrew..."
|
||||
$CMD
|
||||
89
scripts/setup.ps1
Normal file
89
scripts/setup.ps1
Normal file
@@ -0,0 +1,89 @@
|
||||
# GhostCrew PowerShell Setup Script
|
||||
|
||||
Write-Host "GhostCrew Setup" -ForegroundColor Blue
|
||||
Write-Host "AI Penetration Testing" -ForegroundColor Green
|
||||
Write-Host ""
|
||||
|
||||
# Check Python version
|
||||
Write-Host "Checking Python version..." -ForegroundColor Yellow
|
||||
try {
|
||||
$pythonVersion = python --version 2>&1
|
||||
if ($pythonVersion -match "Python (\d+)\.(\d+)") {
|
||||
$major = [int]$Matches[1]
|
||||
$minor = [int]$Matches[2]
|
||||
if ($major -lt 3 -or ($major -eq 3 -and $minor -lt 10)) {
|
||||
Write-Host "Error: Python 3.10 or higher is required" -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
Write-Host "[OK] $pythonVersion" -ForegroundColor Green
|
||||
}
|
||||
} catch {
|
||||
Write-Host "Error: Python not found. Please install Python 3.10+" -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Create virtual environment
|
||||
Write-Host "Creating virtual environment..." -ForegroundColor Yellow
|
||||
if (-not (Test-Path "venv")) {
|
||||
python -m venv venv
|
||||
Write-Host "[OK] Virtual environment created" -ForegroundColor Green
|
||||
} else {
|
||||
Write-Host "[OK] Virtual environment exists" -ForegroundColor Green
|
||||
}
|
||||
|
||||
# Activate virtual environment
|
||||
Write-Host "Activating virtual environment..." -ForegroundColor Yellow
|
||||
& .\venv\Scripts\Activate.ps1
|
||||
|
||||
# Upgrade pip
|
||||
Write-Host "Upgrading pip..." -ForegroundColor Yellow
|
||||
pip install --upgrade pip
|
||||
|
||||
# Install dependencies
|
||||
Write-Host "Installing dependencies..." -ForegroundColor Yellow
|
||||
pip install -e ".[all]"
|
||||
Write-Host "[OK] Dependencies installed" -ForegroundColor Green
|
||||
|
||||
# Install playwright browsers
|
||||
Write-Host "Installing Playwright browsers..." -ForegroundColor Yellow
|
||||
playwright install chromium
|
||||
Write-Host "[OK] Playwright browsers installed" -ForegroundColor Green
|
||||
|
||||
# Create .env file if not exists
|
||||
if (-not (Test-Path ".env")) {
|
||||
Write-Host "Creating .env file..." -ForegroundColor Yellow
|
||||
@"
|
||||
# GhostCrew Configuration
|
||||
# Add your API keys here
|
||||
|
||||
# OpenAI API Key (required for GPT models)
|
||||
OPENAI_API_KEY=
|
||||
|
||||
# Anthropic API Key (required for Claude models)
|
||||
ANTHROPIC_API_KEY=
|
||||
|
||||
# Model Configuration
|
||||
GHOSTCREW_MODEL=gpt-5
|
||||
|
||||
# Debug Mode
|
||||
GHOSTCREW_DEBUG=false
|
||||
|
||||
# Max Iterations
|
||||
GHOSTCREW_MAX_ITERATIONS=50
|
||||
"@ | Set-Content -Path ".env" -Encoding UTF8
|
||||
Write-Host "[OK] .env file created" -ForegroundColor Green
|
||||
Write-Host "[!] Please edit .env and add your API keys" -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
# Create loot directory for reports
|
||||
New-Item -ItemType Directory -Force -Path "loot" | Out-Null
|
||||
Write-Host "[OK] Loot directory created" -ForegroundColor Green
|
||||
|
||||
Write-Host ""
|
||||
Write-Host "Setup complete!" -ForegroundColor Green
|
||||
Write-Host ""
|
||||
Write-Host "To get started:"
|
||||
Write-Host " 1. Edit .env and add your API keys"
|
||||
Write-Host " 2. Activate: .\venv\Scripts\Activate.ps1"
|
||||
Write-Host " 3. Run: ghostcrew or python -m ghostcrew"
|
||||
Write-Host ""
|
||||
97
scripts/setup.sh
Normal file
97
scripts/setup.sh
Normal file
@@ -0,0 +1,97 @@
|
||||
#!/bin/bash
|
||||
# GhostCrew Setup Script
|
||||
|
||||
set -e
|
||||
|
||||
# Colors
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
BLUE='\033[0;34m'
|
||||
NC='\033[0m'
|
||||
|
||||
echo -e "${BLUE}GhostCrew${NC} - AI Penetration Testing"
|
||||
echo ""
|
||||
|
||||
# Check Python version
|
||||
echo -e "${YELLOW}Checking Python version...${NC}"
|
||||
python_version=$(python3 --version 2>&1 | cut -d' ' -f2 | cut -d'.' -f1,2)
|
||||
required_version="3.10"
|
||||
|
||||
if [ "$(printf '%s\n' "$required_version" "$python_version" | sort -V | head -n1)" != "$required_version" ]; then
|
||||
echo -e "${RED}Error: Python $required_version or higher is required (found $python_version)${NC}"
|
||||
exit 1
|
||||
fi
|
||||
echo -e "${GREEN}✓ Python $python_version${NC}"
|
||||
|
||||
# Create virtual environment
|
||||
echo -e "${YELLOW}Creating virtual environment...${NC}"
|
||||
if [ ! -d "venv" ]; then
|
||||
python3 -m venv venv
|
||||
echo -e "${GREEN}✓ Virtual environment created${NC}"
|
||||
else
|
||||
echo -e "${GREEN}✓ Virtual environment exists${NC}"
|
||||
fi
|
||||
|
||||
# Activate virtual environment
|
||||
echo -e "${YELLOW}Activating virtual environment...${NC}"
|
||||
source venv/bin/activate
|
||||
|
||||
# Upgrade pip
|
||||
echo -e "${YELLOW}Upgrading pip...${NC}"
|
||||
pip install --upgrade pip
|
||||
|
||||
# Install dependencies
|
||||
echo -e "${YELLOW}Installing dependencies...${NC}"
|
||||
pip install -e ".[all]"
|
||||
echo -e "${GREEN}✓ Dependencies installed${NC}"
|
||||
|
||||
# Install playwright browsers
|
||||
echo -e "${YELLOW}Installing Playwright browsers...${NC}"
|
||||
playwright install chromium
|
||||
echo -e "${GREEN}✓ Playwright browsers installed${NC}"
|
||||
|
||||
# Create .env file if not exists
|
||||
if [ ! -f ".env" ]; then
|
||||
echo -e "${YELLOW}Creating .env file...${NC}"
|
||||
cat > .env << EOF
|
||||
# GhostCrew Configuration
|
||||
# Add your API keys here
|
||||
|
||||
# OpenAI API Key (required for GPT models)
|
||||
OPENAI_API_KEY=
|
||||
|
||||
# Anthropic API Key (required for Claude models)
|
||||
ANTHROPIC_API_KEY=
|
||||
|
||||
# Model Configuration
|
||||
GHOSTCREW_MODEL=gpt-5
|
||||
|
||||
# Debug Mode
|
||||
GHOSTCREW_DEBUG=false
|
||||
|
||||
# Max Iterations
|
||||
GHOSTCREW_MAX_ITERATIONS=50
|
||||
EOF
|
||||
echo -e "${GREEN}✓ .env file created${NC}"
|
||||
echo -e "${YELLOW}⚠️ Please edit .env and add your API keys${NC}"
|
||||
fi
|
||||
|
||||
# Create loot directory for reports
|
||||
mkdir -p loot
|
||||
echo -e "${GREEN}✓ Loot directory created${NC}"
|
||||
|
||||
echo ""
|
||||
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo -e "${GREEN}Setup complete!${NC}"
|
||||
echo ""
|
||||
echo -e "To get started:"
|
||||
echo -e " 1. Edit ${YELLOW}.env${NC} and add your API keys"
|
||||
echo -e " 2. Activate the virtual environment: ${YELLOW}source venv/bin/activate${NC}"
|
||||
echo -e " 3. Run GhostCrew: ${YELLOW}ghostcrew${NC} or ${YELLOW}python -m ghostcrew${NC}"
|
||||
echo ""
|
||||
echo -e "For Docker usage:"
|
||||
echo -e " ${YELLOW}docker-compose up ghostcrew${NC}"
|
||||
echo -e " ${YELLOW}docker-compose --profile kali up ghostcrew-kali${NC}"
|
||||
echo ""
|
||||
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
1
tests/__init__.py
Normal file
1
tests/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# GhostCrew Tests
|
||||
95
tests/conftest.py
Normal file
95
tests/conftest.py
Normal file
@@ -0,0 +1,95 @@
|
||||
"""Test fixtures for GhostCrew tests."""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Generator, AsyncGenerator
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
|
||||
from ghostcrew.config import Settings
|
||||
from ghostcrew.agents.state import AgentState, AgentStateManager
|
||||
from ghostcrew.tools import get_all_tools, Tool, ToolSchema
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def settings() -> Settings:
|
||||
"""Create test settings."""
|
||||
return Settings(
|
||||
debug=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_state() -> AgentStateManager:
|
||||
"""Create a test agent state manager."""
|
||||
return AgentStateManager()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm() -> MagicMock:
|
||||
"""Create a mock LLM."""
|
||||
mock = AsyncMock()
|
||||
mock.generate.return_value = MagicMock(
|
||||
content="Test response",
|
||||
tool_calls=None,
|
||||
usage={"prompt_tokens": 100, "completion_tokens": 50},
|
||||
model="gpt-5",
|
||||
finish_reason="stop"
|
||||
)
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir(tmp_path: Path) -> Path:
|
||||
"""Create a temporary directory for tests."""
|
||||
return tmp_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def event_loop():
|
||||
"""Create an event loop for async tests."""
|
||||
loop = asyncio.new_event_loop()
|
||||
yield loop
|
||||
loop.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_finding() -> dict:
|
||||
"""Create a sample security finding."""
|
||||
return {
|
||||
"title": "SQL Injection in Login Form",
|
||||
"severity": "high",
|
||||
"target": "http://example.com/login",
|
||||
"description": "The login form is vulnerable to SQL injection attacks.",
|
||||
"evidence": "Parameter 'username' with payload: ' OR '1'='1",
|
||||
"remediation": "Use parameterized queries or prepared statements."
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_tool_result() -> dict:
|
||||
"""Create a sample tool execution result."""
|
||||
return {
|
||||
"tool": "terminal",
|
||||
"success": True,
|
||||
"output": "nmap scan results...",
|
||||
"duration_ms": 1500.0
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_tool() -> Tool:
|
||||
"""Create a sample tool for testing."""
|
||||
async def dummy_execute(arguments: dict, runtime) -> str:
|
||||
return f"Executed with: {arguments}"
|
||||
|
||||
return Tool(
|
||||
name="test_tool",
|
||||
description="A test tool",
|
||||
schema=ToolSchema(
|
||||
properties={"param": {"type": "string", "description": "A parameter"}},
|
||||
required=["param"]
|
||||
),
|
||||
execute_fn=dummy_execute,
|
||||
category="test"
|
||||
)
|
||||
110
tests/test_agents.py
Normal file
110
tests/test_agents.py
Normal file
@@ -0,0 +1,110 @@
|
||||
"""Tests for the agent state management."""
|
||||
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
|
||||
from ghostcrew.agents.state import AgentState, AgentStateManager, StateTransition
|
||||
|
||||
|
||||
class TestAgentState:
|
||||
"""Tests for AgentState enum."""
|
||||
|
||||
def test_state_values(self):
|
||||
"""Test state enum values."""
|
||||
assert AgentState.IDLE.value == "idle"
|
||||
assert AgentState.THINKING.value == "thinking"
|
||||
assert AgentState.EXECUTING.value == "executing"
|
||||
assert AgentState.WAITING_INPUT.value == "waiting_input"
|
||||
assert AgentState.COMPLETE.value == "complete"
|
||||
assert AgentState.ERROR.value == "error"
|
||||
|
||||
def test_all_states_exist(self):
|
||||
"""Test that all expected states exist."""
|
||||
states = list(AgentState)
|
||||
assert len(states) >= 6
|
||||
|
||||
|
||||
class TestAgentStateManager:
|
||||
"""Tests for AgentStateManager class."""
|
||||
|
||||
@pytest.fixture
|
||||
def state_manager(self):
|
||||
"""Create a fresh AgentStateManager for each test."""
|
||||
return AgentStateManager()
|
||||
|
||||
def test_initial_state(self, state_manager):
|
||||
"""Test initial state is IDLE."""
|
||||
assert state_manager.current_state == AgentState.IDLE
|
||||
assert len(state_manager.history) == 0
|
||||
|
||||
def test_valid_transition(self, state_manager):
|
||||
"""Test valid state transition."""
|
||||
result = state_manager.transition_to(AgentState.THINKING)
|
||||
assert result is True
|
||||
assert state_manager.current_state == AgentState.THINKING
|
||||
assert len(state_manager.history) == 1
|
||||
|
||||
def test_invalid_transition(self, state_manager):
|
||||
"""Test invalid state transition."""
|
||||
result = state_manager.transition_to(AgentState.COMPLETE)
|
||||
assert result is False
|
||||
assert state_manager.current_state == AgentState.IDLE
|
||||
|
||||
def test_transition_chain(self, state_manager):
|
||||
"""Test a chain of valid transitions."""
|
||||
assert state_manager.transition_to(AgentState.THINKING)
|
||||
assert state_manager.transition_to(AgentState.EXECUTING)
|
||||
assert state_manager.transition_to(AgentState.THINKING)
|
||||
assert state_manager.transition_to(AgentState.COMPLETE)
|
||||
|
||||
assert state_manager.current_state == AgentState.COMPLETE
|
||||
assert len(state_manager.history) == 4
|
||||
|
||||
def test_force_transition(self, state_manager):
|
||||
"""Test forcing a transition."""
|
||||
state_manager.force_transition(AgentState.ERROR, reason="Test error")
|
||||
assert state_manager.current_state == AgentState.ERROR
|
||||
assert "FORCED" in state_manager.history[-1].reason
|
||||
|
||||
def test_reset(self, state_manager):
|
||||
"""Test resetting state."""
|
||||
state_manager.transition_to(AgentState.THINKING)
|
||||
state_manager.transition_to(AgentState.EXECUTING)
|
||||
|
||||
state_manager.reset()
|
||||
|
||||
assert state_manager.current_state == AgentState.IDLE
|
||||
assert len(state_manager.history) == 0
|
||||
|
||||
def test_is_terminal(self, state_manager):
|
||||
"""Test terminal state detection."""
|
||||
assert state_manager.is_terminal() is False
|
||||
|
||||
state_manager.transition_to(AgentState.THINKING)
|
||||
state_manager.transition_to(AgentState.COMPLETE)
|
||||
|
||||
assert state_manager.is_terminal() is True
|
||||
|
||||
def test_is_active(self, state_manager):
|
||||
"""Test active state detection."""
|
||||
assert state_manager.is_active() is False
|
||||
|
||||
state_manager.transition_to(AgentState.THINKING)
|
||||
assert state_manager.is_active() is True
|
||||
|
||||
|
||||
class TestStateTransition:
|
||||
"""Tests for StateTransition dataclass."""
|
||||
|
||||
def test_create_transition(self):
|
||||
"""Test creating a state transition."""
|
||||
transition = StateTransition(
|
||||
from_state=AgentState.IDLE,
|
||||
to_state=AgentState.THINKING,
|
||||
reason="Starting work"
|
||||
)
|
||||
|
||||
assert transition.from_state == AgentState.IDLE
|
||||
assert transition.to_state == AgentState.THINKING
|
||||
assert transition.reason == "Starting work"
|
||||
assert transition.timestamp is not None
|
||||
100
tests/test_knowledge.py
Normal file
100
tests/test_knowledge.py
Normal file
@@ -0,0 +1,100 @@
|
||||
"""Tests for the RAG knowledge system."""
|
||||
|
||||
import pytest
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from ghostcrew.knowledge.rag import RAGEngine, Document
|
||||
|
||||
|
||||
class TestDocument:
|
||||
"""Tests for Document dataclass."""
|
||||
|
||||
def test_create_document(self):
|
||||
"""Test creating a document."""
|
||||
doc = Document(content="Test content", source="test.md")
|
||||
assert doc.content == "Test content"
|
||||
assert doc.source == "test.md"
|
||||
assert doc.metadata == {}
|
||||
assert doc.doc_id is not None
|
||||
|
||||
def test_document_with_metadata(self):
|
||||
"""Test document with metadata."""
|
||||
doc = Document(
|
||||
content="Test",
|
||||
source="test.md",
|
||||
metadata={"cve_id": "CVE-2021-1234", "severity": "high"}
|
||||
)
|
||||
assert doc.metadata["cve_id"] == "CVE-2021-1234"
|
||||
assert doc.metadata["severity"] == "high"
|
||||
|
||||
def test_document_with_embedding(self):
|
||||
"""Test document with embedding."""
|
||||
embedding = np.random.rand(384)
|
||||
doc = Document(content="Test", source="test.md", embedding=embedding)
|
||||
assert doc.embedding is not None
|
||||
assert len(doc.embedding) == 384
|
||||
|
||||
def test_document_with_custom_id(self):
|
||||
"""Test document with custom doc_id."""
|
||||
doc = Document(content="Test", source="test.md", doc_id="custom-id-123")
|
||||
assert doc.doc_id == "custom-id-123"
|
||||
|
||||
|
||||
class TestRAGEngine:
|
||||
"""Tests for RAGEngine class."""
|
||||
|
||||
@pytest.fixture
|
||||
def rag_engine(self, tmp_path):
|
||||
"""Create a RAG engine for testing."""
|
||||
return RAGEngine(
|
||||
knowledge_path=tmp_path / "knowledge",
|
||||
use_local_embeddings=True
|
||||
)
|
||||
|
||||
def test_create_engine(self, rag_engine):
|
||||
"""Test creating a RAG engine."""
|
||||
assert rag_engine is not None
|
||||
assert len(rag_engine.documents) == 0
|
||||
assert rag_engine.embeddings is None
|
||||
|
||||
def test_get_document_count_empty(self, rag_engine):
|
||||
"""Test document count on empty engine."""
|
||||
assert rag_engine.get_document_count() == 0
|
||||
|
||||
def test_clear(self, rag_engine):
|
||||
"""Test clearing the engine."""
|
||||
rag_engine.documents.append(Document(content="test", source="test.md"))
|
||||
rag_engine.embeddings = np.random.rand(1, 384)
|
||||
rag_engine._indexed = True
|
||||
|
||||
rag_engine.clear()
|
||||
|
||||
assert len(rag_engine.documents) == 0
|
||||
assert rag_engine.embeddings is None
|
||||
assert rag_engine._indexed == False
|
||||
|
||||
|
||||
class TestRAGEngineChunking:
|
||||
"""Tests for text chunking functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def engine(self, tmp_path):
|
||||
"""Create engine for chunking tests."""
|
||||
return RAGEngine(knowledge_path=tmp_path)
|
||||
|
||||
def test_chunk_short_text(self, engine):
|
||||
"""Test chunking text shorter than chunk size."""
|
||||
text = "This is a short paragraph.\n\nThis is another paragraph."
|
||||
chunks = engine._chunk_text(text, source="test.md", chunk_size=1000)
|
||||
|
||||
assert len(chunks) >= 1
|
||||
assert all(isinstance(c, Document) for c in chunks)
|
||||
|
||||
def test_chunk_preserves_source(self, engine):
|
||||
"""Test that chunking preserves source information."""
|
||||
text = "Test paragraph 1.\n\nTest paragraph 2."
|
||||
chunks = engine._chunk_text(text, source="my_source.md")
|
||||
|
||||
assert all(c.source == "my_source.md" for c in chunks)
|
||||
141
tests/test_tools.py
Normal file
141
tests/test_tools.py
Normal file
@@ -0,0 +1,141 @@
|
||||
"""Tests for the tool system."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ghostcrew.tools import (
|
||||
Tool, ToolSchema, register_tool, get_all_tools, get_tool,
|
||||
enable_tool, disable_tool, get_tool_names
|
||||
)
|
||||
|
||||
|
||||
class TestToolRegistry:
|
||||
"""Tests for tool registry functions."""
|
||||
|
||||
def test_tools_loaded(self):
|
||||
"""Test that built-in tools are loaded."""
|
||||
tools = get_all_tools()
|
||||
assert len(tools) > 0
|
||||
|
||||
tool_names = get_tool_names()
|
||||
assert "terminal" in tool_names
|
||||
assert "browser" in tool_names
|
||||
|
||||
def test_get_tool(self):
|
||||
"""Test getting a tool by name."""
|
||||
tool = get_tool("terminal")
|
||||
assert tool is not None
|
||||
assert tool.name == "terminal"
|
||||
assert tool.category == "execution"
|
||||
|
||||
def test_get_nonexistent_tool(self):
|
||||
"""Test getting a tool that doesn't exist."""
|
||||
tool = get_tool("nonexistent_tool_xyz")
|
||||
assert tool is None
|
||||
|
||||
def test_disable_enable_tool(self):
|
||||
"""Test disabling and enabling a tool."""
|
||||
result = disable_tool("terminal")
|
||||
assert result is True
|
||||
|
||||
tool = get_tool("terminal")
|
||||
assert tool.enabled is False
|
||||
|
||||
result = enable_tool("terminal")
|
||||
assert result is True
|
||||
|
||||
tool = get_tool("terminal")
|
||||
assert tool.enabled is True
|
||||
|
||||
def test_disable_nonexistent_tool(self):
|
||||
"""Test disabling a tool that doesn't exist."""
|
||||
result = disable_tool("nonexistent_tool_xyz")
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestToolSchema:
|
||||
"""Tests for ToolSchema class."""
|
||||
|
||||
def test_create_schema(self):
|
||||
"""Test creating a tool schema."""
|
||||
schema = ToolSchema(
|
||||
properties={
|
||||
"command": {"type": "string", "description": "Command to run"}
|
||||
},
|
||||
required=["command"]
|
||||
)
|
||||
|
||||
assert schema.type == "object"
|
||||
assert "command" in schema.properties
|
||||
assert "command" in schema.required
|
||||
|
||||
def test_schema_to_dict(self):
|
||||
"""Test converting schema to dictionary."""
|
||||
schema = ToolSchema(
|
||||
properties={"input": {"type": "string"}},
|
||||
required=["input"]
|
||||
)
|
||||
|
||||
d = schema.to_dict()
|
||||
assert d["type"] == "object"
|
||||
assert d["properties"]["input"]["type"] == "string"
|
||||
assert d["required"] == ["input"]
|
||||
|
||||
|
||||
class TestTool:
|
||||
"""Tests for Tool class."""
|
||||
|
||||
def test_create_tool(self, sample_tool):
|
||||
"""Test creating a tool."""
|
||||
assert sample_tool.name == "test_tool"
|
||||
assert sample_tool.description == "A test tool"
|
||||
assert sample_tool.category == "test"
|
||||
assert sample_tool.enabled is True
|
||||
|
||||
def test_tool_to_llm_format(self, sample_tool):
|
||||
"""Test converting tool to LLM format."""
|
||||
llm_format = sample_tool.to_llm_format()
|
||||
|
||||
assert llm_format["type"] == "function"
|
||||
assert llm_format["function"]["name"] == "test_tool"
|
||||
assert llm_format["function"]["description"] == "A test tool"
|
||||
assert "parameters" in llm_format["function"]
|
||||
|
||||
def test_tool_validate_arguments(self, sample_tool):
|
||||
"""Test argument validation."""
|
||||
is_valid, error = sample_tool.validate_arguments({"param": "value"})
|
||||
assert is_valid is True
|
||||
assert error is None
|
||||
|
||||
is_valid, error = sample_tool.validate_arguments({})
|
||||
assert is_valid is False
|
||||
assert "param" in error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_execute(self, sample_tool):
|
||||
"""Test tool execution."""
|
||||
result = await sample_tool.execute({"param": "test"}, runtime=None)
|
||||
assert "test" in result
|
||||
|
||||
|
||||
class TestRegisterToolDecorator:
|
||||
"""Tests for register_tool decorator."""
|
||||
|
||||
def test_decorator_registers_tool(self):
|
||||
"""Test that decorator registers a new tool."""
|
||||
initial_count = len(get_all_tools())
|
||||
|
||||
@register_tool(
|
||||
name="pytest_test_tool_unique",
|
||||
description="A tool registered in tests",
|
||||
schema=ToolSchema(properties={}, required=[]),
|
||||
category="test"
|
||||
)
|
||||
async def pytest_test_tool_unique(arguments, runtime):
|
||||
return "test result"
|
||||
|
||||
new_count = len(get_all_tools())
|
||||
assert new_count == initial_count + 1
|
||||
|
||||
tool = get_tool("pytest_test_tool_unique")
|
||||
assert tool is not None
|
||||
assert tool.name == "pytest_test_tool_unique"
|
||||
@@ -1 +0,0 @@
|
||||
"""MCP (Model Context Protocol) integration for GHOSTCREW."""
|
||||
@@ -1,445 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from colorama import init, Fore, Style
|
||||
|
||||
init(autoreset=True)
|
||||
|
||||
def find_tool_path(tool_name):
|
||||
"""Auto-discover tool path using system commands"""
|
||||
try:
|
||||
if platform.system() == "Windows":
|
||||
# Use 'where' command on Windows
|
||||
result = subprocess.run(['where', tool_name],
|
||||
capture_output=True, text=True, check=False)
|
||||
if result.returncode == 0:
|
||||
# Get first valid path from results
|
||||
paths = result.stdout.strip().split('\n')
|
||||
for path in paths:
|
||||
path = path.strip()
|
||||
if path and os.path.exists(path):
|
||||
return path
|
||||
else:
|
||||
# Use 'which' command on Linux/Mac
|
||||
path = shutil.which(tool_name)
|
||||
if path and os.path.exists(path):
|
||||
return path
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def get_tool_search_variants(exe_name):
|
||||
"""Get different variants of tool names to search for"""
|
||||
if not exe_name:
|
||||
return []
|
||||
|
||||
# For Windows, just search for the base name - 'where' will find the actual executable
|
||||
base_name = exe_name.replace('.exe', '').replace('.py', '')
|
||||
variants = [base_name]
|
||||
|
||||
# Also try the exact name if it's different
|
||||
if exe_name != base_name:
|
||||
variants.append(exe_name)
|
||||
|
||||
return variants
|
||||
|
||||
def auto_discover_tool_path(server):
|
||||
"""Auto-discover tool path with user confirmation"""
|
||||
if not server.get('exe_name'):
|
||||
return None
|
||||
|
||||
print(f"{Fore.CYAN}Searching for {server['name']}...{Style.RESET_ALL}")
|
||||
|
||||
# Get search variants
|
||||
search_variants = get_tool_search_variants(server['exe_name'])
|
||||
|
||||
# Try to find the tool
|
||||
found_path = None
|
||||
for variant in search_variants:
|
||||
found_path = find_tool_path(variant)
|
||||
if found_path:
|
||||
break
|
||||
|
||||
if found_path:
|
||||
print(f"{Fore.GREEN}Found: {found_path}{Style.RESET_ALL}")
|
||||
choice = input(f" Use this path? (yes/no): ").strip().lower()
|
||||
if choice == 'y' or choice == 'yes':
|
||||
return found_path
|
||||
# If user says no, fall through to manual input
|
||||
else:
|
||||
print(f"{Fore.YELLOW}{server['name']} not found in PATH{Style.RESET_ALL}")
|
||||
|
||||
# Fallback to manual input
|
||||
manual_path = input(f" Enter path to {server['exe_name']} manually (or press Enter to skip): ").strip()
|
||||
return manual_path if manual_path else None
|
||||
|
||||
MCP_SERVERS = [
|
||||
{
|
||||
"name": "AlterX",
|
||||
"key": "AlterX",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-alterx-mcp"],
|
||||
"description": "MCP server for subdomain permutation and wordlist generation using the AlterX tool.",
|
||||
"exe_name": "alterx.exe",
|
||||
"env_var": "ALTERX_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-alterx-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Amass",
|
||||
"key": "Amass",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-amass-mcp"],
|
||||
"description": "MCP server for advanced subdomain enumeration and reconnaissance using the Amass tool.",
|
||||
"exe_name": "amass.exe",
|
||||
"env_var": "AMASS_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-amass-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Arjun",
|
||||
"key": "Arjun",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-arjun-mcp"],
|
||||
"description": "MCP server for discovering hidden HTTP parameters using the Arjun tool.",
|
||||
"exe_name": "arjun",
|
||||
"env_var": "ARJUN_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-arjun-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Assetfinder",
|
||||
"key": "Assetfinder",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-assetfinder-mcp"],
|
||||
"description": "MCP server for passive subdomain discovery using the Assetfinder tool.",
|
||||
"exe_name": "assetfinder.exe",
|
||||
"env_var": "ASSETFINDER_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-assetfinder-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Certificate Transparency",
|
||||
"key": "CrtSh",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-crtsh-mcp"],
|
||||
"description": "MCP server for subdomain discovery using SSL certificate transparency logs (crt.sh).",
|
||||
"exe_name": None, # No executable needed for this service
|
||||
"env_var": None,
|
||||
"homepage": "https://www.npmjs.com/package/gc-crtsh-mcp"
|
||||
},
|
||||
{
|
||||
"name": "FFUF Fuzzer",
|
||||
"key": "FFUF",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-ffuf-mcp"],
|
||||
"description": "MCP server for web fuzzing operations using FFUF (Fuzz Faster U Fool) tool.",
|
||||
"exe_name": "ffuf.exe",
|
||||
"env_var": "FFUF_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-ffuf-mcp"
|
||||
},
|
||||
{
|
||||
"name": "httpx",
|
||||
"key": "HTTPx",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-httpx-mcp"],
|
||||
"description": "MCP server for fast HTTP toolkit and port scanning using the httpx tool.",
|
||||
"exe_name": "httpx.exe",
|
||||
"env_var": "HTTPX_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-httpx-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Hydra",
|
||||
"key": "Hydra",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-hydra-mcp"],
|
||||
"description": "MCP server for password brute-force attacks and credential testing using the Hydra tool.",
|
||||
"exe_name": "hydra.exe",
|
||||
"env_var": "HYDRA_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-hydra-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Katana",
|
||||
"key": "Katana",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-katana-mcp"],
|
||||
"description": "MCP server for fast web crawling with JavaScript parsing using the Katana tool.",
|
||||
"exe_name": "katana.exe",
|
||||
"env_var": "KATANA_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-katana-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Masscan",
|
||||
"key": "Masscan",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-masscan-mcp"],
|
||||
"description": "MCP server for high-speed network port scanning with the Masscan tool.",
|
||||
"exe_name": "masscan.exe",
|
||||
"env_var": "MASSCAN_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-masscan-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Metasploit",
|
||||
"key": "MetasploitMCP",
|
||||
"command": "uvx",
|
||||
"args": ["gc-metasploit", "--transport", "stdio"],
|
||||
"description": "MCP server for Metasploit Framework with exploit execution, payload generation, and session management.",
|
||||
"exe_name": None, # No local executable needed - uses uvx package
|
||||
"env_var": "MSF_PASSWORD",
|
||||
"env_extra": {
|
||||
"MSF_SERVER": "127.0.0.1",
|
||||
"MSF_PORT": "55553",
|
||||
"MSF_SSL": "false",
|
||||
"PAYLOAD_SAVE_DIR": "knowledge"
|
||||
},
|
||||
"homepage": "https://github.com/GH05TCREW/MetasploitMCP"
|
||||
},
|
||||
{
|
||||
"name": "Nmap Scanner",
|
||||
"key": "Nmap",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-nmap-mcp"],
|
||||
"description": "MCP server for interacting with Nmap network scanner to discover hosts and services on a network.",
|
||||
"exe_name": "nmap.exe",
|
||||
"env_var": "NMAP_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-nmap-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Nuclei Scanner",
|
||||
"key": "Nuclei",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-nuclei-mcp"],
|
||||
"description": "MCP server for vulnerability scanning using Nuclei's template-based detection engine.",
|
||||
"exe_name": "nuclei.exe",
|
||||
"env_var": "NUCLEI_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-nuclei-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Scout Suite",
|
||||
"key": "ScoutSuite",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-scoutsuite-mcp"],
|
||||
"description": "MCP server for cloud security auditing using the Scout Suite tool.",
|
||||
"exe_name": "scout.py",
|
||||
"env_var": "SCOUTSUITE_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-scoutsuite-mcp"
|
||||
},
|
||||
{
|
||||
"name": "shuffledns",
|
||||
"key": "ShuffleDNS",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-shuffledns-mcp"],
|
||||
"description": "MCP server for high-speed DNS brute-forcing and resolution using the shuffledns tool.",
|
||||
"exe_name": "shuffledns",
|
||||
"env_var": "SHUFFLEDNS_PATH",
|
||||
"env_extra": {
|
||||
"MASSDNS_PATH": ""
|
||||
},
|
||||
"homepage": "https://www.npmjs.com/package/gc-shuffledns-mcp"
|
||||
},
|
||||
{
|
||||
"name": "SQLMap",
|
||||
"key": "SQLMap",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-sqlmap-mcp"],
|
||||
"description": "MCP server for conducting automated SQL injection detection and exploitation using SQLMap.",
|
||||
"exe_name": "sqlmap.py",
|
||||
"env_var": "SQLMAP_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-sqlmap-mcp"
|
||||
},
|
||||
{
|
||||
"name": "SSL Scanner",
|
||||
"key": "SSLScan",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-sslscan-mcp"],
|
||||
"description": "MCP server for analyzing SSL/TLS configurations and identifying security issues.",
|
||||
"exe_name": "sslscan.exe",
|
||||
"env_var": "SSLSCAN_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-sslscan-mcp"
|
||||
},
|
||||
{
|
||||
"name": "Wayback URLs",
|
||||
"key": "WaybackURLs",
|
||||
"command": "npx",
|
||||
"args": ["-y", "gc-waybackurls-mcp"],
|
||||
"description": "MCP server for discovering historical URLs from the Wayback Machine archive.",
|
||||
"exe_name": "waybackurls.exe",
|
||||
"env_var": "WAYBACKURLS_PATH",
|
||||
"homepage": "https://www.npmjs.com/package/gc-waybackurls-mcp"
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
|
||||
def check_npm_installed():
|
||||
"""Check if npm is installed"""
|
||||
try:
|
||||
result = shutil.which("npm")
|
||||
return result is not None
|
||||
except:
|
||||
return False
|
||||
|
||||
def main():
|
||||
print(f"{Fore.GREEN}===================== GHOSTCREW MCP SERVER CONFIGURATION ====================={Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}This tool will help you configure the MCP servers for your GHOSTCREW installation.{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Auto-discovery will attempt to find tools automatically in your system PATH.{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}You can confirm, decline, or provide custom paths as needed.{Style.RESET_ALL}")
|
||||
print()
|
||||
|
||||
# Check if npm is installed
|
||||
if not check_npm_installed():
|
||||
print(f"{Fore.RED}Warning: npm doesn't appear to be installed. MCP servers use Node.js and npm.{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}You may need to install Node.js from: https://nodejs.org/{Style.RESET_ALL}")
|
||||
cont = input(f"{Fore.YELLOW}Continue anyway? (yes/no): {Style.RESET_ALL}").strip().lower()
|
||||
if cont != "yes":
|
||||
print(f"{Fore.RED}Configuration cancelled.{Style.RESET_ALL}")
|
||||
return
|
||||
|
||||
# Check if mcp.json exists and load it
|
||||
mcp_config = {"servers": []}
|
||||
if os.path.exists("mcp.json"):
|
||||
try:
|
||||
with open("mcp.json", 'r') as f:
|
||||
mcp_config = json.load(f)
|
||||
print(f"{Fore.GREEN}Loaded existing mcp.json with {len(mcp_config.get('servers', []))} server configurations.{Style.RESET_ALL}")
|
||||
except:
|
||||
print(f"{Fore.RED}Error loading existing mcp.json. Starting with empty configuration.{Style.RESET_ALL}")
|
||||
|
||||
configured_servers = []
|
||||
|
||||
print(f"{Fore.CYAN}Available tools:{Style.RESET_ALL}")
|
||||
for i, server in enumerate(MCP_SERVERS):
|
||||
print(f"{i+1}. {server['name']} - {server['description']}")
|
||||
|
||||
print()
|
||||
print(f"{Fore.YELLOW}Select tools to configure (comma-separated numbers, 'all' for all tools, or 'none' to skip):{Style.RESET_ALL}")
|
||||
selection = input().strip().lower()
|
||||
|
||||
selected_indices = []
|
||||
if selection == "all":
|
||||
selected_indices = list(range(len(MCP_SERVERS)))
|
||||
elif selection != "none":
|
||||
try:
|
||||
for part in selection.split(","):
|
||||
idx = int(part.strip()) - 1
|
||||
if 0 <= idx < len(MCP_SERVERS):
|
||||
selected_indices.append(idx)
|
||||
except:
|
||||
print(f"{Fore.RED}Invalid selection. Please enter comma-separated numbers.{Style.RESET_ALL}")
|
||||
return
|
||||
|
||||
for idx in selected_indices:
|
||||
server = MCP_SERVERS[idx]
|
||||
print(f"\n{Fore.CYAN}Configuring {server['name']}:{Style.RESET_ALL}")
|
||||
|
||||
# Unified tool configuration - handles all tools generically
|
||||
env_vars = {}
|
||||
|
||||
# Handle main executable and environment variable
|
||||
if server.get('exe_name'):
|
||||
# Try to auto-discover the executable
|
||||
exe_path = auto_discover_tool_path(server)
|
||||
|
||||
if exe_path:
|
||||
# Verify the path exists
|
||||
if not os.path.exists(exe_path):
|
||||
print(f"{Fore.YELLOW}Warning: The specified path does not exist: {exe_path}{Style.RESET_ALL}")
|
||||
cont = input(f" Continue anyway? (yes/no, default: no): ").strip().lower()
|
||||
if cont != "yes":
|
||||
print(f" {Fore.YELLOW}Skipping {server['name']}.{Style.RESET_ALL}")
|
||||
continue
|
||||
|
||||
# Set the main environment variable
|
||||
if server.get('env_var'):
|
||||
env_vars[server['env_var']] = exe_path
|
||||
else:
|
||||
# Executable not found and user didn't provide manual path
|
||||
print(f"{Fore.YELLOW}Skipping {server['name']} - executable not found.{Style.RESET_ALL}")
|
||||
continue
|
||||
elif server.get('env_var'):
|
||||
# Tool has no executable but needs a main environment variable (like Metasploit)
|
||||
value = input(f"Enter value for {server['env_var']} (default: ): ").strip()
|
||||
|
||||
if value:
|
||||
env_vars[server['env_var']] = value
|
||||
else:
|
||||
print(f"{Fore.YELLOW}Skipping {server['name']} - {server['env_var']} required.{Style.RESET_ALL}")
|
||||
continue
|
||||
else:
|
||||
# Tool requires no executable (like Certificate Transparency)
|
||||
print(f"{Fore.GREEN}{server['name']} requires no local executable.{Style.RESET_ALL}")
|
||||
|
||||
# Handle additional environment variables
|
||||
if 'env_extra' in server:
|
||||
for extra_var, default_value in server['env_extra'].items():
|
||||
if extra_var == "MASSDNS_PATH":
|
||||
# Special auto-discovery for massdns
|
||||
print(f"\n{Fore.CYAN}Also configuring massdns for {server['name']}...{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Searching for massdns...{Style.RESET_ALL}")
|
||||
|
||||
massdns_path = find_tool_path("massdns")
|
||||
if massdns_path:
|
||||
print(f"{Fore.GREEN}Found: {massdns_path}{Style.RESET_ALL}")
|
||||
choice = input(f" Use this path? (yes/no): ").strip().lower()
|
||||
if choice != 'y' and choice != 'yes':
|
||||
massdns_path = input(f" Enter path to massdns manually: ").strip()
|
||||
else:
|
||||
print(f"{Fore.YELLOW}massdns not found in PATH{Style.RESET_ALL}")
|
||||
massdns_path = input(f" Enter path to massdns manually (or press Enter to skip): ").strip()
|
||||
|
||||
if massdns_path:
|
||||
env_vars[extra_var] = massdns_path
|
||||
else:
|
||||
print(f"{Fore.YELLOW}Skipping {server['name']} - massdns path required.{Style.RESET_ALL}")
|
||||
continue
|
||||
else:
|
||||
# Handle all environment variables generically
|
||||
value = input(f"Enter value for {extra_var} (default: {default_value}): ").strip()
|
||||
env_vars[extra_var] = value if value else default_value
|
||||
|
||||
# Add to configured servers
|
||||
configured_servers.append({
|
||||
"name": server['name'],
|
||||
"params": {
|
||||
"command": server['command'],
|
||||
"args": server['args'],
|
||||
"env": env_vars
|
||||
},
|
||||
"cache_tools_list": True
|
||||
})
|
||||
print(f"{Fore.GREEN}{server['name']} configured successfully!{Style.RESET_ALL}")
|
||||
|
||||
# Update mcp.json
|
||||
if "servers" not in mcp_config:
|
||||
mcp_config["servers"] = []
|
||||
|
||||
if configured_servers:
|
||||
# Ask if user wants to replace or append
|
||||
if mcp_config["servers"]:
|
||||
replace = input(f"{Fore.YELLOW}Replace existing configurations or append new ones? (replace/append, default: append): {Style.RESET_ALL}").strip().lower()
|
||||
if replace == "replace":
|
||||
mcp_config["servers"] = configured_servers
|
||||
else:
|
||||
# Remove any duplicates by name
|
||||
existing_names = [s["name"] for s in mcp_config["servers"]]
|
||||
for server in configured_servers:
|
||||
if server["name"] in existing_names:
|
||||
# Replace existing configuration
|
||||
idx = existing_names.index(server["name"])
|
||||
mcp_config["servers"][idx] = server
|
||||
else:
|
||||
# Add new configuration
|
||||
mcp_config["servers"].append(server)
|
||||
else:
|
||||
mcp_config["servers"] = configured_servers
|
||||
|
||||
# Save to mcp.json
|
||||
with open("mcp.json", 'w') as f:
|
||||
json.dump(mcp_config, f, indent=2)
|
||||
|
||||
print(f"\n{Fore.GREEN}Configuration saved to mcp.json with {len(mcp_config['servers'])} server configurations.{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}You can now run the main application with: python main.py{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"\n{Fore.YELLOW}No tools were configured. Keeping existing configuration.{Style.RESET_ALL}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,222 +0,0 @@
|
||||
"""MCP (Model Context Protocol) server management for GHOSTCREW."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import List, Optional, Tuple
|
||||
from colorama import Fore, Style
|
||||
from config.constants import MCP_SESSION_TIMEOUT, MCP_CONFIG_FILE
|
||||
|
||||
|
||||
class MCPManager:
|
||||
"""Manages MCP server connections and configuration."""
|
||||
|
||||
def __init__(self, MCPServerStdio=None, MCPServerSse=None):
|
||||
"""
|
||||
Initialize the MCP manager.
|
||||
|
||||
Args:
|
||||
MCPServerStdio: MCP server stdio class
|
||||
MCPServerSse: MCP server SSE class
|
||||
"""
|
||||
self.MCPServerStdio = MCPServerStdio
|
||||
self.MCPServerSse = MCPServerSse
|
||||
self.server_instances = []
|
||||
self.connected_servers = []
|
||||
|
||||
@staticmethod
|
||||
def get_available_tools(connected_servers: List) -> List[str]:
|
||||
"""Get list of available/connected tool names."""
|
||||
return [server.name for server in connected_servers]
|
||||
|
||||
def load_mcp_config(self) -> List[dict]:
|
||||
"""Load MCP tool configurations from mcp.json."""
|
||||
available_tools = []
|
||||
try:
|
||||
with open(MCP_CONFIG_FILE, 'r', encoding='utf-8') as f:
|
||||
mcp_config = json.load(f)
|
||||
available_tools = mcp_config.get('servers', [])
|
||||
except FileNotFoundError:
|
||||
print(f"{Fore.YELLOW}mcp.json configuration file not found.{Style.RESET_ALL}")
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error loading MCP configuration file: {e}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Proceeding without MCP tools.{Style.RESET_ALL}")
|
||||
return available_tools
|
||||
|
||||
def display_tool_menu(self, available_tools: List[dict]) -> Optional[List[int]]:
|
||||
"""Display MCP tool selection menu and get user choice."""
|
||||
if not available_tools:
|
||||
print(f"{Fore.YELLOW}No MCP tools currently configured.{Style.RESET_ALL}")
|
||||
configure_now = input(f"{Fore.YELLOW}Would you like to add tools? (yes/no, default: no): {Style.RESET_ALL}").strip().lower()
|
||||
if configure_now == 'yes':
|
||||
print(f"\n{Fore.CYAN}Launching tool configuration...{Style.RESET_ALL}")
|
||||
os.system("python tools/configure_mcp.py")
|
||||
print(f"\n{Fore.GREEN}Tool configuration completed.{Style.RESET_ALL}")
|
||||
# Reload configuration and continue
|
||||
return "reload_and_continue"
|
||||
else:
|
||||
print(f"{Fore.YELLOW}Proceeding without MCP tools.{Style.RESET_ALL}")
|
||||
return []
|
||||
|
||||
print(f"\n{Fore.CYAN}Available MCP tools:{Style.RESET_ALL}")
|
||||
for i, server in enumerate(available_tools):
|
||||
print(f"{i+1}. {server['name']}")
|
||||
print(f"{len(available_tools)+1}. Configure new tools")
|
||||
print(f"{len(available_tools)+2}. Connect to all tools")
|
||||
print(f"{len(available_tools)+3}. Skip tool connection")
|
||||
print(f"{len(available_tools)+4}. Clear all MCP tools")
|
||||
|
||||
try:
|
||||
tool_choice = input(f"\n{Fore.YELLOW}Select option: {Style.RESET_ALL}").strip()
|
||||
|
||||
if not tool_choice: # Default to all tools
|
||||
return list(range(len(available_tools)))
|
||||
elif tool_choice == str(len(available_tools)+1): # Configure new tools
|
||||
print(f"\n{Fore.CYAN}Launching tool configuration...{Style.RESET_ALL}")
|
||||
os.system("python tools/configure_mcp.py")
|
||||
print(f"\n{Fore.GREEN}Tool configuration completed.{Style.RESET_ALL}")
|
||||
# Reload configuration and continue
|
||||
return "reload_and_continue"
|
||||
elif tool_choice == str(len(available_tools)+2): # Connect to all tools
|
||||
return list(range(len(available_tools)))
|
||||
elif tool_choice == str(len(available_tools)+3): # Skip tool connection
|
||||
return []
|
||||
elif tool_choice == str(len(available_tools)+4): # Clear all MCP tools
|
||||
if self.clear_mcp_tools():
|
||||
return "reload_and_continue"
|
||||
return []
|
||||
else: # Parse comma-separated list
|
||||
selected_indices = []
|
||||
for part in tool_choice.split(","):
|
||||
idx = int(part.strip()) - 1
|
||||
if 0 <= idx < len(available_tools):
|
||||
selected_indices.append(idx)
|
||||
return selected_indices
|
||||
except ValueError:
|
||||
print(f"{Fore.RED}Invalid selection. Defaulting to all tools.{Style.RESET_ALL}")
|
||||
return list(range(len(available_tools)))
|
||||
|
||||
def clear_mcp_tools(self) -> bool:
|
||||
"""Clear all MCP tools from configuration."""
|
||||
confirm = input(f"{Fore.YELLOW}Are you sure you want to clear all MCP tools? This will empty mcp.json (yes/no): {Style.RESET_ALL}").strip().lower()
|
||||
if confirm == "yes":
|
||||
try:
|
||||
# Create empty mcp.json file
|
||||
with open(MCP_CONFIG_FILE, 'w', encoding='utf-8') as f:
|
||||
json.dump({"servers": []}, f, indent=2)
|
||||
print(f"{Fore.GREEN}Successfully cleared all MCP tools. mcp.json has been reset.{Style.RESET_ALL}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error clearing MCP tools: {e}{Style.RESET_ALL}")
|
||||
return False
|
||||
|
||||
def initialize_servers(self, available_tools: List[dict], selected_indices: List[int]) -> None:
|
||||
"""Initialize selected MCP servers."""
|
||||
# Use the MCP classes passed during initialization
|
||||
if not self.MCPServerStdio or not self.MCPServerSse:
|
||||
raise ValueError("MCP server classes not provided during initialization")
|
||||
|
||||
print(f"{Fore.GREEN}Initializing selected MCP servers...{Style.RESET_ALL}")
|
||||
for idx in selected_indices:
|
||||
if idx < len(available_tools):
|
||||
server = available_tools[idx]
|
||||
print(f"{Fore.CYAN}Initializing {server['name']}...{Style.RESET_ALL}")
|
||||
try:
|
||||
if 'params' in server:
|
||||
mcp_server = self.MCPServerStdio(
|
||||
name=server['name'],
|
||||
params=server['params'],
|
||||
cache_tools_list=server.get('cache_tools_list', True),
|
||||
client_session_timeout_seconds=MCP_SESSION_TIMEOUT
|
||||
)
|
||||
elif 'url' in server:
|
||||
mcp_server = self.MCPServerSse(
|
||||
params={"url": server["url"]},
|
||||
cache_tools_list=server.get('cache_tools_list', True),
|
||||
name=server['name'],
|
||||
client_session_timeout_seconds=MCP_SESSION_TIMEOUT
|
||||
)
|
||||
else:
|
||||
print(f"{Fore.RED}Unknown MCP server configuration: {server}{Style.RESET_ALL}")
|
||||
continue
|
||||
self.server_instances.append(mcp_server)
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Error initializing {server['name']}: {e}{Style.RESET_ALL}")
|
||||
|
||||
async def connect_servers(self) -> List:
|
||||
"""Connect to initialized MCP servers."""
|
||||
if not self.server_instances:
|
||||
return []
|
||||
|
||||
print(f"{Fore.YELLOW}Connecting to MCP servers...{Style.RESET_ALL}")
|
||||
for mcp_server in self.server_instances:
|
||||
try:
|
||||
await mcp_server.connect()
|
||||
print(f"{Fore.GREEN}Successfully connected to MCP server: {mcp_server.name}{Style.RESET_ALL}")
|
||||
self.connected_servers.append(mcp_server)
|
||||
except Exception as e:
|
||||
print(f"{Fore.RED}Failed to connect to MCP server {mcp_server.name}: {e}{Style.RESET_ALL}")
|
||||
|
||||
if self.connected_servers:
|
||||
print(f"{Fore.GREEN}MCP server connection successful! Can use tools provided by {len(self.connected_servers)} servers.{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"{Fore.YELLOW}No MCP servers successfully connected. Proceeding without tools.{Style.RESET_ALL}")
|
||||
|
||||
return self.connected_servers
|
||||
|
||||
async def setup_mcp_tools(self, use_mcp: bool = False) -> Tuple[List, List]:
|
||||
"""
|
||||
Main method to setup MCP tools.
|
||||
|
||||
Args:
|
||||
use_mcp: Whether to use MCP tools
|
||||
|
||||
Returns:
|
||||
Tuple of (server_instances, connected_servers)
|
||||
"""
|
||||
if not use_mcp:
|
||||
print(f"{Fore.YELLOW}Proceeding without MCP tools.{Style.RESET_ALL}")
|
||||
return [], []
|
||||
|
||||
while True: # Loop to handle configuration and reload
|
||||
# Load available tools
|
||||
available_tools = self.load_mcp_config()
|
||||
|
||||
# Get user selection
|
||||
selected_indices = self.display_tool_menu(available_tools)
|
||||
|
||||
# Handle special cases
|
||||
if selected_indices is None:
|
||||
# Restart needed (e.g., after clearing tools)
|
||||
return self.server_instances, []
|
||||
elif selected_indices == "reload_and_continue":
|
||||
# Tools were configured, reload and show menu again
|
||||
continue
|
||||
else:
|
||||
# Normal selection, proceed with initialization
|
||||
break
|
||||
|
||||
# Initialize servers
|
||||
if selected_indices:
|
||||
self.initialize_servers(available_tools, selected_indices)
|
||||
|
||||
# Connect to servers
|
||||
connected = await self.connect_servers()
|
||||
|
||||
return self.server_instances, connected
|
||||
|
||||
async def cleanup_servers(self) -> None:
|
||||
"""Clean up MCP server resources."""
|
||||
if not self.server_instances:
|
||||
return
|
||||
|
||||
print(f"{Fore.YELLOW}Cleaning up MCP server resources...{Style.RESET_ALL}")
|
||||
|
||||
for mcp_server in self.server_instances:
|
||||
print(f"{Fore.YELLOW}Attempting to clean up server: {mcp_server.name}...{Style.RESET_ALL}", flush=True)
|
||||
try:
|
||||
await mcp_server.cleanup()
|
||||
print(f"{Fore.GREEN}Cleanup completed for {mcp_server.name}.{Style.RESET_ALL}", flush=True)
|
||||
except Exception:
|
||||
print(f"{Fore.RED}Failed to cleanup {mcp_server.name}.{Style.RESET_ALL}", flush=True)
|
||||
|
||||
print(f"{Fore.YELLOW}MCP server resource cleanup complete.{Style.RESET_ALL}")
|
||||
@@ -1 +0,0 @@
|
||||
"""User interface components for GHOSTCREW."""
|
||||
@@ -1,122 +0,0 @@
|
||||
"""Conversation history management for GHOSTCREW."""
|
||||
|
||||
from typing import List, Dict, Optional
|
||||
import tiktoken
|
||||
from config.app_config import app_config
|
||||
|
||||
|
||||
class ConversationManager:
|
||||
"""Manages conversation history and dialogue tracking."""
|
||||
|
||||
def __init__(self, max_tokens: int = 4000):
|
||||
"""
|
||||
Initialize the conversation manager.
|
||||
|
||||
Args:
|
||||
max_tokens: Maximum tokens to keep in history
|
||||
"""
|
||||
self.history: List[Dict[str, str]] = []
|
||||
self.max_tokens = max_tokens
|
||||
self.model_name = app_config.model_name
|
||||
|
||||
def add_dialogue(self, user_query: str, ai_response: str = "") -> None:
|
||||
"""
|
||||
Add a dialogue entry to the conversation history.
|
||||
|
||||
Args:
|
||||
user_query: The user's query
|
||||
ai_response: The AI's response (can be empty initially)
|
||||
"""
|
||||
dialogue = {
|
||||
"user_query": user_query,
|
||||
"ai_response": ai_response
|
||||
}
|
||||
self.history.append(dialogue)
|
||||
|
||||
# Trim history if it exceeds token limit
|
||||
self._trim_history()
|
||||
|
||||
def update_last_response(self, ai_response: str) -> None:
|
||||
"""
|
||||
Update the AI response for the last dialogue entry.
|
||||
|
||||
Args:
|
||||
ai_response: The AI's response to update
|
||||
"""
|
||||
if self.history:
|
||||
self.history[-1]["ai_response"] = ai_response
|
||||
|
||||
def get_history(self) -> List[Dict[str, str]]:
|
||||
"""Get the complete conversation history."""
|
||||
return self.history
|
||||
|
||||
def get_history_for_context(self) -> List[Dict[str, str]]:
|
||||
"""Get conversation history suitable for context."""
|
||||
return self.history
|
||||
|
||||
def estimate_tokens(self) -> int:
|
||||
"""
|
||||
Estimate the number of tokens in the conversation history.
|
||||
|
||||
Returns:
|
||||
Estimated token count
|
||||
"""
|
||||
try:
|
||||
encoding = tiktoken.encoding_for_model(self.model_name)
|
||||
return sum(
|
||||
len(encoding.encode(entry['user_query'])) +
|
||||
len(encoding.encode(entry.get('ai_response', '')))
|
||||
for entry in self.history
|
||||
)
|
||||
except Exception:
|
||||
# Fall back to approximate counting if tiktoken fails
|
||||
return sum(
|
||||
len(entry['user_query'].split()) +
|
||||
len(entry.get('ai_response', '').split())
|
||||
for entry in self.history
|
||||
)
|
||||
|
||||
def _trim_history(self) -> None:
|
||||
"""Trim history to keep token count under the limit."""
|
||||
while self.estimate_tokens() > self.max_tokens and len(self.history) > 1:
|
||||
self.history.pop(0)
|
||||
|
||||
def clear_history(self) -> None:
|
||||
"""Clear all conversation history."""
|
||||
self.history = []
|
||||
|
||||
def get_dialogue_count(self) -> int:
|
||||
"""Get the number of dialogues in history."""
|
||||
return len(self.history)
|
||||
|
||||
def get_workflow_conversation(self, start_index: int) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Get conversation history starting from a specific index.
|
||||
|
||||
Args:
|
||||
start_index: The index to start from
|
||||
|
||||
Returns:
|
||||
Subset of conversation history
|
||||
"""
|
||||
return self.history[start_index:]
|
||||
|
||||
def export_history(self) -> str:
|
||||
"""
|
||||
Export conversation history as formatted text.
|
||||
|
||||
Returns:
|
||||
Formatted conversation history
|
||||
"""
|
||||
if not self.history:
|
||||
return "No conversation history available."
|
||||
|
||||
output = []
|
||||
for i, entry in enumerate(self.history, 1):
|
||||
output.append(f"=== Dialogue {i} ===")
|
||||
output.append(f"User: {entry['user_query']}")
|
||||
if entry.get('ai_response'):
|
||||
output.append(f"AI: {entry['ai_response']}")
|
||||
output.append("")
|
||||
|
||||
return "\n".join(output)
|
||||
@@ -1,254 +0,0 @@
|
||||
"""Menu system and user interface components for GHOSTCREW."""
|
||||
|
||||
from typing import Optional, List, Tuple, Dict, Any
|
||||
from colorama import Fore, Style
|
||||
from config.constants import (
|
||||
MAIN_MENU_TITLE, INTERACTIVE_OPTION, AUTOMATED_OPTION,
|
||||
EXIT_OPTION, MULTI_LINE_PROMPT, MULTI_LINE_END_MARKER
|
||||
)
|
||||
|
||||
|
||||
class MenuSystem:
|
||||
"""Handles all menu displays and user input for GHOSTCREW."""
|
||||
|
||||
@staticmethod
|
||||
def display_main_menu(workflows_available: bool, has_connected_servers: bool) -> None:
|
||||
"""Display the main application menu."""
|
||||
print(f"\n{MAIN_MENU_TITLE}")
|
||||
print(f"1. {INTERACTIVE_OPTION}")
|
||||
|
||||
# Check if automated mode should be available
|
||||
if workflows_available and has_connected_servers:
|
||||
print(f"2. {AUTOMATED_OPTION}")
|
||||
elif workflows_available and not has_connected_servers:
|
||||
print(f"2. {Fore.LIGHTBLACK_EX}Workflows (requires MCP tools){Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"2. {Fore.LIGHTBLACK_EX}Workflows (workflows.py not found){Style.RESET_ALL}")
|
||||
|
||||
# Agent mode option
|
||||
if has_connected_servers:
|
||||
print(f"3. {Fore.YELLOW}Agent{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"3. {Fore.LIGHTBLACK_EX}Agent (requires MCP tools){Style.RESET_ALL}")
|
||||
|
||||
print(f"4. {EXIT_OPTION}")
|
||||
|
||||
@staticmethod
|
||||
def get_menu_choice(max_option: int = 4) -> str:
|
||||
"""Get user's menu selection."""
|
||||
return input(f"\n{Fore.GREEN}Select mode (1-{max_option}): {Style.RESET_ALL}").strip()
|
||||
|
||||
@staticmethod
|
||||
def display_interactive_mode_intro() -> None:
|
||||
"""Display introduction for interactive chat mode."""
|
||||
print(f"\n{Fore.CYAN}CHAT MODE{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Type your questions or commands. Use 'multi' for multi-line input.{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Type 'menu' to return to main menu.{Style.RESET_ALL}\n")
|
||||
|
||||
@staticmethod
|
||||
def display_agent_mode_intro() -> None:
|
||||
"""Display introduction for agent mode."""
|
||||
print(f"\n{Fore.CYAN}AGENT MODE{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}{'='*60}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}The AI agent will autonomously conduct a penetration test{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}using a dynamic Pentesting Task Tree (PTT) for strategic{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}decision making and maintaining context throughout the test.{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}{'='*60}{Style.RESET_ALL}\n")
|
||||
|
||||
@staticmethod
|
||||
def get_agent_mode_params() -> Optional[Dict[str, Any]]:
|
||||
"""Get parameters for agent mode initialization."""
|
||||
print(f"{Fore.CYAN}Agent Mode Setup{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Please provide the following information:{Style.RESET_ALL}\n")
|
||||
|
||||
# Get goal
|
||||
print(f"{Fore.YELLOW}1. Primary Goal{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}What is the main objective of this penetration test?{Style.RESET_ALL}")
|
||||
print(f"{Fore.LIGHTBLACK_EX}Example: 'Gain administrative access to the target system'{Style.RESET_ALL}")
|
||||
print(f"{Fore.LIGHTBLACK_EX}Example: 'Identify and exploit vulnerabilities in the web application'{Style.RESET_ALL}")
|
||||
goal = input(f"{Fore.GREEN}Goal: {Style.RESET_ALL}").strip()
|
||||
|
||||
if not goal:
|
||||
print(f"{Fore.RED}Goal is required.{Style.RESET_ALL}")
|
||||
return None
|
||||
|
||||
# Get target
|
||||
print(f"\n{Fore.YELLOW}2. Target Information{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Specify the target system, network, or application.{Style.RESET_ALL}")
|
||||
print(f"{Fore.LIGHTBLACK_EX}Example: '192.168.1.100' or 'example.com' or '192.168.1.0/24'{Style.RESET_ALL}")
|
||||
target = input(f"{Fore.GREEN}Target: {Style.RESET_ALL}").strip()
|
||||
|
||||
if not target:
|
||||
print(f"{Fore.RED}Target is required.{Style.RESET_ALL}")
|
||||
return None
|
||||
|
||||
# Get constraints
|
||||
constraints = {}
|
||||
print(f"\n{Fore.YELLOW}3. Constraints/Scope (Optional){Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Any limitations or specific requirements?{Style.RESET_ALL}")
|
||||
|
||||
# Iteration limit
|
||||
print(f"\n{Fore.WHITE}Iteration Limit:{Style.RESET_ALL}")
|
||||
print("How many iterations should the agent run?")
|
||||
print("Each iteration involves task selection, execution, and tree updates.")
|
||||
print("Recommended: 10-30 iterations for thorough testing")
|
||||
print("Set to 0 to run until goal is achieved or no more actions available")
|
||||
iteration_limit_input = input(f"{Fore.GREEN}Iteration limit (default: 20): {Style.RESET_ALL}").strip()
|
||||
|
||||
try:
|
||||
iteration_limit = int(iteration_limit_input) if iteration_limit_input else 20
|
||||
# Allow 0 for unlimited, but cap maximum at 200 for safety
|
||||
iteration_limit = max(0, min(200, iteration_limit))
|
||||
constraints['iteration_limit'] = iteration_limit
|
||||
except ValueError:
|
||||
constraints['iteration_limit'] = 20
|
||||
print(f"{Fore.YELLOW}Invalid input, using default: 20{Style.RESET_ALL}")
|
||||
|
||||
# Additional notes
|
||||
notes = input(f"{Fore.GREEN}Additional notes/constraints (optional): {Style.RESET_ALL}").strip()
|
||||
if notes:
|
||||
constraints['notes'] = notes
|
||||
|
||||
# Confirm
|
||||
print(f"\n{Fore.CYAN}Configuration Summary:{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Goal: {goal}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Target: {target}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Iteration Limit: {constraints['iteration_limit']}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Constraints: {constraints}{Style.RESET_ALL}")
|
||||
|
||||
confirm = input(f"\n{Fore.YELLOW}Proceed with agent mode? (yes/no): {Style.RESET_ALL}").strip().lower()
|
||||
if confirm != 'yes':
|
||||
print(f"{Fore.YELLOW}Agent mode cancelled.{Style.RESET_ALL}")
|
||||
return None
|
||||
|
||||
return {
|
||||
'goal': goal,
|
||||
'target': target,
|
||||
'constraints': constraints
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_user_input() -> str:
|
||||
"""Get user input with prompt."""
|
||||
print(f"\n{Fore.GREEN}[>]{Style.RESET_ALL} ", end="")
|
||||
return input().strip()
|
||||
|
||||
@staticmethod
|
||||
def get_multi_line_input() -> Optional[str]:
|
||||
"""Get multi-line input from user."""
|
||||
print(f"{Fore.CYAN}Entering multi-line mode. Type your query across multiple lines.{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Press Enter on empty line to submit.{Style.RESET_ALL}")
|
||||
|
||||
lines = []
|
||||
while True:
|
||||
line = input()
|
||||
if line == "": # Empty line ends input
|
||||
break
|
||||
lines.append(line)
|
||||
|
||||
# Only proceed if they actually entered something
|
||||
if not lines:
|
||||
print(f"{Fore.YELLOW}No query entered in multi-line mode.{Style.RESET_ALL}")
|
||||
return None
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def display_no_query_message() -> None:
|
||||
"""Display message when no query is entered."""
|
||||
print(f"{Fore.YELLOW}No query entered. Please type your question.{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_ready_message() -> None:
|
||||
"""Display ready for next query message."""
|
||||
print(f"\n{Fore.CYAN}Ready for next query. Type 'quit', 'multi' for multi-line, or 'menu' for main menu.{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_exit_message() -> None:
|
||||
"""Display exit message."""
|
||||
print(f"\n{Fore.CYAN}Thank you for using GHOSTCREW, exiting...{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_workflow_requirements_message() -> None:
|
||||
"""Display message about automated workflow requirements."""
|
||||
print(f"\n{Fore.YELLOW}Workflows requires MCP tools to be configured and connected.{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Without real security tools, the AI would only generate simulated responses.{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Please configure MCP tools to use this feature.{Style.RESET_ALL}")
|
||||
input(f"{Fore.CYAN}Press Enter to continue...{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_agent_mode_requirements_message() -> None:
|
||||
"""Display message about agent mode requirements."""
|
||||
print(f"\n{Fore.YELLOW}Agent Mode requires MCP tools to be configured and connected.{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}The autonomous agent needs real security tools to execute PTT tasks.{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Please configure MCP tools to use this feature.{Style.RESET_ALL}")
|
||||
input(f"{Fore.CYAN}Press Enter to continue...{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def get_workflow_target() -> Optional[str]:
|
||||
"""Get target input for workflow execution."""
|
||||
target = input(f"{Fore.YELLOW}Enter target (IP, domain, or network): {Style.RESET_ALL}").strip()
|
||||
if not target:
|
||||
print(f"{Fore.RED}Target is required.{Style.RESET_ALL}")
|
||||
return None
|
||||
return target
|
||||
|
||||
@staticmethod
|
||||
def confirm_workflow_execution(workflow_name: str, target: str) -> bool:
|
||||
"""Confirm workflow execution with user."""
|
||||
confirm = input(f"{Fore.YELLOW}Execute '{workflow_name}' on {target}? (yes/no): {Style.RESET_ALL}").strip().lower()
|
||||
return confirm == 'yes'
|
||||
|
||||
@staticmethod
|
||||
def display_workflow_cancelled() -> None:
|
||||
"""Display workflow cancelled message."""
|
||||
print(f"{Fore.YELLOW}Workflow cancelled.{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_workflow_completed() -> None:
|
||||
"""Display workflow completion message."""
|
||||
print(f"\n{Fore.GREEN}Workflow completed successfully!{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def ask_generate_report() -> bool:
|
||||
"""Ask if user wants to generate a report."""
|
||||
response = input(f"\n{Fore.CYAN}Generate markdown report? (yes/no): {Style.RESET_ALL}").strip().lower()
|
||||
return response == 'yes'
|
||||
|
||||
@staticmethod
|
||||
def ask_save_raw_history() -> bool:
|
||||
"""Ask if user wants to save raw conversation history."""
|
||||
response = input(f"{Fore.YELLOW}Save raw conversation history? (yes/no, default: no): {Style.RESET_ALL}").strip().lower()
|
||||
return response == 'yes'
|
||||
|
||||
@staticmethod
|
||||
def display_report_generated(report_path: str) -> None:
|
||||
"""Display report generation success message."""
|
||||
print(f"\n{Fore.GREEN}Report generated: {report_path}{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Open the markdown file in any markdown viewer for best formatting{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_report_error(error: Exception) -> None:
|
||||
"""Display report generation error message."""
|
||||
print(f"\n{Fore.RED}Error generating report: {error}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Raw workflow data is still available in conversation history{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_invalid_choice() -> None:
|
||||
"""Display invalid choice message."""
|
||||
print(f"{Fore.RED}Invalid choice. Please select a valid option.{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_invalid_input() -> None:
|
||||
"""Display invalid input message."""
|
||||
print(f"{Fore.RED}Invalid input. Please enter a number.{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def display_operation_cancelled() -> None:
|
||||
"""Display operation cancelled message."""
|
||||
print(f"\n{Fore.YELLOW}Operation cancelled.{Style.RESET_ALL}")
|
||||
|
||||
@staticmethod
|
||||
def press_enter_to_continue() -> None:
|
||||
"""Wait for user to press enter."""
|
||||
input(f"\n{Fore.CYAN}Press Enter to continue...{Style.RESET_ALL}")
|
||||
@@ -1 +0,0 @@
|
||||
"""Workflow system for GHOSTCREW."""
|
||||
@@ -1,73 +0,0 @@
|
||||
# GHOSTCREW Workflows
|
||||
|
||||
def get_available_workflows():
|
||||
"""
|
||||
Get all available automated workflows.
|
||||
All workflows can use any configured tools - no restrictions.
|
||||
"""
|
||||
|
||||
workflows = {
|
||||
"reconnaissance": {
|
||||
"name": "Reconnaissance and Discovery",
|
||||
"description": "Comprehensive information gathering and target profiling",
|
||||
"steps": [
|
||||
"Perform comprehensive reconnaissance on {target}",
|
||||
"Discover subdomains and DNS information",
|
||||
"Scan for open ports and services",
|
||||
"Identify technology stack and fingerprints",
|
||||
"Gather historical data and archived content"
|
||||
]
|
||||
},
|
||||
|
||||
"web_application": {
|
||||
"name": "Web Application Security Assessment",
|
||||
"description": "Comprehensive web application penetration testing",
|
||||
"steps": [
|
||||
"Discover web directories and hidden content on {target}",
|
||||
"Test for SQL injection vulnerabilities",
|
||||
"Scan for web application vulnerabilities and misconfigurations",
|
||||
"Analyze SSL/TLS configuration and security",
|
||||
"Test for authentication and session management flaws",
|
||||
"Check for file inclusion and upload vulnerabilities"
|
||||
]
|
||||
},
|
||||
|
||||
"network_infrastructure": {
|
||||
"name": "Network Infrastructure Penetration Test",
|
||||
"description": "Network-focused penetration testing and exploitation",
|
||||
"steps": [
|
||||
"Scan network range {target} for live hosts and services",
|
||||
"Perform detailed service enumeration and version detection",
|
||||
"Scan for known vulnerabilities in discovered services",
|
||||
"Test for network service misconfigurations",
|
||||
"Attempt exploitation of discovered vulnerabilities",
|
||||
"Assess network segmentation and access controls"
|
||||
]
|
||||
},
|
||||
|
||||
"full_penetration_test": {
|
||||
"name": "Complete Penetration Test",
|
||||
"description": "Full-scope penetration testing methodology",
|
||||
"steps": [
|
||||
"Phase 1: Quick port scan to identify open services on {target}",
|
||||
"Phase 2: Service version detection on discovered ports",
|
||||
"Phase 3: Web service discovery and directory enumeration",
|
||||
"Phase 4: Focused vulnerability scanning of services",
|
||||
"Phase 5: Targeted exploitation of discovered vulnerabilities",
|
||||
"Phase 6: Post-exploitation enumeration if access gained",
|
||||
"Phase 7: Compile findings and remediation recommendations"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
return workflows
|
||||
|
||||
def get_workflow_by_key(workflow_key):
|
||||
"""Get a specific workflow by its key"""
|
||||
workflows = get_available_workflows()
|
||||
return workflows.get(workflow_key, None)
|
||||
|
||||
def list_workflow_names():
|
||||
"""Get a list of all workflow names for display"""
|
||||
workflows = get_available_workflows()
|
||||
return [(key, workflow["name"]) for key, workflow in workflows.items()]
|
||||
@@ -1,153 +0,0 @@
|
||||
"""Workflow execution engine for GHOSTCREW."""
|
||||
|
||||
import asyncio
|
||||
from typing import List, Dict, Any, Optional
|
||||
from colorama import Fore, Style
|
||||
from datetime import datetime
|
||||
from workflows.workflow_definitions import (
|
||||
get_available_workflows, get_workflow_by_key, list_workflow_names
|
||||
)
|
||||
from config.constants import (
|
||||
ERROR_NO_WORKFLOWS, ERROR_WORKFLOW_NOT_FOUND, WORKFLOW_TARGET_PROMPT,
|
||||
WORKFLOW_CONFIRM_PROMPT, WORKFLOW_CANCELLED_MESSAGE, WORKFLOW_COMPLETED_MESSAGE
|
||||
)
|
||||
from tools.mcp_manager import MCPManager
|
||||
|
||||
|
||||
class WorkflowEngine:
|
||||
"""Handles automated workflow execution."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the workflow engine."""
|
||||
self.workflows_available = self._check_workflows_available()
|
||||
|
||||
@staticmethod
|
||||
def _check_workflows_available() -> bool:
|
||||
"""Check if workflow definitions are available."""
|
||||
try:
|
||||
# Test import to verify module is available
|
||||
from workflows.workflow_definitions import get_available_workflows
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""Check if workflows are available."""
|
||||
return self.workflows_available
|
||||
|
||||
@staticmethod
|
||||
def show_automated_menu() -> Optional[List[tuple]]:
|
||||
"""Display the automated workflow selection menu."""
|
||||
try:
|
||||
print(f"\n{Fore.CYAN}WORKFLOWS{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}{'='*50}{Style.RESET_ALL}")
|
||||
|
||||
workflow_list = list_workflow_names()
|
||||
workflows = get_available_workflows()
|
||||
|
||||
for i, (key, name) in enumerate(workflow_list, 1):
|
||||
description = workflows[key]["description"]
|
||||
step_count = len(workflows[key]["steps"])
|
||||
print(f"{i}. {Fore.YELLOW}{name}{Style.RESET_ALL}")
|
||||
print(f" {Fore.WHITE}{description}{Style.RESET_ALL}")
|
||||
print(f" {Fore.CYAN}Steps: {step_count}{Style.RESET_ALL}")
|
||||
print()
|
||||
|
||||
print(f"{len(workflow_list)+1}. {Fore.RED}Back to Main Menu{Style.RESET_ALL}")
|
||||
|
||||
return workflow_list
|
||||
except Exception:
|
||||
print(f"{Fore.YELLOW}Error loading workflows.{Style.RESET_ALL}")
|
||||
return None
|
||||
|
||||
async def run_automated_workflow(
|
||||
self,
|
||||
workflow: Dict[str, Any],
|
||||
target: str,
|
||||
connected_servers: List[Any],
|
||||
conversation_history: List[Dict[str, str]],
|
||||
kb_instance: Any,
|
||||
run_agent_func: Any
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Execute a workflow.
|
||||
|
||||
Args:
|
||||
workflow: The workflow definition
|
||||
target: The target for the workflow
|
||||
connected_servers: List of connected MCP servers
|
||||
conversation_history: Conversation history list
|
||||
kb_instance: Knowledge base instance
|
||||
run_agent_func: Function to run agent queries
|
||||
|
||||
Returns:
|
||||
List of workflow results
|
||||
"""
|
||||
available_tools = MCPManager.get_available_tools(connected_servers)
|
||||
|
||||
print(f"\n{Fore.CYAN}Starting Automated Workflow: {workflow['name']}{Style.RESET_ALL}")
|
||||
print(f"{Fore.YELLOW}Target: {target}{Style.RESET_ALL}")
|
||||
print(f"{Fore.CYAN}Available Tools: {', '.join(available_tools) if available_tools else 'None'}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}Description: {workflow['description']}{Style.RESET_ALL}")
|
||||
print(f"{Fore.WHITE}{'='*60}{Style.RESET_ALL}")
|
||||
|
||||
results = []
|
||||
|
||||
for i, step in enumerate(workflow['steps'], 1):
|
||||
print(f"\n{Fore.CYAN}Step {i}/{len(workflow['steps'])}{Style.RESET_ALL}")
|
||||
formatted_step = step.format(target=target)
|
||||
print(f"{Fore.WHITE}{formatted_step}{Style.RESET_ALL}")
|
||||
|
||||
# Create comprehensive query for this step
|
||||
enhanced_query = f"""
|
||||
TARGET: {target}
|
||||
STEP: {formatted_step}
|
||||
|
||||
Execute this step and provide the results.
|
||||
"""
|
||||
|
||||
# Execute the step through the agent
|
||||
result = await run_agent_func(
|
||||
enhanced_query,
|
||||
connected_servers,
|
||||
history=conversation_history,
|
||||
streaming=True,
|
||||
kb_instance=kb_instance
|
||||
)
|
||||
|
||||
if result and hasattr(result, "final_output"):
|
||||
results.append({
|
||||
"step": i,
|
||||
"description": formatted_step,
|
||||
"output": result.final_output
|
||||
})
|
||||
|
||||
# Add to conversation history
|
||||
conversation_history.append({
|
||||
"user_query": enhanced_query,
|
||||
"ai_response": result.final_output
|
||||
})
|
||||
|
||||
print(f"{Fore.GREEN}Step {i} completed{Style.RESET_ALL}")
|
||||
|
||||
# Brief delay between steps
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Workflow completion summary
|
||||
print(f"{Fore.CYAN}Steps executed: {len(results)}/{len(workflow['steps'])}{Style.RESET_ALL}")
|
||||
|
||||
return results
|
||||
|
||||
def get_workflow(self, workflow_key: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get a workflow by its key."""
|
||||
try:
|
||||
return get_workflow_by_key(workflow_key)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def get_workflow_list(self) -> List[tuple]:
|
||||
"""Get list of available workflows."""
|
||||
try:
|
||||
return list_workflow_names()
|
||||
except Exception:
|
||||
return []
|
||||
Reference in New Issue
Block a user