diff --git a/.env.example b/.env.example new file mode 100644 index 00000000..e570b8b5 --- /dev/null +++ b/.env.example @@ -0,0 +1 @@ +OPENAI_API_KEY= diff --git a/.gitignore b/.gitignore index 8a225e24..74dc99a6 100644 --- a/.gitignore +++ b/.gitignore @@ -145,3 +145,4 @@ pentestperf cai.png cai.dot cai*.png +test_20* diff --git a/README.md b/README.md index 34f4cb8d..2ca8a967 100644 --- a/README.md +++ b/README.md @@ -355,7 +355,9 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa # Plan of development - [x] Dev container - [x] pre-commit hooks -- [ ] A first example with agents - picoctf_static_flag +- [x] A first example with agents - picoctf_static_flag + - [ ] polish aesthetics (coloring with alpha past content, brief/summarized) + - [ ] add tests - [ ] CI - [ ] Tracing - [ ] Graph/flow and other abstractions diff --git a/cai/core.py b/cai/core.py index 73337962..f041bf8b 100644 --- a/cai/core.py +++ b/cai/core.py @@ -14,7 +14,6 @@ import copy import json from collections import defaultdict from typing import List - # Package/library imports from openai import OpenAI # pylint: disable=import-error @@ -36,14 +35,18 @@ __CTX_VARS_NAME__ = "context_variables" class CAI: """ - Main class for the CAI library. + Cybersecurity AI (CAI) object """ - def __init__(self, client=None, - base_url="http://localhost:8000/v1", api_key="alias"): + def __init__(self, + client=None, + base_url="http://host.docker.internal:8000/v1", + api_key="alias", + ctf=None): if not client: client = OpenAI(base_url=base_url, api_key=api_key) self.client = client + self.ctf = ctf def get_chat_completion( # pylint: disable=too-many-arguments self, @@ -91,7 +94,13 @@ class CAI: def handle_function_result(self, result, debug) -> Result: """ - Handle the result of a function call. + Handle the result of a function call by + converting it into a standardized Result type. + + The Result type encapsulates the possible + return values (Result, Agent, or context variables) + that functions can produce into a consistent + format for the framework to process. """ match result: case Result() as result: @@ -118,7 +127,37 @@ class CAI: debug: bool, ) -> Response: """ - Handle the tool calls for the given agent. + Execute and handle tool calls made by the AI agent. + + Processes a list of tool calls by: + 1. Looking up each function in the provided function map + 2. Handling missing tools gracefully by skipping them + 3. Parsing and validating function arguments + 4. Executing functions with provided arguments and context + 5. Processing results into standardized Response format + 6. Accumulating results from multiple tool calls + + Args: + tool_calls (List[ChatCompletionMessageToolCall]): Tool + calls requested by AI agent + functions (List[AgentFunction]): Available functions + that can be called + context_variables (dict): Context variables to pass + to functions + debug (bool): Flag to enable debug logging + + Returns: + Response: Object containing: + messages (List): Tool call results + agent (Optional[Agent]): Updated agent + if returned by a function + context_variables (dict): Updated context variables + + Note: + Results from multiple tool calls are accumulated + into a single Response. + Context variables are updated iteratively as + functions are called. """ function_map = {f.__name__: f for f in functions} partial_response = Response( @@ -140,12 +179,18 @@ class CAI: continue args = json.loads(tool_call.function.arguments) debug_print( - debug, f"Processing tool call: {name} with arguments {args}") + debug, + "Processing tool call", + name, + "with arguments", + args) func = function_map[name] # pass context_variables to agent functions if __CTX_VARS_NAME__ in func.__code__.co_varnames: args[__CTX_VARS_NAME__] = context_variables + if self.ctf: + args["ctf"] = self.ctf raw_result = function_map[name](**args) result: Result = self.handle_function_result(raw_result, debug) @@ -175,6 +220,9 @@ class CAI: ): """ Run the cai and stream the results. + + The key difference from run() is that this streams results + incrementally, while run() returns everything at once. """ active_agent = agent context_variables = copy.deepcopy(context_variables) diff --git a/cai/tools/cli.py b/cai/tools/cli.py new file mode 100644 index 00000000..d2a71f89 --- /dev/null +++ b/cai/tools/cli.py @@ -0,0 +1,36 @@ +""" +Here are the CLI tools for executing commands, they should usually return +with run_command, but you can be creative +""" + +from cai.tools.common import run_command # pylint: disable=E0401 + + +def list_dir(path: str, args: str = "", ctf=None) -> str: + """ + List the contents of a directory. + by def . + Args: + path: The directory path to list contents from + args: Additional arguments to pass to the ls command + + Returns: + str: The output of running the ls command + """ + command = f'ls {path} {args}' + return run_command(command, ctf=ctf) + + +def cat_file(file_path: str, args: str = "", ctf=None) -> str: + """ + Display the contents of a file. + + Args: + args: Additional arguments to pass to the cat command + file_path: Path to the file to display contents of + + Returns: + str: The output of running the cat command + """ + command = f'cat {args} {file_path} ' + return run_command(command, ctf=ctf) diff --git a/cai/tools/common.py b/cai/tools/common.py new file mode 100644 index 00000000..3885cf28 --- /dev/null +++ b/cai/tools/common.py @@ -0,0 +1,51 @@ +""" +Basic utilities for executing commands inside or outside of virtual +containers. +""" + +import subprocess # nosec B404 +from wasabi import color # pylint: disable=import-error + + +def _run_ctf(ctf, command, stdout=True): + try: + # Ensure the command is executed in a shell that supports command + # chaining + output = ctf.get_shell(command) + # exploit_logger.log_ok() + + if stdout: + print("\033[32m" + output + "\033[0m") + return output + except Exception as e: # pylint: disable=broad-except + print(color(f"Error executing CTF command: {e}", fg="red")) + # exploit_logger.log_error(str(e)) + return f"Error executing CTF command: {str(e)}" + + +def _run_attacker_machine(command, stdout=True): + try: + # nosec B602 - shell=True is required for command chaining + result = subprocess.run( + command, + shell=True, # nosec B602 + capture_output=True, + text=True, + check=True) + output = result.stdout + if stdout: + print("\033[32m" + output + "\033[0m") + return output + except Exception as e: # pylint: disable=broad-except + print(color(f"Error executing local command: {e}", fg="red")) + return f"Error executing local command: {str(e)}" + + +def run_command(command, ctf=None, stdout=True): + """ + Run command either in CTF container or + on the local attacker machine + """ + if ctf: + return _run_ctf(ctf, command, stdout) + return _run_attacker_machine(command, stdout) diff --git a/cai/types.py b/cai/types.py index b1913ec5..52a3b872 100644 --- a/cai/types.py +++ b/cai/types.py @@ -31,7 +31,10 @@ class Agent(BaseModel): # pylint: disable=too-few-public-methods class Response(BaseModel): # pylint: disable=too-few-public-methods """ - Represents a response from the CAI. + Represents a response back to the user from the CAI. + + NOTE: This happens within the run() chain, after "Ending turn" + in the CAI. """ messages: List = [] diff --git a/cai/util.py b/cai/util.py index 6e956cd5..b5d362f8 100644 --- a/cai/util.py +++ b/cai/util.py @@ -4,17 +4,142 @@ This module contains utility functions for the CAI library. import inspect from datetime import datetime +from typing import Any +import json + +# ANSI color codes in a nice, readable palette +COLORS = { + 'timestamp': '\033[38;5;75m', # Light blue + 'bracket': '\033[38;5;247m', # Light gray + 'intro': '\033[38;5;141m', # Light purple + 'object': '\033[38;5;215m', # Light orange + 'arg_key': '\033[38;5;228m', # Light yellow + 'arg_value': '\033[38;5;180m', # Light tan + 'function': '\033[38;5;219m', # Pink + 'tool': '\033[38;5;147m', # Soft purple + 'reset': '\033[0m' +} -def debug_print(debug: bool, *args: str) -> None: +def format_value(value: Any) -> str: """ - Print debug messages if debug mode is enabled. + Format a value for debug printing with appropriate colors. + """ + # Handle ChatCompletionMessage objects + if hasattr( + value, '__class__') and 'ChatCompletionMessage' in value.__class__.__name__: # noqa: E501 # pylint: disable=C0301 + return format_chat_completion(value) + + # Handle lists + if isinstance(value, list): # pylint: disable=R1705 + items = [] + for item in value: + if isinstance(item, dict): + # Format dictionary items in the list + dict_items = [ + f"\n { + COLORS['arg_key']}{k}{ + COLORS['reset']}: { + format_value(v)}" + for k, v in item.items() + ] + items.append("{" + ",".join(dict_items) + "\n }") + else: + items.append(format_value(item)) + return f"[\n {','.join(items)}\n]" + + # Handle dictionaries + elif isinstance(value, dict): + formatted_items = [ + f"{COLORS['arg_key']}{k}{COLORS['reset']}: {format_value(v)}" + for k, v in value.items() + ] + return "{ " + ", ".join(formatted_items) + " }" + + # Handle basic types + else: + return f"{COLORS['arg_value']}{str(value)}{COLORS['reset']}" + + +def format_chat_completion(msg) -> str: + """ + Format a ChatCompletionMessage object with proper indentation and colors. + """ + # Convert message to dict and handle OpenAI types + try: + msg_dict = json.loads(msg.model_dump_json()) + except AttributeError: + msg_dict = msg.__dict__ + + # Clean up the dictionary + msg_dict = {k: v for k, v in msg_dict.items() if v is not None} + + # Format with json.dumps for consistent indentation + formatted_json = json.dumps(msg_dict, indent=2) + + # Color the different parts + colored_lines = [] + for line in formatted_json.split('\n'): + if ':' in line: + key, value = line.split(':', 1) + # Handle nested structures + if value.strip() in ['{', '[', '}', ']']: + colored_lines.append(f"{COLORS['arg_key']}{key}{ + COLORS['reset']}:{value}") + else: + colored_lines.append( + f"{COLORS['arg_key']}{key}{COLORS['reset']}: " + f"{COLORS['arg_value']}{value.strip()}{COLORS['reset']}" + ) + else: + colored_lines.append(line) + + return f"\n {COLORS['object']}ChatCompletionMessage{ + COLORS['reset']}(\n " + '\n '.join(colored_lines) + "\n )" + + +def debug_print(debug: bool, intro: str, *args: Any) -> None: + """ + Print debug messages if debug mode is enabled with color-coded components. """ if not debug: return + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - message = " ".join(map(str, args)) - print(f"\033[97m[\033[90m{timestamp}\033[97m]\033[90m {message}\033[0m") + header = f"{COLORS['bracket']}[{COLORS['timestamp']}{ + timestamp}{COLORS['bracket']}]{COLORS['reset']}" + + # Special handling for tool call processing messages + if "Processing tool call" in intro: + if len(args) >= 2: + tool_name, _, tool_args = args + message = ( + f"{header} { + COLORS['intro']}Processing tool call:{ + COLORS['reset']} " + f"{COLORS['tool']}{tool_name}{COLORS['reset']} " + f"{COLORS['intro']}with arguments{COLORS['reset']} " + f"{format_value(tool_args)}" + ) + else: + message = f"{header} {COLORS['intro']}{intro}{COLORS['reset']}" + else: + formatted_intro = f"{COLORS['intro']}{intro}{COLORS['reset']}" + formatted_args = [] + for arg in args: + if isinstance(arg, str) and arg.startswith( + ('get_', 'list_', 'process_', 'handle_')): + formatted_args.append(f"{COLORS['function']}{ + arg}{COLORS['reset']}") + elif hasattr(arg, '__class__'): + formatted_args.append(format_value(arg)) + else: + formatted_args.append(format_value(arg)) + + message = f"{header} {formatted_intro} { + ' '.join(map(str, formatted_args))}" + + print(message) def merge_fields(target, source): diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py new file mode 100644 index 00000000..f08c9908 --- /dev/null +++ b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py @@ -0,0 +1,47 @@ +""" +This is a short example of how we can solve ctf scenarios programmatically. With CAI framework +""" +# 1. Instance CAI, pentestperf and tools +from cai import CAI, Agent +from cai.tools.cli import list_dir, cat_file +import pentestperf as ptt +from wasabi import color +import os + +# 2. Init the CTF environment +ctf = ptt.ctf( + "picoctf_static_flag", + subnet="192.168.9.0/24", + container_name="ancient", + ip_address="192.168.9.52" +) +ctf.start_ctf() +challenges = list(ctf.get_challenges().keys()) +challenge = challenges[0] + +# 3. Instance the cybersecurity AI Framework client +client = CAI(ctf=ctf) + +# 4. Instance the agent/agents to solve the CTF +ctf_agent = Agent( + model="qwen2.5:14b", + name="A file listing tool", + instructions="You are an agent tailored to list and read files within a file system.", + functions=[list_dir, cat_file], # This code can execute by LLM + # parallel_tool_calls=True, Enable multi function calling +) + + +messages = [{ + "role": "user", + "content": "Instructions: " + ctf.get_instructions() + + "\nChallenge: " + ctf.get_challenges()[challenge] + + "\nTechniques: " + ctf.get_techniques() + + "\nExtract and return only the flag" +}] + +# 5. Run the swarm +response = client.run(agent=ctf_agent, messages=messages, debug=True) +print(response.messages[-1]["content"]) + +ctf.stop_ctf() diff --git a/examples/cybersecurity/1_picoctf_static_flag.py b/examples/cybersecurity/1_picoctf_static_flag.py index aa889355..d4cc6a83 100644 --- a/examples/cybersecurity/1_picoctf_static_flag.py +++ b/examples/cybersecurity/1_picoctf_static_flag.py @@ -1,42 +1,113 @@ from cai import CAI, Agent import pentestperf as ptt +from wasabi import color client = CAI() -english_agent = Agent( + +def run_ctf(ctf, command, stdout=True): + try: + # Ensure the command is executed in a shell that supports command + # chaining + output = ctf.get_shell(command) + # exploit_logger.log_ok() + + if stdout: + print("\033[32m" + output + "\033[0m") + return output + except Exception as e: + print(color(f"Error executing CTF command: {e}", fg="red")) + # exploit_logger.log_error(str(e)) + return f"Error executing CTF command: {str(e)}" + + +def list_dir(path: str, args: str = "") -> str: + """ + List the contents of a directory. + + Args: + path: The directory path to list contents from + args: Additional arguments to pass to the ls command + + Returns: + str: The output of running the ls command + """ + global ctf + command = f'ls {path} {args}' + return run_ctf(ctf, command) + + +def cat_file(args: str, file_path: str) -> str: + """ + Display the contents of a file. + + Args: + args: Additional arguments to pass to the cat command + file_path: Path to the file to display contents of + + Returns: + str: The output of running the cat command + """ + global ctf + command = f'cat {args} {file_path} ' + return run_ctf(ctf, command) + +# def flag_checker(flag: str) -> str: +# """ +# This tool is used to check if the guessed flag is valid. +# """ +# if ctf.get_flag(ctf.get_challenges()[challenge]) == flag: +# return "Flag Found: " + flag +# else: +# return "Flag Not Found: " + flag + + +listing_agent = Agent( model="qwen2.5:14b", name="A file listing tool", - instructions="You only speak English. If another language is detected, invoke transfer_to_spanish_agent.", - # instructions="You only speak English.", - # tool_choice="required", # not working with ollama and qwen2.5 + instructions="You are an agent tailored to list and read files within a file system.", + # parallel_tool_calls=True, ) -spanish_agent = Agent( - name="Spanish Agent", - instructions="You only speak Spanish.", +# flag_extractor = Agent( +# model="qwen2.5:14b", +# name="A flag extractor", +# instructions="You are an agent tailored to extract only the flag from a given content.", +# ) + +# def transfer_to_flag_extractor_agent(): +# """Transfer to flag extractor.""" +# return flag_extractor + + +####################################### + +# PentestPerf setup for "picoctf_static_flag" +ctf = ptt.ctf( + "picoctf_static_flag", + subnet="192.168.9.0/24", + container_name="ancient", + ip_address="192.168.9.52", ) +ctf.start_ctf() +flag = ctf.get_flag() -def transfer_to_spanish_agent(): - """Transfer spanish speaking users immediately.""" - return spanish_agent +listing_agent.functions.append(list_dir) +listing_agent.functions.append(cat_file) +# listing_agent.functions.append(flag_checker) +challenges = list(ctf.get_challenges().keys()) +challenge = challenges[0] -english_agent.functions.append(transfer_to_spanish_agent) -messages = [{"role": "user", "content": "Hola, cómo estás?"}] - - -response = client.run(agent=english_agent, messages=messages, debug=True) +messages = [{"role": "user", "content": "Instructions: " + ctf.get_instructions() + + "\nChallenge: " + ctf.get_challenges()[challenge] + + "\nTechniques: " + ctf.get_techniques() + + "\nExtract and return only the flag"}] +response = client.run(agent=listing_agent, messages=messages, debug=True) print(response.messages[-1]["content"]) +# if ctf.check_flag(response.messages[-1]["content"], ctf.get_challenges()[challenge]): +# print(color(f"Flag found: {response.messages[-1]['content']}", fg="green")) -# # PentestPerf setup for "picoctf_static_flag" -# ctf = ptt.ctf( -# "picoctf_static_flag", -# subnet="192.168.9.0/24", -# container_name="ancient", -# ip_address="192.168.9.52", -# ) -# ctf.start_ctf() -# flag = ctf.get_flag() -# ctf.stop_ctf() +ctf.stop_ctf() diff --git a/setup.cfg b/setup.cfg index 3775a6b7..52ff3782 100644 --- a/setup.cfg +++ b/setup.cfg @@ -19,6 +19,7 @@ install_requires = tqdm pre-commit instructor + wasabi pentestperf @ git+https://gitlab:gldt-1NfZ2XRWdHNip4K45auL@gitlab.com/aliasrobotics/alias_research/pentestperf@main python_requires = >=3.10 diff --git a/tests/__init__.py b/tests/CAI/test_runs/__init__.py similarity index 100% rename from tests/__init__.py rename to tests/CAI/test_runs/__init__.py diff --git a/tests/mock_client.py b/tests/CAI/test_runs/mock_client.py similarity index 100% rename from tests/mock_client.py rename to tests/CAI/test_runs/mock_client.py diff --git a/tests/test_core.py b/tests/CAI/test_runs/test_core.py similarity index 98% rename from tests/test_core.py rename to tests/CAI/test_runs/test_core.py index b2777bdb..6c6cb4b4 100644 --- a/tests/test_core.py +++ b/tests/CAI/test_runs/test_core.py @@ -1,6 +1,6 @@ import pytest from cai import CAI, Agent -from tests.mock_client import MockOpenAIClient, create_mock_response +from .mock_client import MockOpenAIClient, create_mock_response from unittest.mock import Mock import json diff --git a/tests/test_util.py b/tests/CAI/test_runs/test_util.py similarity index 100% rename from tests/test_util.py rename to tests/CAI/test_runs/test_util.py diff --git a/tests/test_runs/test_20240402-113647.json b/tests/test_runs/test_20240402-113647.json deleted file mode 100644 index b020a09d..00000000 --- a/tests/test_runs/test_20240402-113647.json +++ /dev/null @@ -1 +0,0 @@ -[{"task_id": "02b37e8e-e436-445c-abdc-13e227616e07", "role": "user", "content": "If I have 5 ducks, and lose 2 of them. How many do I have left"}, {"task_id": "02b37e8e-e436-445c-abdc-13e227616e07", "role": "assistant", "content": "Response to user: 3 ducks"}]