RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'issue_62' into 'v0.4.0' 

Review agents

See merge request aliasrobotics/alias_research/cai!142
This commit is contained in:
Mery-Sanz 2025-04-24 08:31:47 +00:00
parent 4305f41089 fedef3cf51
commit d87dd27147
30 changed files with 257 additions and 231 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
ci/test/.test.yml
View File

@ -110,6 +110,12 @@
variables:
TEST_PATH: tests/agents/test_max_turns.py
# 🤖 agents test_agent_inference:
# <<: *run_test
# variables:
# TEST_PATH: tests/agents/test_agent_inference.py
# ⚙️ core test_openai_chatcompletions:
# <<: *run_test
# variables:

3
examples/cai/simple_one_tool_test.py
View File

@ -11,7 +11,7 @@ import asyncio
import json
from dotenv import load_dotenv
from openai import AsyncOpenAI
from cai.sdk.agents import Runner, set_default_openai_client
from cai.sdk.agents import Runner, set_default_openai_client, set_tracing_disabled
from cai.agents import get_agent_by_name
from cai.util import fix_litellm_transcription_annotations, color, cli_print_agent_messages
from cai.sdk.agents.models._openai_shared import set_use_responses_by_default
@ -19,6 +19,7 @@ from cai.sdk.agents.models._openai_shared import set_use_responses_by_default
# Load environment variables
load_dotenv()
#set_tracing_disabled(True) #disable tracing or OPENAI_AGENTS_DISABLE_TRACING=1
# NOTE: This is needed when using LiteLLM Proxy Server
#

4
src/cai/agents/__init__.py
View File

@ -123,8 +123,8 @@ def get_available_agents() -> Dict[str, Agent]: # pylint: disable=R0912 # noqa
agent_name = attr_name
if agent_name not in agents_to_display:
agents_to_display[agent_name] = attr
except (ImportError, AttributeError):
pass
except (ImportError, AttributeError) as e:
print(f"Error importing {agent_name}: {e}")
return agents_to_display

146
src/cai/agents/basic.py
View File

