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