Merge branch 'ctf_in_docker_tools' into 'main'

Add tools

See merge request aliasrobotics/alias_research/cai!1

.env.example

@@ -0,0 +1 @@
+OPENAI_API_KEY=

.gitignore

@@ -145,3 +145,4 @@ pentestperf
 cai.png
 cai.dot
 cai*.png
+test_20*

README.md

@@ -355,7 +355,9 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa
 # Plan of development
 - [x] Dev container
 - [x] pre-commit hooks
-- [ ] A first example with agents - picoctf_static_flag
+- [x] A first example with agents - picoctf_static_flag
+  - [ ] polish aesthetics (coloring with alpha past content, brief/summarized)
+  - [ ] add tests
 - [ ] CI
 - [ ] Tracing
 - [ ] Graph/flow and other abstractions

cai/core.py

@@ -14,7 +14,6 @@ import copy
 import json
 from collections import defaultdict
 from typing import List
-
 # Package/library imports
 from openai import OpenAI  # pylint: disable=import-error
 
@@ -36,14 +35,18 @@ __CTX_VARS_NAME__ = "context_variables"
 
 class CAI:
     """
-    Main class for the CAI library.
+    Cybersecurity AI (CAI) object
     """
 
-    def __init__(self, client=None,
-                 base_url="http://localhost:8000/v1", api_key="alias"):
+    def __init__(self,
+                 client=None,
+                 base_url="http://host.docker.internal:8000/v1",
+                 api_key="alias",
+                 ctf=None):
         if not client:
             client = OpenAI(base_url=base_url, api_key=api_key)
         self.client = client
+        self.ctf = ctf
 
     def get_chat_completion(  # pylint: disable=too-many-arguments
         self,
@@ -91,7 +94,13 @@ class CAI:
 
     def handle_function_result(self, result, debug) -> Result:
         """
-        Handle the result of a function call.
+        Handle the result of a function call by
+        converting it into a standardized Result type.
+
+        The Result type encapsulates the possible
+        return values (Result, Agent, or context variables)
+        that functions can produce into a consistent
+        format for the framework to process.
         """
         match result:
             case Result() as result:
@@ -118,7 +127,37 @@ class CAI:
         debug: bool,
     ) -> Response:
         """
-        Handle the tool calls for the given agent.
+        Execute and handle tool calls made by the AI agent.
+
+        Processes a list of tool calls by:
+        1. Looking up each function in the provided function map
+        2. Handling missing tools gracefully by skipping them
+        3. Parsing and validating function arguments
+        4. Executing functions with provided arguments and context
+        5. Processing results into standardized Response format
+        6. Accumulating results from multiple tool calls
+
+        Args:
+            tool_calls (List[ChatCompletionMessageToolCall]): Tool
+                calls requested by AI agent
+            functions (List[AgentFunction]): Available functions
+                that can be called
+            context_variables (dict): Context variables to pass
+                to functions
+            debug (bool): Flag to enable debug logging
+
+        Returns:
+            Response: Object containing:
+                messages (List): Tool call results
+                agent (Optional[Agent]): Updated agent
+                    if returned by a function
+                context_variables (dict): Updated context variables
+
+        Note:
+            Results from multiple tool calls are accumulated
+                into a single Response.
+            Context variables are updated iteratively as
+                functions are called.
         """
         function_map = {f.__name__: f for f in functions}
         partial_response = Response(
@@ -140,12 +179,18 @@ class CAI:
                 continue
             args = json.loads(tool_call.function.arguments)
             debug_print(
-                debug, f"Processing tool call: {name} with arguments {args}")
+                debug,
+                "Processing tool call",
+                name,
+                "with arguments",
+                args)
 
             func = function_map[name]
             # pass context_variables to agent functions
             if __CTX_VARS_NAME__ in func.__code__.co_varnames:
                 args[__CTX_VARS_NAME__] = context_variables
+            if self.ctf:
+                args["ctf"] = self.ctf
             raw_result = function_map[name](**args)
 
             result: Result = self.handle_function_result(raw_result, debug)
@@ -175,6 +220,9 @@ class CAI:
     ):
         """
         Run the cai and stream the results.
+
+        The key difference from run() is that this streams results
+        incrementally, while run() returns everything at once.
         """
         active_agent = agent
         context_variables = copy.deepcopy(context_variables)

cai/tools/cli.py

@@ -0,0 +1,36 @@
+"""
+Here are the CLI tools for executing commands, they should usually return
+with run_command, but you can be creative
+"""
+
+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)

cai/tools/common.py

