RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'pentestperf' into '0.4.0' 

Integrate Pentestperf in v0.4.0

See merge request aliasrobotics/alias_research/cai!179
This commit is contained in:
Luis Javier Navarrete Lozano 2025-05-22 13:21:29 +00:00
parent 56ae641348 7e32792a24
commit 63b2348f91
8 changed files with 442 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

108
.devcontainer/Dockerfile Normal file
View File

@ -0,0 +1,108 @@
FROM kalilinux/kali-rolling
# Set environment variable to non-interactive
ENV DEBIAN_FRONTEND=noninteractive
# [Optional] Uncomment this section to install additional OS packages.
RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
&& apt-get -y install --no-install-recommends \
net-tools python3 python3-pip python3-venv \
curl gnupg nmap iputils-ping ssh git \
graphviz pkg-config libhdf5-dev \
build-essential python3-dev \
asciinema dnsutils \
apt-transport-https ca-certificates \
wget dirb gobuster whatweb gfortran \
cmake libopenblas-dev sshpass seclists \
golang-go
# # Install Metasploit - path 1
# RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > /tmp/msfinstall \
# && chmod 755 /tmp/msfinstall \
# && /tmp/msfinstall
# # Install Metasploit - path 2
# RUN curl https://apt.metasploit.com/metasploit-framework.gpg.key | apt-key add -
# RUN echo "deb https://apt.metasploit.com/ lucid main" | tee /etc/apt/sources.list.d/metasploit-framework.list
# RUN apt-get update && apt-get install metasploit-framework
# Install Metasploit - path 3
RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstall && chmod 755 msfinstall && ./msfinstall
# Update Kali Linux repositories
RUN apt-get update \
&& apt-get upgrade -y
# # Create a virtual environment
# ENV VIRTUAL_ENV=/opt/venv
# RUN python3 -m venv $VIRTUAL_ENV
# ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image.
# # alternatively, this also helps avoid having to pull requirements from the internet every single time
# COPY requirements.txt /tmp/pip-tmp/
# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \
# && rm -rf /tmp/pip-tmp
# FIX for:
#
# Python packages system-wide, try apt install
# python3-xyz, where xyz is the package you are trying to
# install.
#
# see https://github.com/raspiblitz/raspiblitz/issues/4170
RUN mkdir -p /root/.pip && touch /root/.pip/pip.conf
RUN echo "[global]" > /root/.pip/pip.conf && \
echo "break-system-packages = true" >> /root/.pip/pip.conf
RUN pip3 install --upgrade setuptools
RUN pip3 install pandas \
opentelemetry-sdk opentelemetry-exporter-otlp \
mem0ai PyPDF2 sentence_transformers tf-keras \
pytest-repeat parameterized pytest-rerunfailures \
pytest-clarity tiktoken blobfile pytest-timeout \
invoke fabric
# Install rust compiler - dependency of outlines-core
RUN apt-get update && \
apt-get install -y curl && \
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \
. $HOME/.cargo/env && \
rustup default stable && \
rustup component add rustfmt clippy && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
ENV PATH="/root/.cargo/bin:${PATH}"
RUN wget https://raw.githubusercontent.com/adfoster-r7/metasploit-info/master/info/module_metadata.json -O /tmp/module_metadata.json
# NOTE: align dev env in Linux
# work
# RUN echo "192.168.2.1 host.docker.internal" >> /etc/hosts
# Update and upgrade packages
# Preconfigure console-setup and install kali-linux-headless fixers
RUN apt-get update && \
echo "console-setup console-setup/variant select Latin1 and Latin5 - western Europe and Turkic languages" | debconf-set-selections && \
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends kali-linux-headless
# Update pip again -> python3 -m pip install --upgrade pip causes crashes
# Install spacy deps - en_core_web_sm
# RUN python3 -m spacy download en_core_web_sm
# Install docker inside of dev env
RUN apt-get update && \
apt-get install -y apt-transport-https ca-certificates curl gnupg && \
curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg && \
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian bullseye stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && \
apt-get update && \
apt-get install -y docker-ce docker-ce-cli containerd.io cloc
# # Activate virtual environment by default
# RUN echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc
RUN cargo install --git https://github.com/asciinema/agg
# Remove system's python3 libs to avoid conflicts
RUN apt-get remove -y python3-jsonschema && \
apt-get remove -y python3-numpy && \
apt-get autoremove -y

86
.devcontainer/devcontainer.json Normal file
View File

