diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 00000000..43fe72ae --- /dev/null +++ b/.devcontainer/Dockerfile @@ -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 diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 00000000..22e1192e --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -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" + ] +} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml new file mode 100644 index 00000000..92b14fb9 --- /dev/null +++ b/.devcontainer/docker-compose.yml @@ -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 diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt new file mode 100644 index 00000000..2409daf2 --- /dev/null +++ b/.devcontainer/requirements.txt @@ -0,0 +1,6 @@ +pymetasploit3 +networkx==2.5 +requests +wasabi +xmltodict +pydot==1.4.2 diff --git a/.gitignore b/.gitignore index 5101cdeb..593a7ca6 100644 --- a/.gitignore +++ b/.gitignore @@ -30,7 +30,7 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST - +.devcontainer/qdrant_storage/ # PyInstaller *.manifest *.spec diff --git a/src/cai/cli.py b/src/cai/cli.py index 4151faa2..b6cbe6e2 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -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 diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 5ef1cf8c..1ce0850b 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -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, diff --git a/src/cai/util.py b/src/cai/util.py index 161b8e5c..b7223fad 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -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