@@ -0,0 +1,51 @@
+"""
+Basic utilities for executing commands inside or outside of virtual
+containers.
+"""
+
+import subprocess  # nosec B404
+from wasabi import color  # pylint: disable=import-error
+
+
+def _run_ctf(ctf, command, stdout=True):
+    try:
+        # Ensure the command is executed in a shell that supports command
+        # chaining
+        output = ctf.get_shell(command)
+        # exploit_logger.log_ok()
+
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output
+    except Exception as e:  # 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_attacker_machine(command, stdout=True):
+    try:
+        # nosec B602 - shell=True is required for command chaining
+        result = subprocess.run(
+            command,
+            shell=True,  # nosec B602
+            capture_output=True,
+            text=True,
+            check=True)
+        output = result.stdout
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output
+    except Exception as e:  # pylint: disable=broad-except
+        print(color(f"Error executing local command: {e}", fg="red"))
+        return f"Error executing local command: {str(e)}"
+
+
+def run_command(command, ctf=None, stdout=True):
+    """
+    Run command either in CTF container or
+    on the local attacker machine
+    """
+    if ctf:
+        return _run_ctf(ctf, command, stdout)
+    return _run_attacker_machine(command, stdout)

cai/types.py

@@ -31,7 +31,10 @@ class Agent(BaseModel):  # pylint: disable=too-few-public-methods
 
 class Response(BaseModel):  # pylint: disable=too-few-public-methods
     """
-    Represents a response from the CAI.
+    Represents a response back to the user from the CAI.
+
+    NOTE: This happens within the run() chain, after "Ending turn"
+    in the CAI.
     """
 
     messages: List = []

cai/util.py

@@ -4,17 +4,142 @@ This module contains utility functions for the CAI library.
 
 import inspect
 from datetime import datetime
+from typing import Any
+import json
+
+# ANSI color codes in a nice, readable palette
+COLORS = {
+    'timestamp': '\033[38;5;75m',    # Light blue
+    'bracket': '\033[38;5;247m',     # Light gray
+    'intro': '\033[38;5;141m',       # Light purple
+    'object': '\033[38;5;215m',     # Light orange
+    'arg_key': '\033[38;5;228m',    # Light yellow
+    'arg_value': '\033[38;5;180m',  # Light tan
+    'function': '\033[38;5;219m',   # Pink
+    'tool': '\033[38;5;147m',       # Soft purple
+    'reset': '\033[0m'
+}
 
 
-def debug_print(debug: bool, *args: str) -> None:
+def format_value(value: Any) -> str:
     """
-    Print debug messages if debug mode is enabled.
+    Format a value for debug printing with appropriate colors.
+    """
+    # Handle ChatCompletionMessage objects
+    if hasattr(
+            value, '__class__') and 'ChatCompletionMessage' in value.__class__.__name__:  # noqa: E501 # pylint: disable=C0301
+        return format_chat_completion(value)
+
+    # Handle lists
+    if isinstance(value, list):  # pylint: disable=R1705
+        items = []
+        for item in value:
+            if isinstance(item, dict):
+                # Format dictionary items in the list
+                dict_items = [
+                    f"\n    {
+                        COLORS['arg_key']}{k}{
+                        COLORS['reset']}: {
+                        format_value(v)}"
+                    for k, v in item.items()
+                ]
+                items.append("{" + ",".join(dict_items) + "\n  }")
+            else:
+                items.append(format_value(item))
+        return f"[\n  {','.join(items)}\n]"
+
+    # Handle dictionaries
+    elif isinstance(value, dict):
+        formatted_items = [
+            f"{COLORS['arg_key']}{k}{COLORS['reset']}: {format_value(v)}"
+            for k, v in value.items()
+        ]
+        return "{ " + ", ".join(formatted_items) + " }"
+
+    # Handle basic types
+    else:
+        return f"{COLORS['arg_value']}{str(value)}{COLORS['reset']}"
+
+
+def format_chat_completion(msg) -> str:
+    """
+    Format a ChatCompletionMessage object with proper indentation and colors.
+    """
+    # Convert message to dict and handle OpenAI types
+    try:
+        msg_dict = json.loads(msg.model_dump_json())
+    except AttributeError:
+        msg_dict = msg.__dict__
+
+    # Clean up the dictionary
+    msg_dict = {k: v for k, v in msg_dict.items() if v is not None}
+
+    # Format with json.dumps for consistent indentation
+    formatted_json = json.dumps(msg_dict, indent=2)
+
+    # Color the different parts
+    colored_lines = []
+    for line in formatted_json.split('\n'):
+        if ':' in line:
+            key, value = line.split(':', 1)
+            # Handle nested structures
+            if value.strip() in ['{', '[', '}', ']']:
+                colored_lines.append(f"{COLORS['arg_key']}{key}{
+                                     COLORS['reset']}:{value}")
+            else:
+                colored_lines.append(
+                    f"{COLORS['arg_key']}{key}{COLORS['reset']}: "
+                    f"{COLORS['arg_value']}{value.strip()}{COLORS['reset']}"
+                )
+        else:
+            colored_lines.append(line)
+
+    return f"\n  {COLORS['object']}ChatCompletionMessage{
+        COLORS['reset']}(\n    " + '\n    '.join(colored_lines) + "\n  )"
+
+
+def debug_print(debug: bool, intro: str, *args: Any) -> None:
+    """
+    Print debug messages if debug mode is enabled with color-coded components.
     """
     if not debug:
         return
