From 12180ace7d638b53639deb740829eb6948dd0bc2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 14:15:22 +0000 Subject: [PATCH] Fixes on file structure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- README.md | 4 +- .../1_arch_short_picoctf_static_flag.py | 51 ------- .../cybersecurity/1_picoctf_static_flag.py | 132 +++++------------- ...g_python => 2_picoctf_reversing_python.py} | 0 4 files changed, 37 insertions(+), 150 deletions(-) delete mode 100644 examples/cybersecurity/1_arch_short_picoctf_static_flag.py rename examples/cybersecurity/{2_picoctf_reversing_python => 2_picoctf_reversing_python.py} (100%) diff --git a/README.md b/README.md index 4d0895df..cc978204 100644 --- a/README.md +++ b/README.md @@ -373,8 +373,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) - - [ ] re-organize tools + - [x] polish aesthetics (coloring with alpha past content, brief/summarized) + - [x] re-organize tools - [ ] add tests - [ ] CI - [ ] Tracing diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py deleted file mode 100644 index 1c4f1821..00000000 --- a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -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.reconnaissance.filesystem 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 CAI -response = client.run( - agent=ctf_agent, - messages=messages, - debug=True, - brief=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 d4cc6a83..1c4f1821 100644 --- a/examples/cybersecurity/1_picoctf_static_flag.py +++ b/examples/cybersecurity/1_picoctf_static_flag.py @@ -1,113 +1,51 @@ +""" +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.reconnaissance.filesystem import list_dir, cat_file import pentestperf as ptt from wasabi import color +import os -client = CAI() - - -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 are an agent tailored to list and read files within a file system.", - # parallel_tool_calls=True, -) - -# 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" +# 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", + ip_address="192.168.9.52" ) ctf.start_ctf() -flag = ctf.get_flag() - - -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] -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) +# 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 CAI +response = client.run( + agent=ctf_agent, + messages=messages, + debug=True, + brief=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")) - ctf.stop_ctf() diff --git a/examples/cybersecurity/2_picoctf_reversing_python b/examples/cybersecurity/2_picoctf_reversing_python.py similarity index 100% rename from examples/cybersecurity/2_picoctf_reversing_python rename to examples/cybersecurity/2_picoctf_reversing_python.py