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()