Add check_flag and implement time accounting

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

cai/core.py

@@ -13,11 +13,11 @@ and local modules.
 import copy
 import json
 from collections import defaultdict
-from typing import List
+from typing import List, Tuple
 # Package/library imports
+import time
 from openai import OpenAI  # pylint: disable=import-error
 
-
 # Local imports
 from .util import function_to_json, debug_print, merge_chunk
 from .types import (
@@ -331,10 +331,12 @@ class CAI:
         max_turns: int = float("inf"),
         execute_tools: bool = True,
         brief: bool = False,
-    ) -> Response:
+    ) -> Tuple[Response, float]:
         """
-        Run the cai and return the final response.
+        Run the cai and return the final response along
+        with execution time in seconds.
         """
+        start_time = time.time()
         self.brief = brief
         if stream:
             return self.run_and_stream(
@@ -388,8 +390,9 @@ class CAI:
             if partial_response.agent:
                 active_agent = partial_response.agent
 
+        execution_time = time.time() - start_time
         return Response(
             messages=history[init_len:],
             agent=active_agent,
             context_variables=context_variables,
-        )
+        ), execution_time

cai/util.py

@@ -6,6 +6,8 @@ import inspect
 from datetime import datetime
 from typing import Any
 import json
+import os
+from wasabi import color  # pylint: disable=import-error
 
 # ANSI color codes in a nice, readable palette
 COLORS = {
@@ -90,8 +92,8 @@ def format_value(value: Any, prev_value: Any = None, brief: bool = False) -> str
 
     # Handle basic types
     else:
-        color = get_color('arg_value', value, prev_value)
-        return f"{color}{str(value)}{COLORS['reset']}"
+        colorcillo = get_color('arg_value', value, prev_value)
+        return f"{colorcillo}{str(value)}{COLORS['reset']}"
 
 
 def format_chat_completion(msg, prev_msg=None) -> str:  # pylint: disable=unused-argument # noqa: E501
@@ -332,3 +334,47 @@ def function_to_json(func) -> dict:
             },
         },
     }
+
+
+def check_flag(output, ctf, challenge=None):
+    """
+    Check if the CTF flag is present in the output.
+
+    Args:
+        output (str): The output to check for the flag.
+        ctf: The CTF environment object.
+        challenge (str, optional): The specific challenge to check.
+            Defaults to None.
+
+    Returns:
+        tuple: A tuple containing a boolean indicating if the flag was
+            found and the flag itself if found, otherwise None.
+    """
+
+    # Get the challenge from the environment variable or default to the first
+    # challenge
+    challenge_key = os.getenv("CTF_CHALLENGE")
+    challenges = list(ctf.get_challenges().keys())
+    challenge = (
+        challenge_key
+        if challenge_key in challenges
+        else (challenges[0] if len(challenges) > 0 else None)
+    )
+
+    if ctf:
+        if ctf.check_flag(
+            output, challenge
+        ):  # check if the flag is in the output
+            flag = ctf.flags[challenge]
+            print(
+                color(
+                    f"Flag found: {flag}",
+                    fg="green") +
+                " in output " +
+                color(
+                    f"{output}",
+                    fg="blue"))
+            return True, flag
+    else:
+        print(color("CTF environment not found or provided", fg="yellow"))
+    return False, None

examples/cybersecurity/1_picoctf_static_flag.py

@@ -4,6 +4,7 @@ This is a short example of how we can solve ctf scenarios programmatically. With
 # 1. Instance CAI, pentestperf and tools
 from cai import CAI, Agent
 from cai.tools.reconnaissance.filesystem import list_dir, cat_file
+from cai.util import check_flag
 import pentestperf as ptt
 from wasabi import color
 import os
@@ -22,7 +23,7 @@ challenge = challenges[0]
 # 3. Instance the cybersecurity AI Framework client
 client = CAI(ctf=ctf)
 
-# 4. Instance the agent/agents to solve the CTF
+# 4. Create instances the agent/agents to solve the CTF
 ctf_agent = Agent(
     model="qwen2.5:14b",
     name="A file listing tool",
@@ -31,21 +32,42 @@ ctf_agent = Agent(
     # parallel_tool_calls=True, Enable multi function calling
 )
 
+flag_discriminator = Agent(
+    model="qwen2.5:14b",
+    name="A flag discriminator",
+    instructions="You are an agent tailored to extract the flag from the output. ONLY return the flag, nothing else.",
+    # functions=[check_flag],  # This code can execute by LLM
+)
+
+#  handoff from ctf_agent to flag_discriminator
+
+
+def transfer_to_flag_discriminator(**kwargs):
+    """Transfer flag discriminator. Accepts any keyword arguments but ignores them."""
+    return flag_discriminator
+
+
+ctf_agent.functions.append(transfer_to_flag_discriminator)
 
 messages = [{
     "role": "user",
     "content": "Instructions: " + ctf.get_instructions() +
                "\nChallenge: " + ctf.get_challenges()[challenge] +
                "\nTechniques: " + ctf.get_techniques() +
-               "\nExtract and return only the flag"
+               "\nExtract the flag and once finished, handoff to the flag discriminator."
 }]
 
 # 5. Run the CAI
-response = client.run(
+response, time = client.run(
     agent=ctf_agent,
     messages=messages,
     debug=True,
-    brief=True)
+    brief=False)
 print(response.messages[-1]["content"])
+print(f"Time taken: {time} seconds")
 
+# 6. Check if the flag is correct
+success, flag = check_flag(response.messages[-1]["content"], ctf, challenge)
+
+# 7. Stop the CTF environment
 ctf.stop_ctf()