+
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    message = " ".join(map(str, args))
-    print(f"\033[97m[\033[90m{timestamp}\033[97m]\033[90m {message}\033[0m")
+    header = f"{COLORS['bracket']}[{COLORS['timestamp']}{
+        timestamp}{COLORS['bracket']}]{COLORS['reset']}"
+
+    # Special handling for tool call processing messages
+    if "Processing tool call" in intro:
+        if len(args) >= 2:
+            tool_name, _, tool_args = args
+            message = (
+                f"{header} {
+                    COLORS['intro']}Processing tool call:{
+                    COLORS['reset']} "
+                f"{COLORS['tool']}{tool_name}{COLORS['reset']} "
+                f"{COLORS['intro']}with arguments{COLORS['reset']} "
+                f"{format_value(tool_args)}"
+            )
+        else:
+            message = f"{header} {COLORS['intro']}{intro}{COLORS['reset']}"
+    else:
+        formatted_intro = f"{COLORS['intro']}{intro}{COLORS['reset']}"
+        formatted_args = []
+        for arg in args:
+            if isinstance(arg, str) and arg.startswith(
+                    ('get_', 'list_', 'process_', 'handle_')):
+                formatted_args.append(f"{COLORS['function']}{
+                                      arg}{COLORS['reset']}")
+            elif hasattr(arg, '__class__'):
+                formatted_args.append(format_value(arg))
+            else:
+                formatted_args.append(format_value(arg))
+
+        message = f"{header} {formatted_intro} {
+            ' '.join(map(str, formatted_args))}"
+
+    print(message)
 
 
 def merge_fields(target, source):

examples/cybersecurity/1_arch_short_picoctf_static_flag.py

@@ -0,0 +1,47 @@
+"""
+This is a short example of how we can solve ctf scenarios programmatically. With CAI framework
+"""
+# 1. Instance CAI, pentestperf and tools
+from cai import CAI, Agent
+from cai.tools.cli import list_dir, cat_file
+import pentestperf as ptt
+from wasabi import color
+import os
+
+# 2. Init the CTF environment
+ctf = ptt.ctf(
+    "picoctf_static_flag",
+    subnet="192.168.9.0/24",
+    container_name="ancient",
+    ip_address="192.168.9.52"
+)
+ctf.start_ctf()
+challenges = list(ctf.get_challenges().keys())
+challenge = challenges[0]
+
+# 3. Instance the cybersecurity AI Framework client
+client = CAI(ctf=ctf)
+
+# 4. Instance the agent/agents to solve the CTF
+ctf_agent = Agent(
+    model="qwen2.5:14b",
+    name="A file listing tool",
+    instructions="You are an agent tailored to list and read files within a file system.",
+    functions=[list_dir, cat_file],  # This code can execute by LLM
+    # parallel_tool_calls=True, Enable multi function calling
+)
+
+
+messages = [{
+    "role": "user",
+    "content": "Instructions: " + ctf.get_instructions() +
+               "\nChallenge: " + ctf.get_challenges()[challenge] +
+               "\nTechniques: " + ctf.get_techniques() +
+               "\nExtract and return only the flag"
+}]
+
+# 5. Run the swarm
+response = client.run(agent=ctf_agent, messages=messages, debug=True)
+print(response.messages[-1]["content"])
+
+ctf.stop_ctf()

examples/cybersecurity/1_picoctf_static_flag.py

@@ -1,42 +1,113 @@
 from cai import CAI, Agent
 import pentestperf as ptt
+from wasabi import color
 
 client = CAI()
 
