mirror of https://github.com/aliasrobotics/cai.git
114 lines
3.1 KiB
Python
114 lines
3.1 KiB
Python
from cai import CAI, Agent
|
|
import pentestperf as ptt
|
|
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
|
|
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"
|
|
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()
|
|
|
|
|
|
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)
|
|
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()
|