@ -1,146 +0,0 @@
"""
Collection of basic agents
"""
import os
from cai.sdk.agents import Agent # pylint: disable=import-error
from cai.tools.reconnaissance.netcat import netcat # pylint: disable=import-error # noqa: E501
from cai.tools.reconnaissance.nmap import nmap # pylint: disable=import-error # noqa: E501
# First party imports
from cai.tools.reconnaissance.filesystem import (
list_dir,
cat_file,
find_file,
pwd_command
)
from cai.tools.reconnaissance.generic_linux_command import (
generic_linux_command
)
from cai.tools.reconnaissance.crypto_tools import (
decode64,
strings_command,
decode_hex_bytes
)
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
ctf_agent = Agent(
model=model,
name="CTF agent",
description="""Agent focused on CTF challenges.
Expert in cybersecurity and exploitation.
Specializes in executing Linux commands efficiently.""",
instructions="""You are a Cybersecurity expert Leader facing a CTF
challenge.
INSTRUCTIONS:
1. Execute the generic_linux_command tool without any
explanation.
2. Be efficient and strategic when executing commands.
3. Never assume the flag format - it could be any string
4. NEVER GIVE UP. KEEP EXECUTING TOOLS UNTIL YOU FIND
THE FLAG.
5. IF YOU ENCOUNTER A PROBLEM, IT IS HIGHLY LIKELY
THAT YOU ARE ASSUMING WRONG FILENAMES, PORTS,
IP ADDRESSES, ETC. PAY ATTENTION TO YOUR HALLUCINATIONS.
6. YOU MUST USE THE flag_discriminator to check the flag
when you have it in DECODED 64-bit format
7. Do a initial recon (nmap, service enumeration)""",
tools=[
generic_linux_command,
],
)
def transfer_to_ctf_agent(**kwargs): # pylint: disable=W0613
"""Transfer to ctf agent.
Accepts any keyword arguments but ignores them."""
return ctf_agent
network_agent = Agent(
model=model,
name="Network Specialist",
description="""Agent specialized in network reconnaissance.
Expert in using nmap and netcat for service discovery.
Focuses on identifying network vulnerabilities.""",
instructions=("""INSTRUCTIONS:
1. You specialize in using nmap and netcat.
INSTRUCTIONS:
1. You always use real information of the
CTF challenge as arguments for the tools.
If you lack of that information, you find another
way to get it.
2. Keep calling agents until you find the flag.
3. Do not be verbose, just call the tools.
4. YOU NEVER USE PLACEHOLDERS, AS THE TOOLS WILL
NOT WORK WITH THEM.
5. Whenever you find a likely candidate flag, call
the flag_discriminator agent.
6. DO NOT SEND THE SAME COMMAND OVER AND OVER"""),
tools=[netcat, nmap],
)
def transfer_to_network_agent(**kwargs): # pylint: disable=W0613
""" Transfer to network agent.
Accepts any keyword arguments but ignores them."""
return network_agent
crypto_agent = Agent(
model=model,
name="Cryptography Expert",
description="""Agent specialized in cryptographic analysis.
Expert in decoding and analyzing encrypted data.
Focuses on breaking cryptographic challenges in CTFs.""",
instructions=("""
INSTRUCTIONS:
1. Your expertise is in decryption and
data analysis.
2. Use your tools to decode and analyze
encrypted data.
3. DO NOT BE VERBOSE, JUST CALL THE TOOLS AS NEEDED
4. Use your tools to solve the CTF challenge.
5. Do not generate a plan or verbose output.
"""
),
tools=[decode64, strings_command, decode_hex_bytes],
)
def transfer_to_crypto_agent(**kwargs): # pylint: disable=W0613
""" Transfer to crypto agent.
Accepts any keyword arguments but ignores them."""
return crypto_agent
listing_agent = Agent(
model=model,
name="File Explorer",
description="""Agent specialized in filesystem navigation and exploration.
Expert in finding and analyzing files and directories.
Focuses on discovering hidden files and sensitive data.""",
instructions=("""INSTRUCTIONS:
1. You specialize in navigating file systems.
2. Use your tools to list, read, and find files in
a directory.
3. DO NOT USE THE SAME TOOL FOR THE SAME FILE OR
DIRECTORY more than 3 times.
4. DO NOT BE VERBOSE, JUST CALL THE TOOLS AS NEEDED
TO SOLVE THE CTF
5. KEEP CALLING THE TOOLS OR THE CTF Leader AGENT
UNTIL YOU FIND THE FLAG"""),
tools=[
list_dir,
cat_file,
find_file,
pwd_command,
generic_linux_command],
)
def transfer_to_listing_agent(**kwargs): # pylint: disable=W0613
""" Transfer to listing agent.
Accepts any keyword arguments but ignores them."""
return listing_agent

23
src/cai/agents/blue_teamer.py
View File