-english_agent = Agent(
+
+def run_ctf(ctf, command, stdout=True):
+    try:
+        # Ensure the command is executed in a shell that supports command
+        # chaining
+        output = ctf.get_shell(command)
+        # exploit_logger.log_ok()
+
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output
+    except Exception as e:
+        print(color(f"Error executing CTF command: {e}", fg="red"))
+        # exploit_logger.log_error(str(e))
+        return f"Error executing CTF command: {str(e)}"
+
+
+def list_dir(path: str, args: str = "") -> str:
+    """
+    List the contents of a directory.
+
+    Args:
+        path: The directory path to list contents from
+        args: Additional arguments to pass to the ls command
+
+    Returns:
+        str: The output of running the ls command
+    """
+    global ctf
+    command = f'ls {path} {args}'
+    return run_ctf(ctf, command)
+
+
+def cat_file(args: str, file_path: str) -> str:
+    """
+    Display the contents of a file.
+
+    Args:
+        args: Additional arguments to pass to the cat command
+        file_path: Path to the file to display contents of
+
+    Returns:
+        str: The output of running the cat command
+    """
+    global ctf
+    command = f'cat {args} {file_path} '
+    return run_ctf(ctf, command)
+
+# def flag_checker(flag: str) -> str:
+#      """
+#      This tool is used to check if the guessed flag is valid.
+#      """
+#      if ctf.get_flag(ctf.get_challenges()[challenge]) == flag:
+#         return "Flag Found: " + flag
+#      else:
+#         return "Flag Not Found: " + flag
+
+
+listing_agent = Agent(
     model="qwen2.5:14b",
     name="A file listing tool",
-    instructions="You only speak English. If another language is detected, invoke transfer_to_spanish_agent.",
-    # instructions="You only speak English.",
-    # tool_choice="required",  # not working with ollama and qwen2.5
+    instructions="You are an agent tailored to list and read files within a file system.",
+    # parallel_tool_calls=True,
 )
 
-spanish_agent = Agent(
-    name="Spanish Agent",
-    instructions="You only speak Spanish.",
+# flag_extractor = Agent(
+#     model="qwen2.5:14b",
+#     name="A flag extractor",
+#     instructions="You are an agent tailored to extract only the flag from a given content.",
+# )
+
+# def transfer_to_flag_extractor_agent():
+#     """Transfer to flag extractor."""
+#     return flag_extractor
+
+
+#######################################
+
+# PentestPerf setup for "picoctf_static_flag"
+ctf = ptt.ctf(
+    "picoctf_static_flag",
+    subnet="192.168.9.0/24",
+    container_name="ancient",
+    ip_address="192.168.9.52",
 )
+ctf.start_ctf()
+flag = ctf.get_flag()
 
 
-def transfer_to_spanish_agent():
-    """Transfer spanish speaking users immediately."""
-    return spanish_agent
+listing_agent.functions.append(list_dir)
+listing_agent.functions.append(cat_file)
+# listing_agent.functions.append(flag_checker)
 
+challenges = list(ctf.get_challenges().keys())
+challenge = challenges[0]
 
-english_agent.functions.append(transfer_to_spanish_agent)
-messages = [{"role": "user", "content": "Hola, cómo estás?"}]
-
-
-response = client.run(agent=english_agent, messages=messages, debug=True)
+messages = [{"role": "user", "content": "Instructions: " + ctf.get_instructions()
+             + "\nChallenge: " + ctf.get_challenges()[challenge]
+             + "\nTechniques: " + ctf.get_techniques()
+             + "\nExtract and return only the flag"}]
+response = client.run(agent=listing_agent, messages=messages, debug=True)
 print(response.messages[-1]["content"])
 
+# if ctf.check_flag(response.messages[-1]["content"], ctf.get_challenges()[challenge]):
+#    print(color(f"Flag found: {response.messages[-1]['content']}", fg="green"))
 
-# # PentestPerf setup for "picoctf_static_flag"
-# ctf = ptt.ctf(
-#         "picoctf_static_flag",
-#         subnet="192.168.9.0/24",
-#         container_name="ancient",
-#         ip_address="192.168.9.52",
-#     )
-# ctf.start_ctf()
-# flag = ctf.get_flag()
-# ctf.stop_ctf()
+ctf.stop_ctf()

setup.cfg

@@ -19,6 +19,7 @@ install_requires =
     tqdm
     pre-commit
     instructor
+    wasabi
     pentestperf @ git+https://gitlab:gldt-1NfZ2XRWdHNip4K45auL@gitlab.com/aliasrobotics/alias_research/pentestperf@main
 python_requires = >=3.10
 

tests/CAI/test_runs/test_core.py

@@ -1,6 +1,6 @@
 import pytest
 from cai import CAI, Agent
-from tests.mock_client import MockOpenAIClient, create_mock_response
+from .mock_client import MockOpenAIClient, create_mock_response
 from unittest.mock import Mock
 import json
 

tests/test_runs/test_20240402-113647.json

@@ -1 +0,0 @@
-[{"task_id": "02b37e8e-e436-445c-abdc-13e227616e07", "role": "user", "content": "If I have 5 ducks, and lose 2 of them. How many do I have left"}, {"task_id": "02b37e8e-e436-445c-abdc-13e227616e07", "role": "assistant", "content": "Response to user: 3 ducks"}]