From c20cf2d61d09dae6d8b949885e39113619fe61af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Mon, 13 Jan 2025 09:13:50 +0000 Subject: [PATCH] Add check_flag and implement time accounting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/core.py | 13 +++-- cai/util.py | 50 ++++++++++++++++++- .../cybersecurity/1_picoctf_static_flag.py | 30 +++++++++-- 3 files changed, 82 insertions(+), 11 deletions(-) diff --git a/cai/core.py b/cai/core.py index 948530fe..9d87a678 100644 --- a/cai/core.py +++ b/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 diff --git a/cai/util.py b/cai/util.py index 033083fc..6064744b 100644 --- a/cai/util.py +++ b/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 diff --git a/examples/cybersecurity/1_picoctf_static_flag.py b/examples/cybersecurity/1_picoctf_static_flag.py index 1c4f1821..669d0280 100644 --- a/examples/cybersecurity/1_picoctf_static_flag.py +++ b/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()