add info

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

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)

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