@ -0,0 +1,86 @@
// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
// https://github.com/microsoft/vscode-dev-containers/tree/v0.241.1/containers/python-3
{
"name": "cai_devenv",
// "build": {
// "dockerfile": "Dockerfile",
// "context": "..",
// "args": {
// // Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6
// // Append -bullseye or -buster to pin to an OS version.
// // Use -bullseye variants on local on arm64/Apple Silicon.
// "VARIANT": "3.10-bullseye",
// // Options
// "NODE_VERSION": "lts/*"
// }
// },
"dockerComposeFile": ["./docker-compose.yml"],
"service": "devenv",
// "shutdownAction": "none", // don't shut down container when vscode is closed
"workspaceFolder": "/workspace",
// Configure tool-specific properties.
"customizations": {
// Configure properties specific to VS Code.
"vscode": {
// Set *default* container specific settings.json values on container create.
"settings": {
"python.defaultInterpreterPath": "/usr/local/bin/python",
"python.linting.enabled": true,
"python.linting.pylintEnabled": true,
"python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8",
"python.formatting.blackPath": "/usr/local/py-utils/bin/black",
"python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf",
"python.linting.banditPath": "/usr/local/py-utils/bin/bandit",
"python.linting.flake8Path": "/usr/local/py-utils/bin/flake8",
"python.linting.mypyPath": "/usr/local/py-utils/bin/mypy",
"python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle",
"python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle",
"python.linting.pylintPath": "/usr/local/py-utils/bin/pylint"
},
// Add the IDs of extensions you want installed when the container is created.
"extensions": [
"ms-python.python",
"ms-toolsai.jupyter-renderers",
"ms-toolsai.jupyter",
"ms-python.vscode-pylance",
"ms-azuretools.vscode-docker",
"ms-python.debugpy",
"ms-python.black-formatter",
"MS-vsliveshare.vsliveshare"
]
}
},
// Use 'forwardPorts' to make a list of ports inside the container available locally.
// "forwardPorts": [],
// Use 'preCreateCommand' to run commands before the container is created.
"preCreateCommand": "docker network prune -f && docker container prune -f",
// Use 'postCreateCommand' to run commands after the container is created.
// "postCreateCommand": "pip3 install --user -r requirements.txt", // doing it instead in
// the dockerfile to cache
// requirements in container
// Install cai via "pip3 install -e ."
// Extensions, NOTE: only one postCreateCommand is allowed, so we need to install all extensions here
"postCreateCommand": "pip3 install -e /workspace/",
// MSF setup, and RAG setup
"postStartCommand": [
"nohup", "msfrpcd", "-P", "cai", "&",
"&&",
"python3", "cai/ins/rag/agent_helper.py"
],
// // Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
// "remoteUser": "vscode"
"runArgs": [
"--privileged",
"-e", "DISPLAY=host.docker.internal:0",
"-v", "/tmp/.X11-unix:/tmp/.X11-unix"
]
}

78
.devcontainer/docker-compose.yml Normal file
View File

@ -0,0 +1,78 @@
version: '3'
#################
# SERVICES
#################
services:
# Developer environment
devenv:
build:
context: ..
dockerfile: .devcontainer/Dockerfile
# platform: linux/amd64 # TODO: uncomment this when we have a x86_64 machine
volumes:
# Mount the root folder that contains .git
- ..:/workspace:cached
# Mount Docker socket to enable Docker-in-Docker
- /var/run/docker.sock:/var/run/docker.sock
command: /bin/sh -c "while sleep 10; do :; done"
networks:
cainet:
ipv4_address: 192.168.2.5
cap_add:
- NET_ADMIN
- SYS_ADMIN
privileged: true
# # Vulnhub vulnerable machine (Hackable II) SUGGESTION: Comment this and use PentestPerf
# hackableii:
# image: vmayoral/vulnhub:hackableii
# command: |
# /bin/bash -c "
# rm -rf /var/lock && \
# mkdir -p /var/lock && \
# chmod 755 /var/lock && \
# /etc/init.d/apache2 start && \
# /etc/init.d/ssh start && \
# /etc/init.d/runproftpd.sh && \
# /etc/init.d/php7.0-fpm start && \
# while sleep 1; do :; done
# "
# networks:
# cainet:
# ipv4_address: 192.168.2.11
# mac_address: 08:00:27:85:55:86
# # Vulnhub vulnerable machine (Bob) SUGGESTION: Comment this and use PentestPerf
# bob:
# image: vmayoral/vulnhub:bob
# command: |
# /bin/bash -c "rm -r /var/lock; mkdir -p /var/lock; chmod 755 /var/lock; /etc/init.d/apache2 start; /etc/init.d/ssh start; while sleep 10; do :; done"
# ports: # map port in the container to the host systems
# - "8080:80"
# networks:
# cainet:
# ipv4_address: 192.168.2.12
# mac_address: 08:00:27:cb:07:d4
# Qdrant vector database
qdrant:
image: qdrant/qdrant
ports:
- "6333:6333"
- "6334:6334"
volumes:
- ./qdrant_storage:/qdrant/storage:z
networks:
cainet:
ipv4_address: 192.168.2.13
#################
# NETWORKS
#################
networks:
cainet:
ipam:
driver: default
config:
- subnet: 192.168.2.0/24

