mirror of https://github.com/aliasrobotics/cai.git
Fixes on file structure
Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches
parent
f9a39bfa87
commit
12180ace7d
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue