diff --git a/src/cai/cli.py b/src/cai/cli.py index a7283c83..c1766b59 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -134,7 +134,10 @@ from cai.util import ( start_idle_timer, stop_idle_timer, start_active_timer, - stop_active_timer + stop_active_timer, + setup_ctf, + check_flag + ) # CAI REPL imports @@ -151,6 +154,11 @@ from cai.internal.components.metrics import process_metrics # Add import for parallel configs at the top of the file from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig +from cai import is_pentestperf_available +ctf_global = None +if is_pentestperf_available() and os.getenv('CTF_NAME', None): + ctf, messages_ctf = setup_ctf() + ctf_global = ctf # Load environment variables from .env file load_dotenv() @@ -187,7 +195,7 @@ agent = Agent( ) ) -def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')): +def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), force_until_flag=False): """ Run a simple interactive CLI loop for CAI. @@ -283,14 +291,17 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')): except Exception as e: console.print(f"[red]Error switching agent: {str(e)}[/red]") - # Get user input with command completion and history - user_input = get_user_input( - command_completer, - kb, - history_file, - get_toolbar_with_refresh, - current_text - ) + if not force_until_flag: + # Get user input with command completion and history + user_input = get_user_input( + command_completer, + kb, + history_file, + get_toolbar_with_refresh, + current_text + ) + else: + user_input = messages_ctf idle_time += time.time() - idle_start_time # Stop measuring user idle time and start measuring active time @@ -393,6 +404,11 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')): # Prevent duplicate cost display from the COST_TRACKER exit handler os.environ["CAI_COST_DISPLAYED"] = "true" + if (is_pentestperf_available() and os.getenv('CTF_NAME', None)): + ctf.stop_ctf() + return False + + except Exception: pass break diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 2650a827..a338302f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -708,6 +708,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Use the active timer during tool execution stop_idle_timer() start_active_timer() + + from cai.cli import ctf_global + ctf = ctf_global # Parse command into standard parts to ensure consistent naming parts = command.strip().split(' ', 1) diff --git a/src/cai/util.py b/src/cai/util.py index a33cb05d..46f07ea1 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -29,6 +29,9 @@ from rich.box import ROUNDED from rich.table import Table import re import uuid +from cai import is_pentestperf_available +if is_pentestperf_available(): + import pentestperf as ptt # Global timing variables for tracking active and idle time _active_timer_start = None @@ -2721,3 +2724,109 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, # Mark the streaming session as complete if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions: cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True + +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 + +def setup_ctf(): + """Setup CTF environment if CTF_NAME is provided""" + ctf_name = os.getenv('CTF_NAME', None) + if not ctf_name: + print(color("CTF name not provided, necessary to run CTF", fg="white", bg="red")) + sys.exit(1) + + print(color("Setting up CTF: ", fg="black", bg="yellow") + + color(ctf_name, fg="black", bg="yellow")) + + ctf = ptt.ctf( # pylint: disable=I1101 # noqa + ctf_name, + subnet=os.getenv('CTF_SUBNET', "192.168.2.0/24"), + container_name="ctf_target", + ip_address=os.getenv('CTF_IP', "192.168.2.100"), + ) + ctf.start_ctf() + + # Get the challenge from the environment variable or default to the + # first challenge + challenge_key = os.getenv('CTF_CHALLENGE') # TODO: + challenges = list(ctf.get_challenges().keys()) + challenge = challenge_key if challenge_key in challenges else ( + challenges[0] if len(challenges) > 0 else None) + + # Use the user master template + messages = Template( + filename="src/cai/prompts/core/user_master_template.md").render( + ctf=ctf, + challenge=challenge, + ip=ctf.get_ip() if ctf else None, + ) + + + print( color( + "Testing CTF: ", + fg="black", + bg="yellow") + + color( + ctf.name, + fg="black", + bg="yellow")) + if not challenge_key or challenge_key not in challenges: + print( + color( + "No challenge provided or challenge not found. Attempting to use the first challenge.", + fg="white", + bg="blue")) + if challenge: + print( + color( + "Testing challenge: ", + fg="white", + bg="blue") + + color( + "'" + + challenge + + "' (" + + repr( + ctf.flags[challenge]) + + ")", + fg="white", + bg="blue")) + + return ctf, messages