6
.devcontainer/requirements.txt Normal file
View File

@ -0,0 +1,6 @@
pymetasploit3
networkx==2.5
requests
wasabi
xmltodict
pydot==1.4.2

2
.gitignore vendored
View File

@ -30,7 +30,7 @@ share/python-wheels/
.installed.cfg
*.egg
MANIFEST
.devcontainer/qdrant_storage/
# PyInstaller
*.manifest
*.spec

54
src/cai/cli.py
View File

@ -134,7 +134,9 @@ from cai.util import (
start_idle_timer,
stop_idle_timer,
start_active_timer,
stop_active_timer
stop_active_timer,
setup_ctf,
check_flag,
)
# CAI REPL imports
@ -151,6 +153,16 @@ from cai.internal.components.metrics import process_metrics
# Add import for parallel configs at the top of the file
from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig
from cai import is_pentestperf_available
ctf_global = None
messages_ctf = ""
previous_ctf_name = os.getenv('CTF_NAME', None)
if is_pentestperf_available() and os.getenv('CTF_NAME', None):
ctf, messages_ctf = setup_ctf()
ctf_global = ctf
if os.getenv('CTF_INSIDE', 'True').lower() == 'false':
container_id = ""
os.environ['CAI_ACTIVE_CONTAINER'] = container_id
# Load environment variables from .env file
load_dotenv()
@ -187,7 +199,7 @@ agent = Agent(
)
)
def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')):
def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), force_until_flag=False):
"""
Run a simple interactive CLI loop for CAI.
@ -247,6 +259,18 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')):
turn_limit_reached = False
while True:
# Check if the ctf name has changed and instanciate the ctf
global previous_ctf_name
global ctf_global
global messages_ctf
if previous_ctf_name != os.getenv('CTF_NAME', None):
if is_pentestperf_available():
if ctf_global:
ctf_global.stop_ctf()
ctf, messages_ctf = setup_ctf()
ctf_global = ctf
previous_ctf_name = os.getenv('CTF_NAME', None)
# Check if CAI_MAX_TURNS has been updated via /config
current_max_turns = os.getenv('CAI_MAX_TURNS', 'inf')
if current_max_turns != str(prev_max_turns):
@ -303,15 +327,22 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')):
agent.model.set_agent_name(get_agent_short_name(agent))
except Exception as e:
console.print(f"[red]Error switching agent: {str(e)}[/red]")
# Get user input with command completion and history
user_input = get_user_input(
command_completer,
kb,
history_file,
get_toolbar_with_refresh,
current_text
)
if not force_until_flag and messages_ctf == "":
# Get user input with command completion and history
user_input = get_user_input(
command_completer,
kb,
history_file,
get_toolbar_with_refresh,
current_text
) + messages_ctf
else:
user_input = messages_ctf
if os.getenv('CTF_INSIDE', 'True').lower() == 'false':
user_input += ctf_global.get_ip()
idle_time += time.time() - idle_start_time
# Stop measuring user idle time and start measuring active time
@ -475,6 +506,9 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')):
# Prevent duplicate cost display from the COST_TRACKER exit handler
os.environ["CAI_COST_DISPLAYED"] = "true"
if (is_pentestperf_available() and os.getenv('CTF_NAME', None)):
ctf.stop_ctf()
except Exception:
pass
break

14
src/cai/tools/common.py
View File

@ -375,7 +375,7 @@ def terminate_session(session_id):
def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None):
"""Runs command in CTF env, changing to workspace_dir first."""
target_dir = workspace_dir or _get_workspace_dir()
full_command = f"cd '{target_dir}' && {command}"
full_command = f"{command}"
original_cmd_for_msg = command # For logging
context_msg = f"(ctf:{target_dir})"
try:
@ -713,6 +713,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
# Use the active timer during tool execution
stop_idle_timer()
start_active_timer()
from cai.cli import ctf_global
ctf = ctf_global
# Parse command into standard parts to ensure consistent naming
parts = command.strip().split(' ', 1)
@ -761,7 +764,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST'])
# --- Docker Container Execution ---
if active_container and not ctf and not is_ssh_env:
if active_container and not is_ssh_env:
container_id = active_container
container_workspace = _get_container_workspace_path()
context_msg = f"(docker:{container_id[:12]}:{container_workspace})"
@ -1016,7 +1019,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501
# --- CTF Execution ---
if ctf:
if ctf and os.getenv('CTF_INSIDE', "True").lower() == "true":
# If streaming is enabled and we have a call_id, show streaming UI for CTF too
if stream:
# Import the streaming utilities from util
@ -1039,8 +1043,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
call_id = start_tool_streaming(tool_name, tool_args, call_id)
target_dir = _get_workspace_dir()
full_command = f"cd '{target_dir}' && {command}"
#full_command = f"cd '{target_dir}' && {command}"
full_command = command
# Update with "executing" status
update_tool_streaming(
tool_name,

110
src/cai/util.py
View File

@ -3,6 +3,7 @@ Util model for CAI
"""
import os
import sys
import subprocess
import importlib.resources
import pathlib
import json
@ -29,6 +30,9 @@ from rich.box import ROUNDED
from rich.table import Table
import re
import uuid
from cai import is_pentestperf_available
if is_pentestperf_available():
import pentestperf as ptt
# Global timing variables for tracking active and idle time
_active_timer_start = None
@ -2887,3 +2891,109 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None,
# Mark the streaming session as complete
if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions:
cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True
def check_flag(output, ctf, challenge=None):
"""
Check if the CTF flag is present in the output.
Args:
output (str): The output to check for the flag.
ctf: The CTF environment object.
challenge (str, optional): The specific challenge to check.
Defaults to None.
Returns:
tuple: A tuple containing a boolean indicating if the flag was
found and the flag itself if found, otherwise None.
"""
# Get the challenge from the environment variable or default to the first
# challenge
challenge_key = os.getenv("CTF_CHALLENGE")
challenges = list(ctf.get_challenges().keys())
challenge = (
challenge_key
if challenge_key in challenges
else (challenges[0] if len(challenges) > 0 else None))
if ctf:
if ctf.check_flag(
output, challenge
): # check if the flag is in the output
flag = ctf.flags[challenge]
print(
color(
f"Flag found: {flag}",
fg="green") +
" in output " +
color(
f"{output}",
fg="blue"))
return True, flag
else:
print(color("CTF environment not found or provided", fg="yellow"))
return False, None
def setup_ctf():
"""Setup CTF environment if CTF_NAME is provided"""
ctf_name = os.getenv('CTF_NAME', None)
if not ctf_name:
print(color("CTF name not provided, necessary to run CTF", fg="white", bg="red"))
sys.exit(1)
print(color("Setting up CTF: ", fg="black", bg="yellow") +
color(ctf_name, fg="black", bg="yellow"))
ctf = ptt.ctf( # pylint: disable=I1101 # noqa
ctf_name,
subnet=os.getenv('CTF_SUBNET', "192.168.2.0/24"),
container_name="ctf_target",
ip_address=os.getenv('CTF_IP', "192.168.2.100"),
)
ctf.start_ctf()
# Get the challenge from the environment variable or default to the
# first challenge
challenge_key = os.getenv('CTF_CHALLENGE') # TODO:
challenges = list(ctf.get_challenges().keys())
challenge = challenge_key if challenge_key in challenges else (
challenges[0] if len(challenges) > 0 else None)
# Use the user master template
messages = Template(
filename="src/cai/prompts/core/user_master_template.md").render(
ctf=ctf,
challenge=challenge,
ip=ctf.get_ip() if ctf else None,
)
print( color(
"Testing CTF: ",
fg="black",
bg="yellow") +
color(
ctf.name,
fg="black",
bg="yellow"))
if not challenge_key or challenge_key not in challenges:
print(
color(
"No challenge provided or challenge not found. Attempting to use the first challenge.",
fg="white",
bg="blue"))
if challenge:
print(
color(
"Testing challenge: ",
fg="white",
bg="blue") +
color(
"'" +
challenge +
"' (" +
repr(
ctf.flags[challenge]) +
")",
fg="white",
bg="blue"))
return ctf, messages