@ -4,8 +4,10 @@ SSH_HOST
SSH_USER
"""
import os
from cai.sdk.agents import Agent # pylint: disable=import-error
from openai import AsyncOpenAI
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error
from cai.util import load_prompt_template
from dotenv import load_dotenv
from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501
run_ssh_command_with_credentials
)
@ -14,7 +16,7 @@ from cai.tools.reconnaissance.generic_linux_command import ( # pylint: disable=
generic_linux_command
)
from cai.tools.web.search_web import ( # pylint: disable=import-error # noqa: E501
search_web
make_web_search_with_explanation
)
from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error # noqa: E501
@ -23,23 +25,26 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error
# Prompts
blueteam_agent_system_prompt = load_prompt_template("prompts/system_blue_team_agent.md")
# Define functions list based on available API keys
functions = [
# Define tools list based on available API keys
tools = [
generic_linux_command,
run_ssh_command_with_credentials,
execute_code,
]
# Add search_web function if PERPLEXITY_API_KEY environment variable is set
load_dotenv()
# Add search_web tools if PERPLEXITY_API_KEY environment variable is set
if os.getenv('PERPLEXITY_API_KEY'):
functions.append(search_web)
tools.append(make_web_search_with_explanation)
blueteam_agent = Agent(
name="Blue Team Agent",
instructions=blueteam_agent_system_prompt,
description="""Agent that specializes in system defense and security monitoring.
Expert in cybersecurity protection and incident response.""",
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
functions=functions,
parallel_tool_calls=False,
model=OpenAIChatCompletionsModel(
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
openai_client=AsyncOpenAI(),
),
tools=tools,
)

17
src/cai/agents/bug_bounter.py
View File

@ -1,7 +1,8 @@
"""Red Team Base Agent"""
import os
from dotenv import load_dotenv
from cai.sdk.agents import Agent # pylint: disable=import-error
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
from openai import AsyncOpenAI
from cai.util import load_prompt_template
from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501
run_ssh_command_with_credentials
@ -26,8 +27,8 @@ from cai.tools.reconnaissance.shodan import ( # pylint: disable=import-error #
load_dotenv()
# Prompts
bug_bounter_system_prompt = load_prompt_template("prompts/system_bug_bounter.md")
# Define functions list based on available API keys
functions = [
# Define tools list based on available API keys
tools = [
generic_linux_command,
execute_code,
shodan_search,
@ -35,13 +36,17 @@ functions = [
]
if os.getenv('GOOGLE_SEARCH_API_KEY') and os.getenv('GOOGLE_SEARCH_CX'):
functions.append(make_google_search)
tools.append(make_google_search)
bug_bounter_agent = Agent(
name="Bug Bounter",
instructions=bug_bounter_system_prompt,
description="""Agent that specializes in bug bounty hunting and vulnerability discovery.
Expert in web security, API testing, and responsible disclosure.""",
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
tools=functions,
tools=tools,
model=OpenAIChatCompletionsModel(
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
openai_client=AsyncOpenAI(),
)
)

9
src/cai/agents/codeagent.py
View File

@ -31,8 +31,9 @@ from cai.agents.meta.local_python_executor import (
fix_final_answer_code,
truncate_content,
)
from cai.sdk.agents import Agent, Result
from cai.sdk.agents import Agent, Result, OpenAIChatCompletionsModel
from dotenv import load_dotenv
from openai import AsyncOpenAI
class CodeAgentException(Exception):
"""Base exception class for CodeAgent-related errors."""
@ -185,7 +186,7 @@ class CodeAgent(Agent):
name: str = "CodeAgent",
model: str = "qwen2.5:14b",
instructions: Union[str, Callable[[], str]] = None,
functions: List[Callable] = None,
tools: List[Callable] = None,
additional_authorized_imports: Optional[List[str]] = None,
description: str = """Agent focused on writing and executing code.
State-of-the-art in code production.""",
@ -201,7 +202,7 @@ class CodeAgent(Agent):
name: Name of the agent
model: Model to use for the agent
instructions: Instructions for the agent
functions: List of functions available to the agent
tools: List of tools available to the agent
additional_authorized_imports: List of additional imports to allow
max_print_outputs_length: Maximum length of print outputs
reasoning_effort: Level of reasoning effort (low, medium, high)

21
src/cai/agents/flag_discriminator.py
View File

@ -1,24 +1,27 @@
"""
CTF Flag Discriminator Agent
CTF Flag Discriminator Agent with test
"""
import os
from cai.sdk.agents import Agent
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
from openai import AsyncOpenAI
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
flag_discriminator = Agent(
model="gpt-4o" if os.getenv('CAI_MODEL') == "o3-mini" else model,
name="Flag discriminator",
description="Agent focused on extracting the flag from the output",
instructions="""INSTRUCTIONS:
1. You are an agent tailored to extract the flag
from the output.
1. You are an agent tailored to extract the flag from a given output.
2. ONLY return the flag, nothing else.
3. IF THE FLAG IS NOT FOUND, CALL THE ctf_agent
again to CONTINUE the search""",
3. Flags can be in any format (not always CTF{...}).
4. If you do not find a flag, call `ctf_agent` to continue investigating.
""",
model=OpenAIChatCompletionsModel(
model="gpt-4o" if os.getenv('CAI_MODEL') == "o3-mini" else model,
openai_client=AsyncOpenAI(),
)
)
# TRANSFER FUNCTIONS
# Transfer Function
def transfer_to_flag_discriminator(**kwargs): # pylint: disable=W0613
"""Transfer flag discriminator.
Accepts any keyword arguments but ignores them."""

