From b1fc7e9ee2832f87cbcfa7df1fcaee6dadd25c9f Mon Sep 17 00:00:00 2001 From: luijait2 Date: Thu, 9 Jan 2025 15:05:54 +0000 Subject: [PATCH 01/12] Add a first example CTF Peer coded with @Mery-Sanz, @Lidia9, @vmayoral Signed-off-by: luijait2 --- cai/core.py | 6 +- .../cybersecurity/1_picoctf_static_flag.py | 121 ++++++++++++++---- setup.cfg | 1 + 3 files changed, 101 insertions(+), 27 deletions(-) diff --git a/cai/core.py b/cai/core.py index 73337962..f113a925 100644 --- a/cai/core.py +++ b/cai/core.py @@ -39,8 +39,10 @@ class CAI: Main class for the CAI library. """ - 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"): if not client: client = OpenAI(base_url=base_url, api_key=api_key) self.client = client diff --git a/examples/cybersecurity/1_picoctf_static_flag.py b/examples/cybersecurity/1_picoctf_static_flag.py index aa889355..70ab8661 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 From 77c8f36df408b5e006bd1db3082db6a790655dfe Mon Sep 17 00:00:00 2001 From: luijait2 Date: Thu, 9 Jan 2025 16:36:38 +0000 Subject: [PATCH 02/12] Add CTF_IN_DOCKER --- .env.example | 2 + README.md | 3 +- cai/core.py | 9 ++++- cai/tools/cli.py | 30 +++++++++++++++ cai/tools/common.py | 37 +++++++++++++++++++ .../cybersecurity/1_picoctf_static_flag.py | 29 ++++++--------- 6 files changed, 90 insertions(+), 20 deletions(-) create mode 100644 .env.example create mode 100644 cai/tools/cli.py create mode 100644 cai/tools/common.py diff --git a/.env.example b/.env.example new file mode 100644 index 00000000..9f656e6b --- /dev/null +++ b/.env.example @@ -0,0 +1,2 @@ +OPENAI_API_KEY= +CTF_IN_DOCKER= \ No newline at end of file diff --git a/README.md b/README.md index 34f4cb8d..49474e4e 100644 --- a/README.md +++ b/README.md @@ -355,7 +355,8 @@ 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 + - [ ] ... - [ ] CI - [ ] Tracing - [ ] Graph/flow and other abstractions diff --git a/cai/core.py b/cai/core.py index f113a925..35a4b597 100644 --- a/cai/core.py +++ b/cai/core.py @@ -31,6 +31,7 @@ from .types import ( Result, ) +from .tools.common import run_command __CTX_VARS_NAME__ = "context_variables" @@ -42,10 +43,12 @@ class CAI: def __init__(self, client=None, base_url="http://host.docker.internal:8000/v1", - api_key="alias"): + 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, @@ -147,7 +150,9 @@ class CAI: 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 + 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) diff --git a/cai/tools/cli.py b/cai/tools/cli.py new file mode 100644 index 00000000..44ad72c5 --- /dev/null +++ b/cai/tools/cli.py @@ -0,0 +1,30 @@ + +from .common import run_command + +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(args: str, file_path: 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) \ No newline at end of file diff --git a/cai/tools/common.py b/cai/tools/common.py new file mode 100644 index 00000000..e8e4a78a --- /dev/null +++ b/cai/tools/common.py @@ -0,0 +1,37 @@ +from wasabi import color +import subprocess +import os + +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 _run_attacker_machine(command, stdout=True): + try: + result = subprocess.run(command, shell=True, capture_output=True, text=True) + output = result.stdout + if stdout: + print("\033[32m" + output + "\033[0m") + return output + except Exception as e: + 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 local attacker machine based on CTF_IN_DOCKER flag + """ + if os.getenv("CTF_IN_DOCKER").lower() == "true" and ctf: + return _run_ctf(ctf, command, stdout) + else: + return _run_attacker_machine(command, stdout) diff --git a/examples/cybersecurity/1_picoctf_static_flag.py b/examples/cybersecurity/1_picoctf_static_flag.py index 70ab8661..527d98ed 100644 --- a/examples/cybersecurity/1_picoctf_static_flag.py +++ b/examples/cybersecurity/1_picoctf_static_flag.py @@ -4,11 +4,9 @@ from wasabi import color client = CAI() - def run_ctf(ctf, command, stdout=True): try: - # Ensure the command is executed in a shell that supports command - # chaining + # Ensure the command is executed in a shell that supports command chaining output = ctf.get_shell(command) # exploit_logger.log_ok() @@ -20,7 +18,6 @@ def run_ctf(ctf, command, stdout=True): # 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. @@ -36,7 +33,6 @@ def list_dir(path: str, args: str) -> str: 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. @@ -61,7 +57,6 @@ def cat_file(args: str, file_path: str) -> str: # else: # return "Flag Not Found: " + flag - listing_agent = Agent( model="qwen2.5:14b", name="A file listing tool", @@ -84,11 +79,11 @@ listing_agent = Agent( # 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", -) + "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() @@ -100,14 +95,14 @@ listing_agent.functions.append(cat_file) challenges = list(ctf.get_challenges().keys()) challenge = challenges[0] -messages = [{"role": "user", "content": "Instructions: " + ctf.get_instructions() - + "\nChallenge: " + ctf.get_challenges()[challenge] - + "\nTechniques: " + ctf.get_techniques() - + "\nExtract and return only the flag"}] +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")) +#if ctf.check_flag(response.messages[-1]["content"], ctf.get_challenges()[challenge]): +# print(color(f"Flag found: {response.messages[-1]['content']}", fg="green")) ctf.stop_ctf() From cc06d9d5b202563ac98f988a1162fbb239958f52 Mon Sep 17 00:00:00 2001 From: luijait2 Date: Thu, 9 Jan 2025 16:59:43 +0000 Subject: [PATCH 03/12] Added CTF_IN_DOCKER variable --- cai/core.py | 7 +-- cai/tools/common.py | 17 +++++-- .../1_arch_short_picoctf_static_flag.py | 46 +++++++++++++++++++ 3 files changed, 63 insertions(+), 7 deletions(-) create mode 100644 examples/cybersecurity/1_arch_short_picoctf_static_flag.py diff --git a/cai/core.py b/cai/core.py index 35a4b597..93cbddc4 100644 --- a/cai/core.py +++ b/cai/core.py @@ -14,7 +14,7 @@ import copy import json from collections import defaultdict from typing import List - +import os # Package/library imports from openai import OpenAI # pylint: disable=import-error @@ -31,7 +31,6 @@ from .types import ( Result, ) -from .tools.common import run_command __CTX_VARS_NAME__ = "context_variables" @@ -150,9 +149,11 @@ class CAI: 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 + args[__CTX_VARS_NAME__] = context_variables if self.ctf: args["ctf"] = self.ctf + else: + os.environ["CTF_IN_DOCKER"] = "false" raw_result = function_map[name](**args) result: Result = self.handle_function_result(raw_result, debug) diff --git a/cai/tools/common.py b/cai/tools/common.py index e8e4a78a..d8baff8f 100644 --- a/cai/tools/common.py +++ b/cai/tools/common.py @@ -2,9 +2,11 @@ from wasabi import color import subprocess import os + def _run_ctf(ctf, command, stdout=True): try: - # Ensure the command is executed in a shell that supports command chaining + # Ensure the command is executed in a shell that supports command + # chaining output = ctf.get_shell(command) # exploit_logger.log_ok() @@ -16,9 +18,14 @@ def _run_ctf(ctf, command, stdout=True): # exploit_logger.log_error(str(e)) return f"Error executing CTF command: {str(e)}" + def _run_attacker_machine(command, stdout=True): try: - result = subprocess.run(command, shell=True, capture_output=True, text=True) + result = subprocess.run( + command, + shell=True, + capture_output=True, + text=True) output = result.stdout if stdout: print("\033[32m" + output + "\033[0m") @@ -27,11 +34,13 @@ def _run_attacker_machine(command, stdout=True): 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 local attacker machine based on CTF_IN_DOCKER flag + Run command either in CTF container or + on local attacker machine based on CTF_IN_DOCKER flag """ - if os.getenv("CTF_IN_DOCKER").lower() == "true" and ctf: + if os.getenv("CTF_IN_DOCKER", "false").lower() == "true" and ctf: return _run_ctf(ctf, command, stdout) else: return _run_attacker_machine(command, stdout) 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..241de965 --- /dev/null +++ b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py @@ -0,0 +1,46 @@ +""" +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 + +# 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() From 4bce6cfc067b4d1dc8c98d743fe7844ea157ef6e Mon Sep 17 00:00:00 2001 From: luijait2 Date: Thu, 9 Jan 2025 17:14:06 +0000 Subject: [PATCH 04/12] Push --- .gitignore | 1 + README.md | 4 ++++ cai/tools/cli.py | 8 ++++++- cai/tools/common.py | 24 ++++++++++++------- .../1_arch_short_picoctf_static_flag.py | 2 ++ tests/{ => CAI/test_runs}/__init__.py | 0 tests/{ => CAI/test_runs}/mock_client.py | 0 tests/{ => CAI/test_runs}/test_core.py | 2 +- tests/{ => CAI/test_runs}/test_util.py | 0 tests/test_runs/test_20240402-113647.json | 1 - 10 files changed, 31 insertions(+), 11 deletions(-) rename tests/{ => CAI/test_runs}/__init__.py (100%) rename tests/{ => CAI/test_runs}/mock_client.py (100%) rename tests/{ => CAI/test_runs}/test_core.py (98%) rename tests/{ => CAI/test_runs}/test_util.py (100%) delete mode 100644 tests/test_runs/test_20240402-113647.json 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 49474e4e..fb7e1525 100644 --- a/README.md +++ b/README.md @@ -361,3 +361,7 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa - [ ] Tracing - [ ] Graph/flow and other abstractions - [ ] Plan/router + +### ENV Vars + +CTF_IN_DOCKER="True" diff --git a/cai/tools/cli.py b/cai/tools/cli.py index 44ad72c5..38deaa2c 100644 --- a/cai/tools/cli.py +++ b/cai/tools/cli.py @@ -1,6 +1,11 @@ +""" +Here are the CLI tools for executing commands, they should usually return +with run_command, but you can be creative +""" from .common import run_command + def list_dir(path: str, args: str, ctf=None) -> str: """ List the contents of a directory. @@ -15,6 +20,7 @@ def list_dir(path: str, args: str, ctf=None) -> str: command = f'ls {path} {args}' return run_command(command, ctf=ctf) + def cat_file(args: str, file_path: str, ctf=None) -> str: """ Display the contents of a file. @@ -27,4 +33,4 @@ def cat_file(args: str, file_path: str, ctf=None) -> str: str: The output of running the cat command """ command = f'cat {args} {file_path} ' - return run_command(command, ctf=ctf) \ No newline at end of file + return run_command(command, ctf=ctf) diff --git a/cai/tools/common.py b/cai/tools/common.py index d8baff8f..5a5f01e9 100644 --- a/cai/tools/common.py +++ b/cai/tools/common.py @@ -1,6 +1,13 @@ -from wasabi import color -import subprocess +"""Basic utilities for executing commands inside and outside virtual +containers. It's important to note that for proper functioning, the +CTF_IN_DOCKER environment variable must be configured, and commands can +only be executed through CTF instances with pentestperf within the code. +See examples/cybersecurity/1_arch_short_picoctf_static_flag.py for an +example.""" + import os +import subprocess # nosec B404 +from wasabi import color # pylint: disable=import-error def _run_ctf(ctf, command, stdout=True): @@ -13,7 +20,7 @@ def _run_ctf(ctf, command, stdout=True): if stdout: print("\033[32m" + output + "\033[0m") return output - except Exception as e: + 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)}" @@ -21,16 +28,18 @@ def _run_ctf(ctf, command, stdout=True): def _run_attacker_machine(command, stdout=True): try: + # nosec B602 - shell=True is required for command chaining result = subprocess.run( command, - shell=True, + shell=True, # nosec B602 capture_output=True, - text=True) + text=True, + check=True) output = result.stdout if stdout: print("\033[32m" + output + "\033[0m") return output - except Exception as e: + 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)}" @@ -42,5 +51,4 @@ def run_command(command, ctf=None, stdout=True): """ if os.getenv("CTF_IN_DOCKER", "false").lower() == "true" and ctf: return _run_ctf(ctf, command, stdout) - else: - return _run_attacker_machine(command, stdout) + return _run_attacker_machine(command, stdout) diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py index 241de965..6bc6d32e 100644 --- a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py +++ b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py @@ -6,8 +6,10 @@ 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 +os.environ["CTF_IN_DOCKER"] = "true" ctf = ptt.ctf( "picoctf_static_flag", subnet="192.168.9.0/24", 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"}] From f50638388737f05f775f37bd5c82693017d0c1af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 08:01:06 +0000 Subject: [PATCH 05/12] Address bug which led to an error when allucinating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix implies ensuring that even with allucationations, args adopts a value Signed-off-by: Víctor Mayoral Vilches --- .../cybersecurity/1_picoctf_static_flag.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/examples/cybersecurity/1_picoctf_static_flag.py b/examples/cybersecurity/1_picoctf_static_flag.py index 527d98ed..d4cc6a83 100644 --- a/examples/cybersecurity/1_picoctf_static_flag.py +++ b/examples/cybersecurity/1_picoctf_static_flag.py @@ -4,9 +4,11 @@ from wasabi import color client = CAI() + def run_ctf(ctf, command, stdout=True): try: - # Ensure the command is executed in a shell that supports command chaining + # Ensure the command is executed in a shell that supports command + # chaining output = ctf.get_shell(command) # exploit_logger.log_ok() @@ -18,7 +20,8 @@ def run_ctf(ctf, command, stdout=True): # exploit_logger.log_error(str(e)) return f"Error executing CTF command: {str(e)}" -def list_dir(path: str, args: str) -> str: + +def list_dir(path: str, args: str = "") -> str: """ List the contents of a directory. @@ -33,6 +36,7 @@ def list_dir(path: str, args: str) -> str: 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. @@ -57,6 +61,7 @@ def cat_file(args: str, file_path: str) -> str: # else: # return "Flag Not Found: " + flag + listing_agent = Agent( model="qwen2.5:14b", name="A file listing tool", @@ -79,11 +84,11 @@ listing_agent = Agent( # 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", - ) + "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() @@ -95,14 +100,14 @@ listing_agent.functions.append(cat_file) challenges = list(ctf.get_challenges().keys()) challenge = challenges[0] -messages = [{"role": "user", "content": "Instructions: " + ctf.get_instructions() - + "\nChallenge: " + ctf.get_challenges()[challenge] - + "\nTechniques: " + ctf.get_techniques() - + "\nExtract and return only the flag"}] +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")) +# if ctf.check_flag(response.messages[-1]["content"], ctf.get_challenges()[challenge]): +# print(color(f"Flag found: {response.messages[-1]['content']}", fg="green")) ctf.stop_ctf() From 4d4e4e3790534192073591b3801261a62dfef34f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 08:03:11 +0000 Subject: [PATCH 06/12] Various improvements in docs and debug prints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/core.py | 51 +++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/cai/core.py b/cai/core.py index 93cbddc4..50f7c1c6 100644 --- a/cai/core.py +++ b/cai/core.py @@ -36,7 +36,7 @@ __CTX_VARS_NAME__ = "context_variables" class CAI: """ - Main class for the CAI library. + Cybersecurity AI (CAI) object """ def __init__(self, @@ -95,7 +95,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: @@ -122,7 +128,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( @@ -144,7 +180,11 @@ 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 @@ -183,6 +223,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) From 4888ce375c473f5763601b2dd34a44645e4ea286 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 08:06:12 +0000 Subject: [PATCH 07/12] Address bugs due to allucinations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/tools/cli.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cai/tools/cli.py b/cai/tools/cli.py index 38deaa2c..d2a71f89 100644 --- a/cai/tools/cli.py +++ b/cai/tools/cli.py @@ -3,10 +3,10 @@ Here are the CLI tools for executing commands, they should usually return with run_command, but you can be creative """ -from .common import run_command +from cai.tools.common import run_command # pylint: disable=E0401 -def list_dir(path: str, args: str, ctf=None) -> str: +def list_dir(path: str, args: str = "", ctf=None) -> str: """ List the contents of a directory. by def . @@ -21,7 +21,7 @@ def list_dir(path: str, args: str, ctf=None) -> str: return run_command(command, ctf=ctf) -def cat_file(args: str, file_path: str, ctf=None) -> str: +def cat_file(file_path: str, args: str = "", ctf=None) -> str: """ Display the contents of a file. From 2ee152d1de0f00685ecfbe36c056619c90d31414 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 08:06:41 +0000 Subject: [PATCH 08/12] Minor docs improments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/types.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 = [] From 7151674f31e7520bf6442ff241ea7dbfc528c736 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 08:08:51 +0000 Subject: [PATCH 09/12] Improve how debug_print works, human-readable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/util.py | 161 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 157 insertions(+), 4 deletions(-) diff --git a/cai/util.py b/cai/util.py index 6e956cd5..3be81592 100644 --- a/cai/util.py +++ b/cai/util.py @@ -4,17 +4,170 @@ 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. + """ + parts = [] + parts.append( + f"\n { + COLORS['object']}ChatCompletionMessage{ + COLORS['reset']}(") + + # Format standard attributes + if hasattr(msg, 'content'): + parts.append( + f"\n { + COLORS['arg_key']}content{ + COLORS['reset']}: { + COLORS['arg_value']}{ + msg.content}{ + COLORS['reset']}") + if hasattr(msg, 'role'): + parts.append( + f"\n { + COLORS['arg_key']}role{ + COLORS['reset']}: { + COLORS['arg_value']}{ + msg.role}{ + COLORS['reset']}") + + # Format tool calls if present + if hasattr(msg, 'tool_calls') and msg.tool_calls: + tool_calls_str = [] + for tool in msg.tool_calls: + tool_str = ( + f"\n {COLORS['object']}ToolCall{COLORS['reset']}(" + f"\n { + COLORS['arg_key']}id{ + COLORS['reset']}: { + COLORS['arg_value']}{ + tool.id}{ + COLORS['reset']}" + f"\n { + COLORS['arg_key']}name{ + COLORS['reset']}: { + COLORS['arg_value']}{ + tool.function.name}{ + COLORS['reset']}" + f"\n { + COLORS['arg_key']}arguments{ + COLORS['reset']}: { + format_value( + json.loads( + tool.function.arguments))}" + f"\n )" + ) + tool_calls_str.append(tool_str) + parts.append( + f"\n { + COLORS['arg_key']}tool_calls{ + COLORS['reset']}: [{ + ','.join(tool_calls_str)}\n ]") + + parts.append("\n )") + return "".join(parts) + + +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): From 0abbfa1fd8d8448317d5db985e8df4ce2267169c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 08:19:19 +0000 Subject: [PATCH 10/12] Further improvements on debug_print MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/util.py | 84 ++++++++++++++++++----------------------------------- 1 file changed, 28 insertions(+), 56 deletions(-) diff --git a/cai/util.py b/cai/util.py index 3be81592..b5d362f8 100644 --- a/cai/util.py +++ b/cai/util.py @@ -65,65 +65,37 @@ def format_chat_completion(msg) -> str: """ Format a ChatCompletionMessage object with proper indentation and colors. """ - parts = [] - parts.append( - f"\n { - COLORS['object']}ChatCompletionMessage{ - COLORS['reset']}(") + # Convert message to dict and handle OpenAI types + try: + msg_dict = json.loads(msg.model_dump_json()) + except AttributeError: + msg_dict = msg.__dict__ - # Format standard attributes - if hasattr(msg, 'content'): - parts.append( - f"\n { - COLORS['arg_key']}content{ - COLORS['reset']}: { - COLORS['arg_value']}{ - msg.content}{ - COLORS['reset']}") - if hasattr(msg, 'role'): - parts.append( - f"\n { - COLORS['arg_key']}role{ - COLORS['reset']}: { - COLORS['arg_value']}{ - msg.role}{ - COLORS['reset']}") + # Clean up the dictionary + msg_dict = {k: v for k, v in msg_dict.items() if v is not None} - # Format tool calls if present - if hasattr(msg, 'tool_calls') and msg.tool_calls: - tool_calls_str = [] - for tool in msg.tool_calls: - tool_str = ( - f"\n {COLORS['object']}ToolCall{COLORS['reset']}(" - f"\n { - COLORS['arg_key']}id{ - COLORS['reset']}: { - COLORS['arg_value']}{ - tool.id}{ - COLORS['reset']}" - f"\n { - COLORS['arg_key']}name{ - COLORS['reset']}: { - COLORS['arg_value']}{ - tool.function.name}{ - COLORS['reset']}" - f"\n { - COLORS['arg_key']}arguments{ - COLORS['reset']}: { - format_value( - json.loads( - tool.function.arguments))}" - f"\n )" - ) - tool_calls_str.append(tool_str) - parts.append( - f"\n { - COLORS['arg_key']}tool_calls{ - COLORS['reset']}: [{ - ','.join(tool_calls_str)}\n ]") + # Format with json.dumps for consistent indentation + formatted_json = json.dumps(msg_dict, indent=2) - parts.append("\n )") - return "".join(parts) + # 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: From ab68a2f0ef42134578432e6a62bb92cdf378904c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 09:22:31 +0000 Subject: [PATCH 11/12] Remove CTF_IN_DOCKER env. variable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- .env.example | 1 - README.md | 4 ---- cai/core.py | 3 --- cai/tools/common.py | 15 ++++++--------- .../1_arch_short_picoctf_static_flag.py | 1 - 5 files changed, 6 insertions(+), 18 deletions(-) diff --git a/.env.example b/.env.example index 9f656e6b..e570b8b5 100644 --- a/.env.example +++ b/.env.example @@ -1,2 +1 @@ OPENAI_API_KEY= -CTF_IN_DOCKER= \ No newline at end of file diff --git a/README.md b/README.md index fb7e1525..49474e4e 100644 --- a/README.md +++ b/README.md @@ -361,7 +361,3 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa - [ ] Tracing - [ ] Graph/flow and other abstractions - [ ] Plan/router - -### ENV Vars - -CTF_IN_DOCKER="True" diff --git a/cai/core.py b/cai/core.py index 50f7c1c6..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 -import os # Package/library imports from openai import OpenAI # pylint: disable=import-error @@ -192,8 +191,6 @@ class CAI: args[__CTX_VARS_NAME__] = context_variables if self.ctf: args["ctf"] = self.ctf - else: - os.environ["CTF_IN_DOCKER"] = "false" raw_result = function_map[name](**args) result: Result = self.handle_function_result(raw_result, debug) diff --git a/cai/tools/common.py b/cai/tools/common.py index 5a5f01e9..3885cf28 100644 --- a/cai/tools/common.py +++ b/cai/tools/common.py @@ -1,11 +1,8 @@ -"""Basic utilities for executing commands inside and outside virtual -containers. It's important to note that for proper functioning, the -CTF_IN_DOCKER environment variable must be configured, and commands can -only be executed through CTF instances with pentestperf within the code. -See examples/cybersecurity/1_arch_short_picoctf_static_flag.py for an -example.""" +""" +Basic utilities for executing commands inside or outside of virtual +containers. +""" -import os import subprocess # nosec B404 from wasabi import color # pylint: disable=import-error @@ -47,8 +44,8 @@ def _run_attacker_machine(command, stdout=True): def run_command(command, ctf=None, stdout=True): """ Run command either in CTF container or - on local attacker machine based on CTF_IN_DOCKER flag + on the local attacker machine """ - if os.getenv("CTF_IN_DOCKER", "false").lower() == "true" and ctf: + if ctf: return _run_ctf(ctf, command, stdout) return _run_attacker_machine(command, stdout) diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py index 6bc6d32e..f08c9908 100644 --- a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py +++ b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py @@ -9,7 +9,6 @@ from wasabi import color import os # 2. Init the CTF environment -os.environ["CTF_IN_DOCKER"] = "true" ctf = ptt.ctf( "picoctf_static_flag", subnet="192.168.9.0/24", From 63544e24164f20cb019e929a1dda5b9f92549853 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 09:35:39 +0000 Subject: [PATCH 12/12] Update in README goals MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 49474e4e..2ca8a967 100644 --- a/README.md +++ b/README.md @@ -356,7 +356,8 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa - [x] Dev container - [x] pre-commit hooks - [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