change CTF_NAME variable in the execution time

This commit is contained in:
Mery-Sanz 2025-05-20 10:01:38 +02:00
parent e398f46302
commit 9713a775f2
2 changed files with 18 additions and 2 deletions

View File

@ -157,6 +157,7 @@ from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig
from cai import is_pentestperf_available
ctf_global = None
messages_ctf = ""
previous_ctf_name = os.getenv('CTF_NAME', None)
if is_pentestperf_available() and os.getenv('CTF_NAME', None):
ctf, messages_ctf = setup_ctf()
ctf_global = ctf
@ -253,6 +254,18 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
agent.model.set_agent_name(get_agent_short_name(agent))
while turn_count < max_turns:
# Check if the ctf name has changed and instanciate the ctf
global previous_ctf_name
global ctf_global
global messages_ctf
if previous_ctf_name != os.getenv('CTF_NAME', None):
if is_pentestperf_available():
if ctf_global:
ctf_global.stop_ctf()
ctf, messages_ctf = setup_ctf()
ctf_global = ctf
previous_ctf_name = os.getenv('CTF_NAME', None)
try:
# Start measuring user idle time
start_idle_timer()
@ -291,7 +304,9 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
agent.model.set_agent_name(get_agent_short_name(agent))
except Exception as e:
console.print(f"[red]Error switching agent: {str(e)}[/red]")
if not force_until_flag:
# Get user input with command completion and history
user_input = get_user_input(

View File

@ -374,7 +374,7 @@ def terminate_session(session_id):
def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None):
"""Runs command in CTF env, changing to workspace_dir first."""
target_dir = workspace_dir or _get_workspace_dir()
full_command = f"cd '{target_dir}' && {command}"
full_command = f"{command}"
original_cmd_for_msg = command # For logging
context_msg = f"(ctf:{target_dir})"
try:
@ -1014,7 +1014,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501
# --- CTF Execution ---
if ctf and os.getenv('CTF_INSIDE', True).lower() == "true":
if ctf and os.getenv('CTF_INSIDE', "True").lower() == "true":
# If streaming is enabled and we have a call_id, show streaming UI for CTF too
if stream:
# Import the streaming utilities from util