15
src/cai/agents/mail.py
View File

@ -2,9 +2,13 @@
Mail Agent module for checking email configuration security.
"""
import os
from openai import AsyncOpenAI
import dns.resolver # pylint: disable=import-error
from cai.sdk.agents import Agent
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
from cai.tools.misc.cli_utils import execute_cli_command
from cai.sdk.agents import function_tool
def get_txt_record(domain, record_type='TXT'):
@ -58,7 +62,7 @@ def check_dkim(domain: str, selector: str = "default"):
return txt_records[0]
return None
@function_tool
def check_mail_spoofing_vulnerability(
domain: str,
dkim_selector: str = "default") -> dict:
@ -100,7 +104,6 @@ def check_mail_spoofing_vulnerability(
dns_smtp_agent = Agent(
model="gpt-4o",
name="DNS_SMTP_Agent",
description="Agent focused on assessing spoofing DMARC.",
instructions=(
@ -110,5 +113,9 @@ dns_smtp_agent = Agent(
"detailed reports. Use execute_cli_command for basic scans. "
"USE ONLY TOOL CALLS, DONT RETURN REASON."
),
tools=[check_mail_spoofing_vulnerability, execute_cli_command]
tools=[check_mail_spoofing_vulnerability, execute_cli_command],
model=OpenAIChatCompletionsModel(
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
openai_client=AsyncOpenAI(),
)
)

23
src/cai/agents/memory.py
View File

@ -82,7 +82,7 @@ Environment Variables enabling the episodic memory store
"""
import os
from cai.sdk.agents import Agent
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
from cai.tools.misc.rag import add_to_memory_semantic, add_to_memory_episodic
from cai.rag.vector_db import get_previous_memory
@ -145,32 +145,35 @@ QUERY_PROMPT = """INSTRUCTIONS:
"""
semantic_builder = Agent(
model=model,
name="Semantic_Builder",
instructions=ADD_MEMORY_PROMPT,
description="""Agent that stores semantic memories from security assessments
and CTF exercises in semantic format.""",
tool_choice="required",
temperature=0,
functions=[add_to_memory_semantic],
parallel_tool_calls=True
tools=[add_to_memory_semantic],
model=OpenAIChatCompletionsModel(
model=model_name,
openai_client=AsyncOpenAI(),
)
)
episodic_builder = Agent(
model=model,
name="Episodic_Builder",
instructions=ADD_MEMORY_PROMPT,
description="""Agent that stores episodic memories from security assessments
and CTF exercises in episodic format.""",
tool_choice="required",
temperature=0,
functions=[add_to_memory_episodic],
parallel_tool_calls=True
tools=[add_to_memory_episodic],
model=OpenAIChatCompletionsModel(
model=model_name,
openai_client=AsyncOpenAI(),
)
)
query_agent = Agent(
model=model,
name="Query_Agent",
description="""Agent that queries the memory system to retrieve relevant
historical information from previous security assessments
@ -178,4 +181,8 @@ query_agent = Agent(
instructions=QUERY_PROMPT,
tool_choice="required",
temperature=0,
model=OpenAIChatCompletionsModel(
model=model_name,
openai_client=AsyncOpenAI(),
)
)

8
src/cai/agents/patterns/red_team.py
View File

@ -37,10 +37,10 @@ def thought_agent_handoff(ctf=None): # pylint: disable=unused-argument
return thought_agent
# Register handoff functions to enable inter-agent communication pathways
redteam_agent.functions.append(transfer_to_dns_agent)
dns_smtp_agent.functions.append(redteam_agent_handoff)
thought_agent.functions.append(redteam_agent_handoff)
# Register handoff to enable inter-agent communication pathways
redteam_agent.handoffs.append(transfer_to_dns_agent)
dns_smtp_agent.handoffs.append(redteam_agent_handoff)
thought_agent.handoffs.append(redteam_agent_handoff)
# Initialize the swarm pattern with the thought agent as the entry point
redteam_swarm_pattern = thought_agent

38
src/cai/agents/red_teamer.py
View File

@ -1,10 +1,11 @@
"""Red Team Base Agent"""
import os
from dotenv import load_dotenv
from cai.sdk.agents import Agent # pylint: disable=import-error
from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501
run_ssh_command_with_credentials
)
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
from openai import AsyncOpenAI
# from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501
# run_ssh_command_with_credentials
# )
from cai.tools.reconnaissance.generic_linux_command import ( # pylint: disable=import-error # noqa: E501
generic_linux_command
@ -18,25 +19,36 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error
)
from cai.util import load_prompt_template
load_dotenv()
model_name = os.getenv("CAI_MODEL", "qwen2.5:14b")
# Prompts
redteam_agent_system_prompt = load_prompt_template("prompts/system_red_team_agent.md")
# Define functions list based on available API keys
functions = [
# Define tools list based on available API keys
tools = [
generic_linux_command,
run_ssh_command_with_credentials,
#run_ssh_command_with_credentials,
execute_code,
]
# Add make_web_search_with_explanation function if PERPLEXITY_API_KEY environment variable is set
if os.getenv('PERPLEXITY_API_KEY'):
functions.append(make_web_search_with_explanation)
tools.append(make_web_search_with_explanation)
redteam_agent = Agent(
name="Red Team Agent",
description="""Agent that mimics a red teamer in a security assessment.
Expert in cybersecurity, recon, and exploitation.""",
instructions=redteam_agent_system_prompt,
description="""Agent that mimic pentester/red teamer in a security assessment.
Expert in cybersecurity and exploitation.""",
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
tools=functions,
)
tools=tools,
model=OpenAIChatCompletionsModel(
model=model_name,
openai_client=AsyncOpenAI(),
),
)
# Transfer function
def transfer_to_redteam_agent(**kwargs): # pylint: disable=W0613
"""Transfer to red team agent.
Accepts any keyword arguments but ignores them."""
return redteam_agent

8
src/cai/agents/thought.py
View File

@ -6,7 +6,8 @@ using reasoner as a tool call
support meta agent may better @cai.sdk.agents.meta.reasoner_support
"""
from cai.tools.misc.reasoning import thought
from cai.sdk.agents import Agent # pylint: disable=import-error
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error
from openai import AsyncOpenAI
from cai.util import load_prompt_template
import os
@ -15,7 +16,10 @@ thought_agent_system_prompt = load_prompt_template("prompts/system_thought_route
# Thought Process Agent for analysis and planning
thought_agent = Agent(
name="ThoughAgent",
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
model=OpenAIChatCompletionsModel(
model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
openai_client=AsyncOpenAI(),
),
description="""Agent focused on analyzing and planning the next steps
in a security assessment or CTF challenge.""",
instructions=thought_agent_system_prompt,

7
src/cai/tools/command_and_control/sshpass.py
View File

@ -12,8 +12,11 @@ It handles privilege escalation very well and is autonomous regarding SSH passwo
something that hasn't been seen in other cybersecurity frameworks yet (Feb 2025)
""" # noqa: E501
from cai.tools.misc.cli_utils import execute_cli_command # pylint: disable=E0401 # noqa: E501
from cai.tools.common import run_command # pylint: disable=E0401 # noqa: E501
from cai.sdk.agents import function_tool
@function_tool
def run_ssh_command_with_credentials(
host: str,
username: str,
@ -43,4 +46,4 @@ def run_ssh_command_with_credentials(
f"{username}@{host} -p {port} "
f"'{escaped_command}'"
)
return execute_cli_command(ssh_command)
return run_command(ssh_command)

3
src/cai/tools/misc/cli_utils.py
View File

@ -3,8 +3,9 @@ CLI utilities module for executing shell commands and processing their output.
"""
from cai.tools.common import run_command # pylint: disable=E0401
from cai.sdk.agents import function_tool
@function_tool
def execute_cli_command(command: str) -> str:
"""
Execute a CLI command and return the output.

2
src/cai/tools/misc/code_interpreter.py
View File

@ -5,8 +5,10 @@ Module for executing Python code and capturing its output.
import io
import sys
from typing import Dict
from cai.sdk.agents import function_tool
@function_tool
def execute_python_code(code: str, context: Dict = None) -> str:
"""
Execute Python code and return the output.

8
src/cai/tools/misc/rag.py
View File

@ -5,12 +5,12 @@ querying and adding data to vector databases.
import os
import uuid
from cai.rag.vector_db import QdrantConnector
from cai.sdk.agents import function_tool
# CTF BASED MEMORY
collection_name = os.getenv('CAI_MEMORY_COLLECTION', "default")
@function_tool
def query_memory(query: str, top_k: int = 3, **kwargs) -> str: # pylint: disable=unused-argument,line-too-long # noqa: E501
"""
Query memory to retrieve relevant context. From Previous CTFs executions.
@ -42,7 +42,7 @@ def query_memory(query: str, top_k: int = 3, **kwargs) -> str: # pylint: disabl
except Exception: # pylint: disable=broad-exception-caught
return results
@function_tool
def add_to_memory_episodic(texts: str, step: int = 0, **kwargs) -> str: # pylint: disable=unused-argument,line-too-long # noqa: E501
"""
This is a persistent memory to add relevant context to our memory.
@ -76,7 +76,7 @@ def add_to_memory_episodic(texts: str, step: int = 0, **kwargs) -> str: # pylin
except Exception as e: # pylint: disable=broad-exception-caught
return f"Error adding documents to vector database: {str(e)}"
@function_tool
def add_to_memory_semantic(texts: str, step: int = 0, **kwargs) -> str: # pylint: disable=unused-argument,line-too-long # noqa: E501
"""
This is a persistent memory to add relevant context to our memory.

6
src/cai/tools/misc/reasoning.py
View File

@ -3,8 +3,10 @@ Reasoning tools module for tracking thoughts, findings and analysis
Provides utilities for recording and retrieving key information discovered
during CTF progression.
"""
from cai.sdk.agents import function_tool
@function_tool
def thought(breakdowns: str = "", reflection: str = "", # pylint: disable=too-many-arguments # noqa: E501
action: str = "", next_step: str = "", key_clues: str = "",
ctf=None) -> str: # pylint: disable=unused-argument # noqa: E501
@ -34,7 +36,7 @@ def thought(breakdowns: str = "", reflection: str = "", # pylint: disable=too-m
output.append(f"Key Clues: {key_clues}")
return "\n".join(output)
@function_tool
def write_key_findings(findings: str) -> str:
"""
Write key findings to a state.txt file to track important CTF details.
@ -58,7 +60,7 @@ def write_key_findings(findings: str) -> str:
except OSError as e:
return f"Error writing to state.txt: {str(e)}"
@function_tool
def read_key_findings() -> str:
"""
Read key findings from the state.txt file to retrieve important data

2
src/cai/tools/others/scripting.py
View File

@ -5,8 +5,10 @@ This is used to create and execute a script in python
# run_command is used in other parts of the codebase that import this module
# pylint: disable=too-many-locals,too-many-branches
from cai.sdk.agents import function_tool
@function_tool
def scripting_tool(
command: str = "",
args: str = "",

9
src/cai/tools/reconnaissance/crypto_tools.py
View File

@ -2,6 +2,9 @@
Here are crypto tools
"""
from cai.tools.common import run_command
from cai.sdk.agents import function_tool
# # URLDecodeTool
# # HexDumpTool
@ -9,7 +12,7 @@ from cai.tools.common import run_command
# # ROT13DecodeTool
# # BinaryAnalysisTool
@function_tool
def strings_command(file_path: str, ctf=None) -> str:
"""
Extract printable strings from a binary file.
@ -24,7 +27,7 @@ def strings_command(file_path: str, ctf=None) -> str:
command = f'strings {file_path}'
return run_command(command, ctf=ctf)
@function_tool
def decode64(input_data: str, ctf=None) -> str:
"""
Decode a base64-encoded string.
@ -39,7 +42,7 @@ def decode64(input_data: str, ctf=None) -> str:
command = f"base64 --decode {input_data}"
return run_command(command, ctf=ctf)
@function_tool
def decode_hex_bytes(input_data: str) -> str:
"""
Decode a string of hex bytes into ASCII text.

3
src/cai/tools/reconnaissance/curl.py
View File

@ -2,8 +2,9 @@
Here are the curl tools.
"""
from cai.tools.common import run_command # pylint: disable=import-error
from cai.sdk.agents import function_tool
@function_tool
def curl(args: str = "", target: str = "", ctf=None) -> str:
"""
A simple curl tool to make HTTP requests to a specified target.

3
src/cai/tools/reconnaissance/exec_code.py
View File

@ -2,7 +2,10 @@
Tool for executing code via LLM tool calls.
"""
from cai.tools.common import run_command # pylint: disable=import-error
from cai.sdk.agents import function_tool
@function_tool
def execute_code(code: str = "", language: str = "python",
filename: str = "exploit", timeout: int = 100, ctf=None) -> str:
"""

9
src/cai/tools/reconnaissance/filesystem.py
View File

@ -3,8 +3,9 @@ Here are the CLI tools for executing commands.
"""
from cai.tools.common import run_command # pylint: disable=E0401
from cai.sdk.agents import function_tool
@function_tool
def list_dir(path: str, args: str = "", ctf=None) -> str:
"""
List the contents of a directory.
@ -19,7 +20,7 @@ def list_dir(path: str, args: str = "", ctf=None) -> str:
command = f'ls {path} {args}'
return run_command(command, ctf=ctf)
@function_tool
def cat_file(file_path: str, args: str = "", ctf=None) -> str:
"""
Display the contents of a file.
@ -44,7 +45,7 @@ def cat_file(file_path: str, args: str = "", ctf=None) -> str:
# FilePermissionsTool
# FileCompressionTool
@function_tool
def pwd_command(ctf=None) -> str:
"""
Retrieve the current working directory.
@ -55,7 +56,7 @@ def pwd_command(ctf=None) -> str:
command = 'pwd'
return run_command(command, ctf=ctf)
@function_tool
def find_file(file_path: str, args: str = "", ctf=None) -> str:
"""
Find a file in the filesystem.

3
src/cai/tools/reconnaissance/netcat.py
View File

@ -2,8 +2,9 @@
Here are the tools for netcat command
"""
from cai.tools.common import run_command # pylint: disable=import-error
from cai.sdk.agents import function_tool
@function_tool
def netcat(host: str, port: int, data: str = '',
args: str = '', ctf=None) -> str:
"""

3
src/cai/tools/reconnaissance/netstat.py
View File

@ -3,8 +3,9 @@
Netstat tool
"""
from cai.tools.common import run_command # pylint: disable=import-error
from cai.sdk.agents import function_tool
@function_tool
def netstat(args: str = '', ctf=None) -> str:
"""
netstat tool to list all listening ports and their associated programs.

3
src/cai/tools/reconnaissance/nmap.py
View File

@ -3,8 +3,9 @@ Here are the nmap tools.
"""
from cai.tools.common import run_command # pylint: disable=E0401
from cai.sdk.agents import function_tool
@function_tool
def nmap(args: str, target: str, ctf=None) -> str:
"""
A simple nmap tool to scan a specified target.

4
src/cai/tools/reconnaissance/shodan.py
View File

@ -8,8 +8,10 @@ import os
import requests
from typing import Dict, List, Optional, Any
from dotenv import load_dotenv
from cai.sdk.agents import function_tool
@function_tool
def shodan_search(query: str, limit: int = 10) -> str:
"""
Search Shodan for information based on the provided query.
@ -41,7 +43,7 @@ def shodan_search(query: str, limit: int = 10) -> str:
return formatted_results
@function_tool
def shodan_host_info(ip: str) -> str:
"""
Get detailed information about a specific host from Shodan.

3
src/cai/tools/reconnaissance/wget.py
View File

@ -4,8 +4,9 @@
Wget tool
"""
from cai.tools.common import run_command # pylint: disable=import-error
from cai.sdk.agents import function_tool
@function_tool
def wget(url: str, args: str = '', ctf=None) -> str:
"""
Wget tool to download files from the web.

6
src/cai/tools/web/search_web.py
View File

@ -6,7 +6,10 @@ from cai.tools.web.google_search import (
google_dork_search,
google_search
)
from cai.sdk.agents import function_tool
@function_tool
def query_perplexity(query: str = "", context: str = "") -> str:
"""
Query the Perplexity AI API with a user prompt.
@ -55,7 +58,7 @@ def query_perplexity(query: str = "", context: str = "") -> str:
)
return response.choices[0].message.content
@function_tool
def make_web_search_with_explanation(context: str = "", query: str = "") -> str:
"""
Executes an intelligent web search via the AI service for relevant
@ -74,6 +77,7 @@ def make_web_search_with_explanation(context: str = "", query: str = "") -> str:
"""
return query_perplexity(query, context)
@function_tool
def make_google_search(query: str, dorks = False) -> str:
"""
Search Google for information.

93
tests/agents/test_agent_inference.py Normal file
View File

@ -0,0 +1,93 @@
import os
import pytest
from cai.sdk.agents import Runner
from cai.agents import get_agent_by_name
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_blue_team_agent_inference():
"""
Non-streaming inference test for the blueteam_agent.
"""
prompt = "Monitor login attempts for suspicious activity what we can do?"
result = await Runner.run(get_agent_by_name("blueteam_agent"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "login" in final_output.lower(), f"Expected 'login' in output, got: {final_output}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_bug_bounter_agent_inference():
"""
Non-streaming inference test for the bug_bounter_agent.
"""
prompt = "Find vulnerabilities in web application sample.com"
result = await Runner.run(get_agent_by_name("bug_bounter_agent"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "sample.com" in final_output.lower(), f"Expected 'sample.com' in output, got: {final_output}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_dns_smtp_agent_inference():
"""
Non-streaming inference test for the dns_smtp_agent.
"""
prompt = "check DKIM record for example.com"
result = await Runner.run(get_agent_by_name("dns_smtp_agent"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "dkim" in final_output.lower(), f"Expected 'dkim' in output, got: {final_output}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_one_tool_agent_inference():
"""
Non-streaming inference test for the one_tool_agent.
"""
prompt = "use nmap to scan 192.168.2.0/24"
result = await Runner.run(get_agent_by_name("one_tool_agent"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "generic_linux_command" in final_output.lower(), f"Expected 'generic_linux_command' in output, got: {final_output}"
assert "nmap" in final_output.lower(), f"Expected 'nmap' in output, got: {final_output}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_red_team_agent_inference():
"""
Non-streaming inference test for the redteam_agent.
"""
prompt = "perform penetration test on example.com ssh service"
result = await Runner.run(get_agent_by_name("redteam_agent"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "example.com" in final_output.lower(), f"Expected 'example.com' in output, got: {final_output}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_flag_discriminator_inference():
"""
Non-streaming inference test for the one_tool_agent.
"""
prompt = "Hello! Can you find tell me which is the flag in this string: 'Hi there, your reward flag{1234}"
result = await Runner.run(get_agent_by_name("flag_discriminator"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "flag{1234}" in final_output.lower(), f"Expected 'flag{1234}' in output, got: {final_output}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_thought_agent_inference():
"""
Non-streaming inference test for the one_tool_agent.
"""
prompt = """The phases of the cybersecurity kill chain are:
Reconnaissance, Exploitation, Lateral Movement, Data Exfiltration, and Command and Control (C2).
In which phase of the kill chain does this example fall:
Gathering initial intelligence about the target using OSINT techniques?"""
result = await Runner.run(get_agent_by_name("thought_agent"), prompt)
final_output = result.final_output or ""
assert final_output, "Expected non-empty final output"
assert "reconnaissance" in final_output.lower(), f"Expected 'reconnaissance' in output, got: {final_output}"