From 48c6046a03c17fef65eaee54dfd45b4e5f534544 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:15:20 +0200 Subject: [PATCH 01/25] add info --- src/cai/cli.py | 36 +++++++++---- src/cai/tools/common.py | 3 ++ src/cai/util.py | 109 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 10 deletions(-) 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 From 906b874f2eaaf0319bee662d1b0aaa9b1f8b2611 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:26:25 +0200 Subject: [PATCH 02/25] inside ctf is working --- src/cai/tools/common.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index a338302f..577f122f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1014,7 +1014,7 @@ 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: + if ctf and os.getenv('CTF_INSIDE', True) == "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 @@ -1037,8 +1037,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg call_id = start_tool_streaming(tool_name, tool_args, call_id) target_dir = _get_workspace_dir() - full_command = f"cd '{target_dir}' && {command}" - + #full_command = f"cd '{target_dir}' && {command}" + full_command = command # Update with "executing" status update_tool_streaming( tool_name, From 10668f1b44eb4596bf85abb5c7f8a40c67c1d359 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:44:53 +0200 Subject: [PATCH 03/25] add true lower --- src/cai/tools/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 577f122f..dbad944a 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1014,7 +1014,7 @@ 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) == "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 From bfe87ebf17962ae9edb4b756ec69f4dc485fd8c6 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:36:19 +0200 Subject: [PATCH 04/25] solve issue --- src/cai/tools/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index dbad944a..e31aba92 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -759,7 +759,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) # --- Docker Container Execution --- - if active_container and not ctf and not is_ssh_env: + if active_container and not is_ssh_env: container_id = active_container container_workspace = _get_container_workspace_path() context_msg = f"(docker:{container_id[:12]}:{container_workspace})" From 6d92e2923e4a6fa184878b5e9e618dadd75a8edb Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:56:30 +0200 Subject: [PATCH 05/25] fix --- src/cai/cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index c1766b59..8a6c3269 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -156,6 +156,7 @@ from cai.internal.components.metrics import process_metrics from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig from cai import is_pentestperf_available ctf_global = None +messages_ctf = "" if is_pentestperf_available() and os.getenv('CTF_NAME', None): ctf, messages_ctf = setup_ctf() ctf_global = ctf @@ -299,7 +300,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), history_file, get_toolbar_with_refresh, current_text - ) + ) + messages_ctf else: user_input = messages_ctf idle_time += time.time() - idle_start_time From 18439c504f5291cb7008f59cfb40f1a5f015c8cf Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:15:20 +0200 Subject: [PATCH 06/25] add info --- src/cai/cli.py | 36 +++++++++---- src/cai/tools/common.py | 3 ++ src/cai/util.py | 109 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 10 deletions(-) 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 1846d56b..836a41fb 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 d576202f..82081f95 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 @@ -2747,3 +2750,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 From cdafb569c23827bc96af0f594c4f5c4c7a7f8df7 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:26:25 +0200 Subject: [PATCH 07/25] inside ctf is working --- src/cai/tools/common.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 836a41fb..11c66772 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1014,7 +1014,7 @@ 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: + if ctf and os.getenv('CTF_INSIDE', True) == "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 @@ -1037,8 +1037,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg call_id = start_tool_streaming(tool_name, tool_args, call_id) target_dir = _get_workspace_dir() - full_command = f"cd '{target_dir}' && {command}" - + #full_command = f"cd '{target_dir}' && {command}" + full_command = command # Update with "executing" status update_tool_streaming( tool_name, From 4c2c979e4831651989a06d5dbdeb77fcae1df04c Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:44:53 +0200 Subject: [PATCH 08/25] add true lower --- src/cai/tools/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 11c66772..61bc2a5f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1014,7 +1014,7 @@ 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) == "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 From 46e00813bbf3bbbffd2dd195640b7fdc07d70e6d Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:36:19 +0200 Subject: [PATCH 09/25] solve issue --- src/cai/tools/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 61bc2a5f..8b9e3e46 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -759,7 +759,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) # --- Docker Container Execution --- - if active_container and not ctf and not is_ssh_env: + if active_container and not is_ssh_env: container_id = active_container container_workspace = _get_container_workspace_path() context_msg = f"(docker:{container_id[:12]}:{container_workspace})" From fd60febfbac6c1f68a07ef809834f143fded97a9 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:56:30 +0200 Subject: [PATCH 10/25] fix --- src/cai/cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index c1766b59..8a6c3269 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -156,6 +156,7 @@ from cai.internal.components.metrics import process_metrics from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig from cai import is_pentestperf_available ctf_global = None +messages_ctf = "" if is_pentestperf_available() and os.getenv('CTF_NAME', None): ctf, messages_ctf = setup_ctf() ctf_global = ctf @@ -299,7 +300,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), history_file, get_toolbar_with_refresh, current_text - ) + ) + messages_ctf else: user_input = messages_ctf idle_time += time.time() - idle_start_time From 9713a775f20238de11e1fa9ff280307702b8ed5a Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 20 May 2025 10:01:38 +0200 Subject: [PATCH 11/25] change CTF_NAME variable in the execution time --- src/cai/cli.py | 15 +++++++++++++++ src/cai/tools/common.py | 5 +++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index 8a6c3269..38cdc8bf 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -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( diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 8b9e3e46..3c497f0a 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -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 From 8803a05ef41950a40ab6e1e715cc6688a8f66e94 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 20 May 2025 16:05:20 +0200 Subject: [PATCH 12/25] docker virtualization for ctf_inside false --- .devcontainer/Dockerfile | 108 ++++++++++++++++++ .devcontainer/devcontainer.json | 86 ++++++++++++++ .devcontainer/docker-compose.yml | 78 +++++++++++++ .../qdrant_storage/aliases/data.json | 1 + .devcontainer/qdrant_storage/raft_state.json | 1 + .devcontainer/requirements.txt | 6 + src/cai/cli.py | 15 ++- src/cai/util.py | 38 ++++++ 8 files changed, 331 insertions(+), 2 deletions(-) create mode 100644 .devcontainer/Dockerfile create mode 100644 .devcontainer/devcontainer.json create mode 100644 .devcontainer/docker-compose.yml create mode 100644 .devcontainer/qdrant_storage/aliases/data.json create mode 100644 .devcontainer/qdrant_storage/raft_state.json create mode 100644 .devcontainer/requirements.txt diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 00000000..43fe72ae --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,108 @@ +FROM kalilinux/kali-rolling + +# Set environment variable to non-interactive +ENV DEBIAN_FRONTEND=noninteractive + +# [Optional] Uncomment this section to install additional OS packages. +RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ + && apt-get -y install --no-install-recommends \ + net-tools python3 python3-pip python3-venv \ + curl gnupg nmap iputils-ping ssh git \ + graphviz pkg-config libhdf5-dev \ + build-essential python3-dev \ + asciinema dnsutils \ + apt-transport-https ca-certificates \ + wget dirb gobuster whatweb gfortran \ + cmake libopenblas-dev sshpass seclists \ + golang-go + +# # Install Metasploit - path 1 +# RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > /tmp/msfinstall \ +# && chmod 755 /tmp/msfinstall \ +# && /tmp/msfinstall + +# # Install Metasploit - path 2 +# RUN curl https://apt.metasploit.com/metasploit-framework.gpg.key | apt-key add - +# RUN echo "deb https://apt.metasploit.com/ lucid main" | tee /etc/apt/sources.list.d/metasploit-framework.list +# RUN apt-get update && apt-get install metasploit-framework + +# Install Metasploit - path 3 +RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstall && chmod 755 msfinstall && ./msfinstall + +# Update Kali Linux repositories +RUN apt-get update \ + && apt-get upgrade -y + +# # Create a virtual environment +# ENV VIRTUAL_ENV=/opt/venv +# RUN python3 -m venv $VIRTUAL_ENV +# ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +# # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image. +# # alternatively, this also helps avoid having to pull requirements from the internet every single time +# COPY requirements.txt /tmp/pip-tmp/ +# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ +# && rm -rf /tmp/pip-tmp + +# FIX for: +# +# Python packages system-wide, try apt install +# python3-xyz, where xyz is the package you are trying to +# install. +# +# see https://github.com/raspiblitz/raspiblitz/issues/4170 +RUN mkdir -p /root/.pip && touch /root/.pip/pip.conf +RUN echo "[global]" > /root/.pip/pip.conf && \ + echo "break-system-packages = true" >> /root/.pip/pip.conf + +RUN pip3 install --upgrade setuptools +RUN pip3 install pandas \ + opentelemetry-sdk opentelemetry-exporter-otlp \ + mem0ai PyPDF2 sentence_transformers tf-keras \ + pytest-repeat parameterized pytest-rerunfailures \ + pytest-clarity tiktoken blobfile pytest-timeout \ + invoke fabric + +# Install rust compiler - dependency of outlines-core +RUN apt-get update && \ + apt-get install -y curl && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \ + . $HOME/.cargo/env && \ + rustup default stable && \ + rustup component add rustfmt clippy && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +ENV PATH="/root/.cargo/bin:${PATH}" +RUN wget https://raw.githubusercontent.com/adfoster-r7/metasploit-info/master/info/module_metadata.json -O /tmp/module_metadata.json + +# NOTE: align dev env in Linux +# work +# RUN echo "192.168.2.1 host.docker.internal" >> /etc/hosts +# Update and upgrade packages +# Preconfigure console-setup and install kali-linux-headless fixers +RUN apt-get update && \ + echo "console-setup console-setup/variant select Latin1 and Latin5 - western Europe and Turkic languages" | debconf-set-selections && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends kali-linux-headless + +# Update pip again -> python3 -m pip install --upgrade pip causes crashes +# Install spacy deps - en_core_web_sm +# RUN python3 -m spacy download en_core_web_sm + +# Install docker inside of dev env +RUN apt-get update && \ + apt-get install -y apt-transport-https ca-certificates curl gnupg && \ + curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg && \ + echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian bullseye stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && \ + apt-get update && \ + apt-get install -y docker-ce docker-ce-cli containerd.io cloc + +# # Activate virtual environment by default +# RUN echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc + +RUN cargo install --git https://github.com/asciinema/agg + +# Remove system's python3 libs to avoid conflicts +RUN apt-get remove -y python3-jsonschema && \ + apt-get remove -y python3-numpy && \ + apt-get autoremove -y diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 00000000..22e1192e --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,86 @@ +// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at: +// https://github.com/microsoft/vscode-dev-containers/tree/v0.241.1/containers/python-3 +{ + "name": "cai_devenv", + + // "build": { + // "dockerfile": "Dockerfile", + // "context": "..", + // "args": { + // // Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6 + // // Append -bullseye or -buster to pin to an OS version. + // // Use -bullseye variants on local on arm64/Apple Silicon. + // "VARIANT": "3.10-bullseye", + // // Options + // "NODE_VERSION": "lts/*" + // } + // }, + + "dockerComposeFile": ["./docker-compose.yml"], + "service": "devenv", + // "shutdownAction": "none", // don't shut down container when vscode is closed + "workspaceFolder": "/workspace", + + // Configure tool-specific properties. + "customizations": { + // Configure properties specific to VS Code. + "vscode": { + // Set *default* container specific settings.json values on container create. + "settings": { + "python.defaultInterpreterPath": "/usr/local/bin/python", + "python.linting.enabled": true, + "python.linting.pylintEnabled": true, + "python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8", + "python.formatting.blackPath": "/usr/local/py-utils/bin/black", + "python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf", + "python.linting.banditPath": "/usr/local/py-utils/bin/bandit", + "python.linting.flake8Path": "/usr/local/py-utils/bin/flake8", + "python.linting.mypyPath": "/usr/local/py-utils/bin/mypy", + "python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle", + "python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", + "python.linting.pylintPath": "/usr/local/py-utils/bin/pylint" + }, + + // Add the IDs of extensions you want installed when the container is created. + "extensions": [ + "ms-python.python", + "ms-toolsai.jupyter-renderers", + "ms-toolsai.jupyter", + "ms-python.vscode-pylance", + "ms-azuretools.vscode-docker", + "ms-python.debugpy", + "ms-python.black-formatter", + "MS-vsliveshare.vsliveshare" + ] + } + }, + + // Use 'forwardPorts' to make a list of ports inside the container available locally. + // "forwardPorts": [], + + // Use 'preCreateCommand' to run commands before the container is created. + "preCreateCommand": "docker network prune -f && docker container prune -f", + + // Use 'postCreateCommand' to run commands after the container is created. + // "postCreateCommand": "pip3 install --user -r requirements.txt", // doing it instead in + // the dockerfile to cache + // requirements in container + // Install cai via "pip3 install -e ." + // Extensions, NOTE: only one postCreateCommand is allowed, so we need to install all extensions here + "postCreateCommand": "pip3 install -e /workspace/", + + // MSF setup, and RAG setup + "postStartCommand": [ + "nohup", "msfrpcd", "-P", "cai", "&", + "&&", + "python3", "cai/ins/rag/agent_helper.py" + ], + // // Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root. + // "remoteUser": "vscode" + + "runArgs": [ + "--privileged", + "-e", "DISPLAY=host.docker.internal:0", + "-v", "/tmp/.X11-unix:/tmp/.X11-unix" + ] +} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml new file mode 100644 index 00000000..92b14fb9 --- /dev/null +++ b/.devcontainer/docker-compose.yml @@ -0,0 +1,78 @@ +version: '3' + +################# +# SERVICES +################# +services: + # Developer environment + devenv: + build: + context: .. + dockerfile: .devcontainer/Dockerfile + # platform: linux/amd64 # TODO: uncomment this when we have a x86_64 machine + volumes: + # Mount the root folder that contains .git + - ..:/workspace:cached + # Mount Docker socket to enable Docker-in-Docker + - /var/run/docker.sock:/var/run/docker.sock + command: /bin/sh -c "while sleep 10; do :; done" + networks: + cainet: + ipv4_address: 192.168.2.5 + cap_add: + - NET_ADMIN + - SYS_ADMIN + privileged: true + + # # Vulnhub vulnerable machine (Hackable II) SUGGESTION: Comment this and use PentestPerf + # hackableii: + # image: vmayoral/vulnhub:hackableii + # command: | + # /bin/bash -c " + # rm -rf /var/lock && \ + # mkdir -p /var/lock && \ + # chmod 755 /var/lock && \ + # /etc/init.d/apache2 start && \ + # /etc/init.d/ssh start && \ + # /etc/init.d/runproftpd.sh && \ + # /etc/init.d/php7.0-fpm start && \ + # while sleep 1; do :; done + # " + # networks: + # cainet: + # ipv4_address: 192.168.2.11 + # mac_address: 08:00:27:85:55:86 + + # # Vulnhub vulnerable machine (Bob) SUGGESTION: Comment this and use PentestPerf + # bob: + # image: vmayoral/vulnhub:bob + # command: | + # /bin/bash -c "rm -r /var/lock; mkdir -p /var/lock; chmod 755 /var/lock; /etc/init.d/apache2 start; /etc/init.d/ssh start; while sleep 10; do :; done" + # ports: # map port in the container to the host systems + # - "8080:80" + # networks: + # cainet: + # ipv4_address: 192.168.2.12 + # mac_address: 08:00:27:cb:07:d4 + + # Qdrant vector database + qdrant: + image: qdrant/qdrant + ports: + - "6333:6333" + - "6334:6334" + volumes: + - ./qdrant_storage:/qdrant/storage:z + networks: + cainet: + ipv4_address: 192.168.2.13 + +################# +# NETWORKS +################# +networks: + cainet: + ipam: + driver: default + config: + - subnet: 192.168.2.0/24 diff --git a/.devcontainer/qdrant_storage/aliases/data.json b/.devcontainer/qdrant_storage/aliases/data.json new file mode 100644 index 00000000..9e26dfee --- /dev/null +++ b/.devcontainer/qdrant_storage/aliases/data.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/.devcontainer/qdrant_storage/raft_state.json b/.devcontainer/qdrant_storage/raft_state.json new file mode 100644 index 00000000..0ee33f67 --- /dev/null +++ b/.devcontainer/qdrant_storage/raft_state.json @@ -0,0 +1 @@ +{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6138497487128594],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6138497487128594,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6138497487128594} \ No newline at end of file diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt new file mode 100644 index 00000000..2409daf2 --- /dev/null +++ b/.devcontainer/requirements.txt @@ -0,0 +1,6 @@ +pymetasploit3 +networkx==2.5 +requests +wasabi +xmltodict +pydot==1.4.2 diff --git a/src/cai/cli.py b/src/cai/cli.py index 38cdc8bf..1bba0c6b 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -136,8 +136,9 @@ from cai.util import ( start_active_timer, stop_active_timer, setup_ctf, - check_flag - + check_flag, + build_docker_container, + remove_container ) # CAI REPL imports @@ -161,6 +162,10 @@ 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 + if os.getenv('CTF_INSIDE', 'True').lower() == 'false': + # instanciate dev container + container_id = build_docker_container("./.devcontainer") + os.environ['CAI_ACTIVE_CONTAINER'] = container_id # Load environment variables from .env file load_dotenv() @@ -266,6 +271,10 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), ctf, messages_ctf = setup_ctf() ctf_global = ctf previous_ctf_name = os.getenv('CTF_NAME', None) + if os.getenv('CTF_INSIDE', 'True') == 'False': + # instanciate dev container + container_id = build_docker_container("./.devcontainer") + os.environ['CAI_ACTIVE_CONTAINER'] = container_id try: # Start measuring user idle time start_idle_timer() @@ -422,6 +431,8 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), if (is_pentestperf_available() and os.getenv('CTF_NAME', None)): ctf.stop_ctf() + if os.getenv('CTF_INSIDE', 'True').lower() == 'false': + remove_container(os.environ['CAI_ACTIVE_CONTAINER']) return False diff --git a/src/cai/util.py b/src/cai/util.py index 82081f95..9a769037 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -3,6 +3,7 @@ Util model for CAI """ import os import sys +import subprocess import importlib.resources import pathlib import json @@ -2856,3 +2857,40 @@ def setup_ctf(): bg="blue")) return ctf, messages + + +def build_docker_container(devcontainer_path: str) -> str: + dockerfile_path = os.path.join(devcontainer_path, "Dockerfile") + image_name = "cai-container3" + container_name = "cai-container3" + + # Try to remove existing container with the same name + try: + subprocess.run(["docker", "rm", "-f", container_name], check=True) + except subprocess.CalledProcessError: + pass # If it doesn't exist, continue normally + + # Build the image + subprocess.run([ + "docker", "build", + "-t", image_name, + "-f", dockerfile_path, + devcontainer_path + ], check=True) + + # Run the container + result = subprocess.run([ + "docker", "run", "-d", "--name", container_name, image_name, "tail", "-f", "/dev/null" + ], stdout=subprocess.PIPE, check=True, text=True) + + container_id = result.stdout.strip() + return container_id + +def remove_container(name_or_id: str): + try: + # Force removal (stops if running, then removes) + subprocess.run(["docker", "rm", "-f", name_or_id], check=True) + print(f"Container '{name_or_id}' removed successfully.") + except subprocess.CalledProcessError as e: + print(f"Error removing container {name_or_id}:", e) + From 5c290e7a40beb2ae72738e66f562054fb94d6e85 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Wed, 21 May 2025 10:07:58 +0200 Subject: [PATCH 13/25] remove devcontainer --- .devcontainer/Dockerfile | 108 ------------------ .devcontainer/devcontainer.json | 86 -------------- .devcontainer/docker-compose.yml | 78 ------------- .../qdrant_storage/aliases/data.json | 1 - .devcontainer/qdrant_storage/raft_state.json | 1 - .devcontainer/requirements.txt | 6 - src/cai/cli.py | 15 +-- src/cai/util.py | 37 ------ 8 files changed, 2 insertions(+), 330 deletions(-) delete mode 100644 .devcontainer/Dockerfile delete mode 100644 .devcontainer/devcontainer.json delete mode 100644 .devcontainer/docker-compose.yml delete mode 100644 .devcontainer/qdrant_storage/aliases/data.json delete mode 100644 .devcontainer/qdrant_storage/raft_state.json delete mode 100644 .devcontainer/requirements.txt diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile deleted file mode 100644 index 43fe72ae..00000000 --- a/.devcontainer/Dockerfile +++ /dev/null @@ -1,108 +0,0 @@ -FROM kalilinux/kali-rolling - -# Set environment variable to non-interactive -ENV DEBIAN_FRONTEND=noninteractive - -# [Optional] Uncomment this section to install additional OS packages. -RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ - && apt-get -y install --no-install-recommends \ - net-tools python3 python3-pip python3-venv \ - curl gnupg nmap iputils-ping ssh git \ - graphviz pkg-config libhdf5-dev \ - build-essential python3-dev \ - asciinema dnsutils \ - apt-transport-https ca-certificates \ - wget dirb gobuster whatweb gfortran \ - cmake libopenblas-dev sshpass seclists \ - golang-go - -# # Install Metasploit - path 1 -# RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > /tmp/msfinstall \ -# && chmod 755 /tmp/msfinstall \ -# && /tmp/msfinstall - -# # Install Metasploit - path 2 -# RUN curl https://apt.metasploit.com/metasploit-framework.gpg.key | apt-key add - -# RUN echo "deb https://apt.metasploit.com/ lucid main" | tee /etc/apt/sources.list.d/metasploit-framework.list -# RUN apt-get update && apt-get install metasploit-framework - -# Install Metasploit - path 3 -RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstall && chmod 755 msfinstall && ./msfinstall - -# Update Kali Linux repositories -RUN apt-get update \ - && apt-get upgrade -y - -# # Create a virtual environment -# ENV VIRTUAL_ENV=/opt/venv -# RUN python3 -m venv $VIRTUAL_ENV -# ENV PATH="$VIRTUAL_ENV/bin:$PATH" - -# # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image. -# # alternatively, this also helps avoid having to pull requirements from the internet every single time -# COPY requirements.txt /tmp/pip-tmp/ -# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ -# && rm -rf /tmp/pip-tmp - -# FIX for: -# -# Python packages system-wide, try apt install -# python3-xyz, where xyz is the package you are trying to -# install. -# -# see https://github.com/raspiblitz/raspiblitz/issues/4170 -RUN mkdir -p /root/.pip && touch /root/.pip/pip.conf -RUN echo "[global]" > /root/.pip/pip.conf && \ - echo "break-system-packages = true" >> /root/.pip/pip.conf - -RUN pip3 install --upgrade setuptools -RUN pip3 install pandas \ - opentelemetry-sdk opentelemetry-exporter-otlp \ - mem0ai PyPDF2 sentence_transformers tf-keras \ - pytest-repeat parameterized pytest-rerunfailures \ - pytest-clarity tiktoken blobfile pytest-timeout \ - invoke fabric - -# Install rust compiler - dependency of outlines-core -RUN apt-get update && \ - apt-get install -y curl && \ - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \ - . $HOME/.cargo/env && \ - rustup default stable && \ - rustup component add rustfmt clippy && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* - -ENV PATH="/root/.cargo/bin:${PATH}" -RUN wget https://raw.githubusercontent.com/adfoster-r7/metasploit-info/master/info/module_metadata.json -O /tmp/module_metadata.json - -# NOTE: align dev env in Linux -# work -# RUN echo "192.168.2.1 host.docker.internal" >> /etc/hosts -# Update and upgrade packages -# Preconfigure console-setup and install kali-linux-headless fixers -RUN apt-get update && \ - echo "console-setup console-setup/variant select Latin1 and Latin5 - western Europe and Turkic languages" | debconf-set-selections && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends kali-linux-headless - -# Update pip again -> python3 -m pip install --upgrade pip causes crashes -# Install spacy deps - en_core_web_sm -# RUN python3 -m spacy download en_core_web_sm - -# Install docker inside of dev env -RUN apt-get update && \ - apt-get install -y apt-transport-https ca-certificates curl gnupg && \ - curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg && \ - echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian bullseye stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && \ - apt-get update && \ - apt-get install -y docker-ce docker-ce-cli containerd.io cloc - -# # Activate virtual environment by default -# RUN echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc - -RUN cargo install --git https://github.com/asciinema/agg - -# Remove system's python3 libs to avoid conflicts -RUN apt-get remove -y python3-jsonschema && \ - apt-get remove -y python3-numpy && \ - apt-get autoremove -y diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json deleted file mode 100644 index 22e1192e..00000000 --- a/.devcontainer/devcontainer.json +++ /dev/null @@ -1,86 +0,0 @@ -// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at: -// https://github.com/microsoft/vscode-dev-containers/tree/v0.241.1/containers/python-3 -{ - "name": "cai_devenv", - - // "build": { - // "dockerfile": "Dockerfile", - // "context": "..", - // "args": { - // // Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6 - // // Append -bullseye or -buster to pin to an OS version. - // // Use -bullseye variants on local on arm64/Apple Silicon. - // "VARIANT": "3.10-bullseye", - // // Options - // "NODE_VERSION": "lts/*" - // } - // }, - - "dockerComposeFile": ["./docker-compose.yml"], - "service": "devenv", - // "shutdownAction": "none", // don't shut down container when vscode is closed - "workspaceFolder": "/workspace", - - // Configure tool-specific properties. - "customizations": { - // Configure properties specific to VS Code. - "vscode": { - // Set *default* container specific settings.json values on container create. - "settings": { - "python.defaultInterpreterPath": "/usr/local/bin/python", - "python.linting.enabled": true, - "python.linting.pylintEnabled": true, - "python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8", - "python.formatting.blackPath": "/usr/local/py-utils/bin/black", - "python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf", - "python.linting.banditPath": "/usr/local/py-utils/bin/bandit", - "python.linting.flake8Path": "/usr/local/py-utils/bin/flake8", - "python.linting.mypyPath": "/usr/local/py-utils/bin/mypy", - "python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle", - "python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", - "python.linting.pylintPath": "/usr/local/py-utils/bin/pylint" - }, - - // Add the IDs of extensions you want installed when the container is created. - "extensions": [ - "ms-python.python", - "ms-toolsai.jupyter-renderers", - "ms-toolsai.jupyter", - "ms-python.vscode-pylance", - "ms-azuretools.vscode-docker", - "ms-python.debugpy", - "ms-python.black-formatter", - "MS-vsliveshare.vsliveshare" - ] - } - }, - - // Use 'forwardPorts' to make a list of ports inside the container available locally. - // "forwardPorts": [], - - // Use 'preCreateCommand' to run commands before the container is created. - "preCreateCommand": "docker network prune -f && docker container prune -f", - - // Use 'postCreateCommand' to run commands after the container is created. - // "postCreateCommand": "pip3 install --user -r requirements.txt", // doing it instead in - // the dockerfile to cache - // requirements in container - // Install cai via "pip3 install -e ." - // Extensions, NOTE: only one postCreateCommand is allowed, so we need to install all extensions here - "postCreateCommand": "pip3 install -e /workspace/", - - // MSF setup, and RAG setup - "postStartCommand": [ - "nohup", "msfrpcd", "-P", "cai", "&", - "&&", - "python3", "cai/ins/rag/agent_helper.py" - ], - // // Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root. - // "remoteUser": "vscode" - - "runArgs": [ - "--privileged", - "-e", "DISPLAY=host.docker.internal:0", - "-v", "/tmp/.X11-unix:/tmp/.X11-unix" - ] -} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml deleted file mode 100644 index 92b14fb9..00000000 --- a/.devcontainer/docker-compose.yml +++ /dev/null @@ -1,78 +0,0 @@ -version: '3' - -################# -# SERVICES -################# -services: - # Developer environment - devenv: - build: - context: .. - dockerfile: .devcontainer/Dockerfile - # platform: linux/amd64 # TODO: uncomment this when we have a x86_64 machine - volumes: - # Mount the root folder that contains .git - - ..:/workspace:cached - # Mount Docker socket to enable Docker-in-Docker - - /var/run/docker.sock:/var/run/docker.sock - command: /bin/sh -c "while sleep 10; do :; done" - networks: - cainet: - ipv4_address: 192.168.2.5 - cap_add: - - NET_ADMIN - - SYS_ADMIN - privileged: true - - # # Vulnhub vulnerable machine (Hackable II) SUGGESTION: Comment this and use PentestPerf - # hackableii: - # image: vmayoral/vulnhub:hackableii - # command: | - # /bin/bash -c " - # rm -rf /var/lock && \ - # mkdir -p /var/lock && \ - # chmod 755 /var/lock && \ - # /etc/init.d/apache2 start && \ - # /etc/init.d/ssh start && \ - # /etc/init.d/runproftpd.sh && \ - # /etc/init.d/php7.0-fpm start && \ - # while sleep 1; do :; done - # " - # networks: - # cainet: - # ipv4_address: 192.168.2.11 - # mac_address: 08:00:27:85:55:86 - - # # Vulnhub vulnerable machine (Bob) SUGGESTION: Comment this and use PentestPerf - # bob: - # image: vmayoral/vulnhub:bob - # command: | - # /bin/bash -c "rm -r /var/lock; mkdir -p /var/lock; chmod 755 /var/lock; /etc/init.d/apache2 start; /etc/init.d/ssh start; while sleep 10; do :; done" - # ports: # map port in the container to the host systems - # - "8080:80" - # networks: - # cainet: - # ipv4_address: 192.168.2.12 - # mac_address: 08:00:27:cb:07:d4 - - # Qdrant vector database - qdrant: - image: qdrant/qdrant - ports: - - "6333:6333" - - "6334:6334" - volumes: - - ./qdrant_storage:/qdrant/storage:z - networks: - cainet: - ipv4_address: 192.168.2.13 - -################# -# NETWORKS -################# -networks: - cainet: - ipam: - driver: default - config: - - subnet: 192.168.2.0/24 diff --git a/.devcontainer/qdrant_storage/aliases/data.json b/.devcontainer/qdrant_storage/aliases/data.json deleted file mode 100644 index 9e26dfee..00000000 --- a/.devcontainer/qdrant_storage/aliases/data.json +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/.devcontainer/qdrant_storage/raft_state.json b/.devcontainer/qdrant_storage/raft_state.json deleted file mode 100644 index 0ee33f67..00000000 --- a/.devcontainer/qdrant_storage/raft_state.json +++ /dev/null @@ -1 +0,0 @@ -{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6138497487128594],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6138497487128594,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6138497487128594} \ No newline at end of file diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt deleted file mode 100644 index 2409daf2..00000000 --- a/.devcontainer/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -pymetasploit3 -networkx==2.5 -requests -wasabi -xmltodict -pydot==1.4.2 diff --git a/src/cai/cli.py b/src/cai/cli.py index 1bba0c6b..a73ffc0d 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -137,8 +137,6 @@ from cai.util import ( stop_active_timer, setup_ctf, check_flag, - build_docker_container, - remove_container ) # CAI REPL imports @@ -163,9 +161,8 @@ if is_pentestperf_available() and os.getenv('CTF_NAME', None): ctf, messages_ctf = setup_ctf() ctf_global = ctf if os.getenv('CTF_INSIDE', 'True').lower() == 'false': - # instanciate dev container - container_id = build_docker_container("./.devcontainer") - os.environ['CAI_ACTIVE_CONTAINER'] = container_id + container_id = "" + os.environ['CAI_ACTIVE_CONTAINER'] = container_id # Load environment variables from .env file load_dotenv() @@ -271,10 +268,6 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), ctf, messages_ctf = setup_ctf() ctf_global = ctf previous_ctf_name = os.getenv('CTF_NAME', None) - if os.getenv('CTF_INSIDE', 'True') == 'False': - # instanciate dev container - container_id = build_docker_container("./.devcontainer") - os.environ['CAI_ACTIVE_CONTAINER'] = container_id try: # Start measuring user idle time start_idle_timer() @@ -431,10 +424,6 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), if (is_pentestperf_available() and os.getenv('CTF_NAME', None)): ctf.stop_ctf() - if os.getenv('CTF_INSIDE', 'True').lower() == 'false': - remove_container(os.environ['CAI_ACTIVE_CONTAINER']) - return False - except Exception: pass diff --git a/src/cai/util.py b/src/cai/util.py index 9a769037..3366c873 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -2857,40 +2857,3 @@ def setup_ctf(): bg="blue")) return ctf, messages - - -def build_docker_container(devcontainer_path: str) -> str: - dockerfile_path = os.path.join(devcontainer_path, "Dockerfile") - image_name = "cai-container3" - container_name = "cai-container3" - - # Try to remove existing container with the same name - try: - subprocess.run(["docker", "rm", "-f", container_name], check=True) - except subprocess.CalledProcessError: - pass # If it doesn't exist, continue normally - - # Build the image - subprocess.run([ - "docker", "build", - "-t", image_name, - "-f", dockerfile_path, - devcontainer_path - ], check=True) - - # Run the container - result = subprocess.run([ - "docker", "run", "-d", "--name", container_name, image_name, "tail", "-f", "/dev/null" - ], stdout=subprocess.PIPE, check=True, text=True) - - container_id = result.stdout.strip() - return container_id - -def remove_container(name_or_id: str): - try: - # Force removal (stops if running, then removes) - subprocess.run(["docker", "rm", "-f", name_or_id], check=True) - print(f"Container '{name_or_id}' removed successfully.") - except subprocess.CalledProcessError as e: - print(f"Error removing container {name_or_id}:", e) - From c0cd5025796c899131362b39d8b12d43a7122dc9 Mon Sep 17 00:00:00 2001 From: luijait Date: Wed, 21 May 2025 15:59:12 +0200 Subject: [PATCH 14/25] Pricing streaming --- .../agents/models/openai_chatcompletions.py | 69 +++++++++++++++++-- src/cai/util.py | 62 ++++++++++++++--- 2 files changed, 116 insertions(+), 15 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 95e4be9f..73ac1c76 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -705,6 +705,15 @@ class OpenAIChatCompletionsModel(Model): # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) + # Pre-check price limit using estimated input tokens and a conservative estimate for output + # This prevents starting a stream that would immediately exceed the price limit + if hasattr(COST_TRACKER, "check_price_limit"): + # Use a conservative estimate for output tokens (roughly equal to input) + estimated_cost = calculate_model_cost(str(self.model), + estimated_input_tokens, + estimated_input_tokens) # Conservative estimate + COST_TRACKER.check_price_limit(estimated_cost) + response, stream = await self._fetch_response( system_instructions, input, @@ -807,12 +816,25 @@ class OpenAIChatCompletionsModel(Model): # Update streaming display if enabled - always do this for text content if streaming_context: - # Create token stats to pass with each update + # Calculate cost for current interaction + current_cost = calculate_model_cost(str(self.model), estimated_input_tokens, estimated_output_tokens) + + # Check price limit before updating (conservative check) + if hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0: + COST_TRACKER.check_price_limit(current_cost) + + # Update session total cost for real-time display + # This is a temporary estimate during streaming that will be properly updated at the end + estimated_session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + + # Create token stats with both current interaction cost and updated total cost token_stats = { 'input_tokens': estimated_input_tokens, 'output_tokens': estimated_output_tokens, - 'cost': calculate_model_cost(str(self.model), estimated_input_tokens, estimated_output_tokens) + 'cost': current_cost, + 'total_cost': estimated_session_total + current_cost } + update_agent_streaming_content(streaming_context, content, token_stats) # More accurate token counting for text content @@ -820,6 +842,32 @@ class OpenAIChatCompletionsModel(Model): token_count, _ = count_tokens_with_tiktoken(output_text) estimated_output_tokens = token_count + # Periodically check price limit during streaming + # This allows early termination if price limit is reached mid-stream + if estimated_output_tokens > 0 and estimated_output_tokens % 50 == 0: # Check every ~50 tokens + current_estimated_cost = calculate_model_cost( + str(self.model), estimated_input_tokens, estimated_output_tokens) + + # Check price limit + if hasattr(COST_TRACKER, "check_price_limit"): + COST_TRACKER.check_price_limit(current_estimated_cost) + + # Update the COST_TRACKER with the running cost for accurate display + if hasattr(COST_TRACKER, "interaction_cost"): + COST_TRACKER.interaction_cost = current_estimated_cost + + # Also update streaming context if available for live display + if streaming_context: + # Get the latest session total cost + session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + current_estimated_cost + updated_token_stats = { + 'input_tokens': estimated_input_tokens, + 'output_tokens': estimated_output_tokens, + 'cost': current_estimated_cost, + 'total_cost': session_total + } + update_agent_streaming_content(streaming_context, "", updated_token_stats) + if not state.text_content_index_and_output: # Initialize a content tracker for streaming text state.text_content_index_and_output = ( @@ -1321,8 +1369,21 @@ class OpenAIChatCompletionsModel(Model): total_cost = float(total_cost if total_cost is not None else 0.0) # Update the global COST_TRACKER with the cost of this specific interaction - if hasattr(COST_TRACKER, "add_interaction_cost") and interaction_cost > 0.0: - COST_TRACKER.add_interaction_cost(interaction_cost) + # and check price limit for streaming mode (similar to non-streaming mode) + if interaction_cost > 0.0: + # Check price limit before adding the new cost + if hasattr(COST_TRACKER, "check_price_limit"): + COST_TRACKER.check_price_limit(interaction_cost) + + # Now add the cost to session total + if hasattr(COST_TRACKER, "update_session_cost"): + COST_TRACKER.update_session_cost(interaction_cost) + elif hasattr(COST_TRACKER, "add_interaction_cost"): + COST_TRACKER.add_interaction_cost(interaction_cost) + + # Ensure the total cost includes the session total for display + if hasattr(COST_TRACKER, "session_total_cost"): + total_cost = COST_TRACKER.session_total_cost # Store the total cost for future recording self.total_cost = total_cost diff --git a/src/cai/util.py b/src/cai/util.py index 938ecde8..e161204a 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -256,6 +256,20 @@ class CostTracker: old_total = self.session_total_cost self.session_total_cost += new_cost + def add_interaction_cost(self, new_cost: float) -> None: + """ + Add an interaction cost to the session total and check price limit. + This is a convenience method that combines check_price_limit and update_session_cost. + """ + # Check price limit first + self.check_price_limit(new_cost) + + # Then update the session cost + self.session_total_cost += new_cost + + # Update the last interaction cost for tracking + self.last_interaction_cost = new_cost + def log_final_cost(self) -> None: """Display final cost information at exit""" # Skip displaying cost if already shown in the session summary @@ -1522,16 +1536,26 @@ def update_agent_streaming_content(context, text_delta, token_stats=None): return False try: - # Parse the text_delta to get just the content if needed - parsed_delta = parse_message_content(text_delta) - - # Skip empty updates to avoid showing an empty panel - if not parsed_delta or parsed_delta.strip() == "": + # Only parse and add text if we have actual content to add + # Skip when text_delta is empty and we're just updating token stats + if text_delta: + # Parse the text_delta to get just the content if needed + parsed_delta = parse_message_content(text_delta) + + # Skip empty updates to avoid showing an empty panel + if not parsed_delta or parsed_delta.strip() == "": + # Update token stats if provided + if token_stats: + # Just update the footer, not the content + pass + else: + # Add the parsed text to the content + context["content"].append(parsed_delta) + # If no text_delta but we have token_stats, just update stats + elif not token_stats: + # No text and no stats - nothing to update return True - # Add the parsed text to the content - context["content"].append(parsed_delta) - # Update the footer with token stats if provided if token_stats: # Create token stats display @@ -1547,13 +1571,24 @@ def update_agent_streaming_content(context, text_delta, token_stats=None): # Add token stats input_tokens = token_stats.get('input_tokens', 0) output_tokens = token_stats.get('output_tokens', 0) - total_cost = token_stats.get('cost', 0.0) + interaction_cost = token_stats.get('cost', 0.0) + + # Get session total cost - either from token_stats or directly from COST_TRACKER + session_total_cost = token_stats.get('total_cost', 0.0) + if session_total_cost == 0.0 and hasattr(COST_TRACKER, 'session_total_cost'): + session_total_cost = COST_TRACKER.session_total_cost if input_tokens > 0: footer_stats.append(" | ", style="dim") footer_stats.append(f"I:{input_tokens} O:{output_tokens}", style="green") - if total_cost > 0: - footer_stats.append(f" (${total_cost:.4f})", style="bold cyan") + + # Show both interaction cost and total session cost + if interaction_cost > 0: + footer_stats.append(f" (${interaction_cost:.4f})", style="bold cyan") + + # Add the total cost information on the same line + footer_stats.append(" | Session: ", style="dim") + footer_stats.append(f"${session_total_cost:.4f}", style="bold magenta") # Add context usage indicator model_name = context.get("model", os.environ.get('CAI_MODEL', 'qwen2.5:14b')) @@ -1684,6 +1719,11 @@ def finish_agent_streaming(context, final_stats=None): compact_tokens.append(f"I:{interaction_input_tokens} O:{interaction_output_tokens} ", style="green") compact_tokens.append(f"(${interaction_cost:.4f}) ", style="bold cyan") + # Include the total session cost + session_total_cost = COST_TRACKER.session_total_cost if hasattr(COST_TRACKER, 'session_total_cost') else total_cost + compact_tokens.append(" | Session: ", style="dim") + compact_tokens.append(f"${session_total_cost:.4f}", style="bold magenta") + # Añadir un indicador de uso de contexto context_pct = interaction_input_tokens / get_model_input_tokens(model_name) * 100 if context_pct < 50: From ec8179083d0ef849bb3ecf115a0d49497431e83a Mon Sep 17 00:00:00 2001 From: luijait Date: Wed, 21 May 2025 16:24:52 +0200 Subject: [PATCH 15/25] local model fix --- .../agents/models/openai_chatcompletions.py | 84 +++++++++++++++---- src/cai/util.py | 65 +++++++++++++- 2 files changed, 134 insertions(+), 15 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 73ac1c76..21ffad1b 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -816,23 +816,39 @@ class OpenAIChatCompletionsModel(Model): # Update streaming display if enabled - always do this for text content if streaming_context: - # Calculate cost for current interaction + # Check if this is a local model and should have zero cost + model_str = str(self.model).lower() + is_local_model = ( + "ollama" in model_str or + "qwen" in model_str or + "llama" in model_str or + "mistral" in model_str or + ":" in model_str or + self.is_ollama + ) + + # Calculate cost for current interaction (will be 0 for local models) current_cost = calculate_model_cost(str(self.model), estimated_input_tokens, estimated_output_tokens) - # Check price limit before updating (conservative check) - if hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0: + # Check price limit only for non-local models + if not is_local_model and hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0: COST_TRACKER.check_price_limit(current_cost) # Update session total cost for real-time display # This is a temporary estimate during streaming that will be properly updated at the end estimated_session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + # For local models, don't add to the total cost + display_total_cost = estimated_session_total + if not is_local_model: + display_total_cost += current_cost + # Create token stats with both current interaction cost and updated total cost token_stats = { 'input_tokens': estimated_input_tokens, 'output_tokens': estimated_output_tokens, 'cost': current_cost, - 'total_cost': estimated_session_total + current_cost + 'total_cost': display_total_cost } update_agent_streaming_content(streaming_context, content, token_stats) @@ -845,11 +861,26 @@ class OpenAIChatCompletionsModel(Model): # Periodically check price limit during streaming # This allows early termination if price limit is reached mid-stream if estimated_output_tokens > 0 and estimated_output_tokens % 50 == 0: # Check every ~50 tokens - current_estimated_cost = calculate_model_cost( - str(self.model), estimated_input_tokens, estimated_output_tokens) + # Check if this is a local model (Ollama, Qwen, etc.) that should have zero cost + model_str = str(self.model).lower() + is_local_model = ( + "ollama" in model_str or + "qwen" in model_str or + "llama" in model_str or + "mistral" in model_str or + ":" in model_str or + self.is_ollama + ) - # Check price limit - if hasattr(COST_TRACKER, "check_price_limit"): + # For local models, cost should always be zero + if is_local_model: + current_estimated_cost = 0.0 + else: + current_estimated_cost = calculate_model_cost( + str(self.model), estimated_input_tokens, estimated_output_tokens) + + # Check price limit only for non-local models + if not is_local_model and hasattr(COST_TRACKER, "check_price_limit"): COST_TRACKER.check_price_limit(current_estimated_cost) # Update the COST_TRACKER with the running cost for accurate display @@ -858,8 +889,12 @@ class OpenAIChatCompletionsModel(Model): # Also update streaming context if available for live display if streaming_context: - # Get the latest session total cost - session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + current_estimated_cost + # For local models, don't add to the session total + if is_local_model: + session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + else: + session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + current_estimated_cost + updated_token_stats = { 'input_tokens': estimated_input_tokens, 'output_tokens': estimated_output_tokens, @@ -1359,10 +1394,31 @@ class OpenAIChatCompletionsModel(Model): total_input = getattr(self, 'total_input_tokens', 0) total_output = getattr(self, 'total_output_tokens', 0) - # Calculate costs using the same token counts - ensure model is a string + # Check if this is a local model and should have zero cost + model_str = str(self.model).lower() + is_local_model = ( + "ollama" in model_str or + "qwen" in model_str or + "llama" in model_str or + "mistral" in model_str or + ":" in model_str or + self.is_ollama + ) + + # Calculate costs - use zero for local models model_name = str(self.model) - interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output) - total_cost = calculate_model_cost(model_name, total_input, total_output) + if is_local_model: + # For local models, cost is always zero + interaction_cost = 0.0 + total_cost = getattr(COST_TRACKER, 'session_total_cost', 0.0) # Keep existing total + + # Ensure the cost tracking system knows this is a free model + if hasattr(COST_TRACKER, "reset_cost_for_local_model"): + COST_TRACKER.reset_cost_for_local_model(model_name) + else: + # For paid models, calculate as normal + interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output) + total_cost = calculate_model_cost(model_name, total_input, total_output) # Explicit conversion to float with fallback to ensure they're never None or 0 interaction_cost = float(interaction_cost if interaction_cost is not None else 0.0) @@ -1370,7 +1426,7 @@ class OpenAIChatCompletionsModel(Model): # Update the global COST_TRACKER with the cost of this specific interaction # and check price limit for streaming mode (similar to non-streaming mode) - if interaction_cost > 0.0: + if not is_local_model and interaction_cost > 0.0: # Check price limit before adding the new cost if hasattr(COST_TRACKER, "check_price_limit"): COST_TRACKER.check_price_limit(interaction_cost) diff --git a/src/cai/util.py b/src/cai/util.py index e161204a..3521e23e 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -261,6 +261,11 @@ class CostTracker: Add an interaction cost to the session total and check price limit. This is a convenience method that combines check_price_limit and update_session_cost. """ + # Skip updating costs if the cost is zero (common with local models) + if new_cost <= 0: + self.last_interaction_cost = 0.0 + return + # Check price limit first self.check_price_limit(new_cost) @@ -270,6 +275,31 @@ class CostTracker: # Update the last interaction cost for tracking self.last_interaction_cost = new_cost + def reset_cost_for_local_model(self, model_name: str) -> bool: + """ + Reset interaction cost tracking when switching to a local model. + Returns True if the model was identified as local and cost was reset. + """ + # Check if this is a local/free model + model_str = model_name.lower() + is_local_model = ( + "ollama" in model_str or + "qwen" in model_str or + "llama" in model_str or + "mistral" in model_str or + ":" in model_str or + (os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false') + ) + + if is_local_model: + # Reset the current interaction costs but keep total session costs + self.interaction_cost = 0.0 + self.last_interaction_cost = 0.0 + # Don't reset session_total_cost as that includes previous paid models + return True + + return False + def log_final_cost(self) -> None: """Display final cost information at exit""" # Skip displaying cost if already shown in the session summary @@ -281,7 +311,25 @@ class CostTracker: # Use the centralized function to standardize model names model_name = get_model_name(model_name) - # Check cache first + # Check if using Ollama or local model + model_str = model_name.lower() + is_local_model = ( + "ollama" in model_str or + "qwen" in model_str or + "llama" in model_str or + "mistral" in model_str or + ":" in model_str or # Ollama uses formats like qwen2.5:7b + (os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false') + ) + + # For local models, always return zero cost + if is_local_model: + # Set and cache zero cost for local models + free_pricing = (0.0, 0.0) + self.model_pricing_cache[model_name] = free_pricing + return free_pricing + + # Check cache for non-local models if model_name in self.model_pricing_cache: return self.model_pricing_cache[model_name] @@ -335,6 +383,21 @@ class CostTracker: # Standardize model name using the central function model_name = get_model_name(model) + # Check if this is a local model (always free) first + model_str = model_name.lower() + is_local_model = ( + "ollama" in model_str or + "qwen" in model_str or + "llama" in model_str or + "mistral" in model_str or + ":" in model_str or + (os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false') + ) + + # For local models, always return zero cost + if is_local_model: + return 0.0 + # Generate a cache key cache_key = f"{model_name}_{input_tokens}_{output_tokens}" From a421eb7fdf482bacef86ce8f3a5fe460e5ebf44c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Wed, 21 May 2025 17:23:55 +0200 Subject: [PATCH 16/25] Fixup replay so that it doesn't crash if no content MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- tools/replay.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/replay.py b/tools/replay.py index b681ca88..75d20746 100644 --- a/tools/replay.py +++ b/tools/replay.py @@ -154,7 +154,8 @@ def replay_conversation(messages: List[Dict], replay_delay: float = 0.5, usage: time.sleep(replay_delay) role = message.get("role", "") - content = message.get("content", "").strip() + content = message.get("content") + content = str(content).strip() if content is not None else "" sender = message.get("sender", role) model = message.get("model", file_model) From b4571bb4b7e29df47a11230b954f709b07b4692a Mon Sep 17 00:00:00 2001 From: luijait Date: Wed, 21 May 2025 22:32:11 +0200 Subject: [PATCH 17/25] Support pricing in alias0 --- src/cai/util.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cai/util.py b/src/cai/util.py index 3521e23e..b7223fad 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -283,6 +283,7 @@ class CostTracker: # Check if this is a local/free model model_str = model_name.lower() is_local_model = ( + "alias" not in model_str and "ollama" in model_str or "qwen" in model_str or "llama" in model_str or @@ -386,6 +387,7 @@ class CostTracker: # Check if this is a local model (always free) first model_str = model_name.lower() is_local_model = ( + "alias" not in model_str and "ollama" in model_str or "qwen" in model_str or "llama" in model_str or From 15efd5809ab3ad1bced8a7a5002eb2fc6ea52186 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 22 May 2025 08:58:20 +0200 Subject: [PATCH 18/25] Robust pricing --- .../agents/models/openai_chatcompletions.py | 71 +++++++++++++++++-- 1 file changed, 67 insertions(+), 4 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 21ffad1b..4a5448a4 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -348,6 +348,21 @@ class OpenAIChatCompletionsModel(Model): # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) + # Pre-check price limit using estimated input tokens and a conservative estimate for output + # This prevents starting a request that would immediately exceed the price limit + if hasattr(COST_TRACKER, "check_price_limit"): + # Use a conservative estimate for output tokens (roughly equal to input) + estimated_cost = calculate_model_cost(str(self.model), + estimated_input_tokens, + estimated_input_tokens) # Conservative estimate + try: + COST_TRACKER.check_price_limit(estimated_cost) + except Exception as e: + # Stop active timer and start idle timer before re-raising the exception + stop_active_timer() + start_idle_timer() + raise + try: response = await self._fetch_response( system_instructions, @@ -712,7 +727,19 @@ class OpenAIChatCompletionsModel(Model): estimated_cost = calculate_model_cost(str(self.model), estimated_input_tokens, estimated_input_tokens) # Conservative estimate - COST_TRACKER.check_price_limit(estimated_cost) + try: + COST_TRACKER.check_price_limit(estimated_cost) + except Exception as e: + # Ensure streaming context is cleaned up in case of errors + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + except Exception: + pass + # Stop active timer and start idle timer before re-raising the exception + stop_active_timer() + start_idle_timer() + raise response, stream = await self._fetch_response( system_instructions, @@ -832,7 +859,19 @@ class OpenAIChatCompletionsModel(Model): # Check price limit only for non-local models if not is_local_model and hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0: - COST_TRACKER.check_price_limit(current_cost) + try: + COST_TRACKER.check_price_limit(current_cost) + except Exception as e: + # Ensure streaming context is cleaned up + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + except Exception: + pass + # Stop timers and re-raise the exception + stop_active_timer() + start_idle_timer() + raise # Update session total cost for real-time display # This is a temporary estimate during streaming that will be properly updated at the end @@ -881,7 +920,19 @@ class OpenAIChatCompletionsModel(Model): # Check price limit only for non-local models if not is_local_model and hasattr(COST_TRACKER, "check_price_limit"): - COST_TRACKER.check_price_limit(current_estimated_cost) + try: + COST_TRACKER.check_price_limit(current_estimated_cost) + except Exception as e: + # Ensure streaming context is cleaned up + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + except Exception: + pass + # Stop timers and re-raise the exception + stop_active_timer() + start_idle_timer() + raise # Update the COST_TRACKER with the running cost for accurate display if hasattr(COST_TRACKER, "interaction_cost"): @@ -1429,7 +1480,19 @@ class OpenAIChatCompletionsModel(Model): if not is_local_model and interaction_cost > 0.0: # Check price limit before adding the new cost if hasattr(COST_TRACKER, "check_price_limit"): - COST_TRACKER.check_price_limit(interaction_cost) + try: + COST_TRACKER.check_price_limit(interaction_cost) + except Exception as e: + # Ensure streaming context is cleaned up + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + except Exception: + pass + # Stop timers and re-raise the exception + stop_active_timer() + start_idle_timer() + raise # Now add the cost to session total if hasattr(COST_TRACKER, "update_session_cost"): From 45e2214f16d232721fa44dcb5531e8207695fcd7 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:15:20 +0200 Subject: [PATCH 19/25] add info --- .devcontainer/Dockerfile | 108 ++++++++++++++++++ .devcontainer/devcontainer.json | 86 ++++++++++++++ .devcontainer/docker-compose.yml | 78 +++++++++++++ .../qdrant_storage/aliases/data.json | 1 + .devcontainer/qdrant_storage/raft_state.json | 1 + .devcontainer/requirements.txt | 6 + src/cai/cli.py | 2 +- 7 files changed, 281 insertions(+), 1 deletion(-) create mode 100644 .devcontainer/Dockerfile create mode 100644 .devcontainer/devcontainer.json create mode 100644 .devcontainer/docker-compose.yml create mode 100644 .devcontainer/qdrant_storage/aliases/data.json create mode 100644 .devcontainer/qdrant_storage/raft_state.json create mode 100644 .devcontainer/requirements.txt diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 00000000..43fe72ae --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,108 @@ +FROM kalilinux/kali-rolling + +# Set environment variable to non-interactive +ENV DEBIAN_FRONTEND=noninteractive + +# [Optional] Uncomment this section to install additional OS packages. +RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ + && apt-get -y install --no-install-recommends \ + net-tools python3 python3-pip python3-venv \ + curl gnupg nmap iputils-ping ssh git \ + graphviz pkg-config libhdf5-dev \ + build-essential python3-dev \ + asciinema dnsutils \ + apt-transport-https ca-certificates \ + wget dirb gobuster whatweb gfortran \ + cmake libopenblas-dev sshpass seclists \ + golang-go + +# # Install Metasploit - path 1 +# RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > /tmp/msfinstall \ +# && chmod 755 /tmp/msfinstall \ +# && /tmp/msfinstall + +# # Install Metasploit - path 2 +# RUN curl https://apt.metasploit.com/metasploit-framework.gpg.key | apt-key add - +# RUN echo "deb https://apt.metasploit.com/ lucid main" | tee /etc/apt/sources.list.d/metasploit-framework.list +# RUN apt-get update && apt-get install metasploit-framework + +# Install Metasploit - path 3 +RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstall && chmod 755 msfinstall && ./msfinstall + +# Update Kali Linux repositories +RUN apt-get update \ + && apt-get upgrade -y + +# # Create a virtual environment +# ENV VIRTUAL_ENV=/opt/venv +# RUN python3 -m venv $VIRTUAL_ENV +# ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +# # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image. +# # alternatively, this also helps avoid having to pull requirements from the internet every single time +# COPY requirements.txt /tmp/pip-tmp/ +# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ +# && rm -rf /tmp/pip-tmp + +# FIX for: +# +# Python packages system-wide, try apt install +# python3-xyz, where xyz is the package you are trying to +# install. +# +# see https://github.com/raspiblitz/raspiblitz/issues/4170 +RUN mkdir -p /root/.pip && touch /root/.pip/pip.conf +RUN echo "[global]" > /root/.pip/pip.conf && \ + echo "break-system-packages = true" >> /root/.pip/pip.conf + +RUN pip3 install --upgrade setuptools +RUN pip3 install pandas \ + opentelemetry-sdk opentelemetry-exporter-otlp \ + mem0ai PyPDF2 sentence_transformers tf-keras \ + pytest-repeat parameterized pytest-rerunfailures \ + pytest-clarity tiktoken blobfile pytest-timeout \ + invoke fabric + +# Install rust compiler - dependency of outlines-core +RUN apt-get update && \ + apt-get install -y curl && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \ + . $HOME/.cargo/env && \ + rustup default stable && \ + rustup component add rustfmt clippy && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +ENV PATH="/root/.cargo/bin:${PATH}" +RUN wget https://raw.githubusercontent.com/adfoster-r7/metasploit-info/master/info/module_metadata.json -O /tmp/module_metadata.json + +# NOTE: align dev env in Linux +# work +# RUN echo "192.168.2.1 host.docker.internal" >> /etc/hosts +# Update and upgrade packages +# Preconfigure console-setup and install kali-linux-headless fixers +RUN apt-get update && \ + echo "console-setup console-setup/variant select Latin1 and Latin5 - western Europe and Turkic languages" | debconf-set-selections && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends kali-linux-headless + +# Update pip again -> python3 -m pip install --upgrade pip causes crashes +# Install spacy deps - en_core_web_sm +# RUN python3 -m spacy download en_core_web_sm + +# Install docker inside of dev env +RUN apt-get update && \ + apt-get install -y apt-transport-https ca-certificates curl gnupg && \ + curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg && \ + echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian bullseye stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && \ + apt-get update && \ + apt-get install -y docker-ce docker-ce-cli containerd.io cloc + +# # Activate virtual environment by default +# RUN echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc + +RUN cargo install --git https://github.com/asciinema/agg + +# Remove system's python3 libs to avoid conflicts +RUN apt-get remove -y python3-jsonschema && \ + apt-get remove -y python3-numpy && \ + apt-get autoremove -y diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 00000000..22e1192e --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,86 @@ +// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at: +// https://github.com/microsoft/vscode-dev-containers/tree/v0.241.1/containers/python-3 +{ + "name": "cai_devenv", + + // "build": { + // "dockerfile": "Dockerfile", + // "context": "..", + // "args": { + // // Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6 + // // Append -bullseye or -buster to pin to an OS version. + // // Use -bullseye variants on local on arm64/Apple Silicon. + // "VARIANT": "3.10-bullseye", + // // Options + // "NODE_VERSION": "lts/*" + // } + // }, + + "dockerComposeFile": ["./docker-compose.yml"], + "service": "devenv", + // "shutdownAction": "none", // don't shut down container when vscode is closed + "workspaceFolder": "/workspace", + + // Configure tool-specific properties. + "customizations": { + // Configure properties specific to VS Code. + "vscode": { + // Set *default* container specific settings.json values on container create. + "settings": { + "python.defaultInterpreterPath": "/usr/local/bin/python", + "python.linting.enabled": true, + "python.linting.pylintEnabled": true, + "python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8", + "python.formatting.blackPath": "/usr/local/py-utils/bin/black", + "python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf", + "python.linting.banditPath": "/usr/local/py-utils/bin/bandit", + "python.linting.flake8Path": "/usr/local/py-utils/bin/flake8", + "python.linting.mypyPath": "/usr/local/py-utils/bin/mypy", + "python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle", + "python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", + "python.linting.pylintPath": "/usr/local/py-utils/bin/pylint" + }, + + // Add the IDs of extensions you want installed when the container is created. + "extensions": [ + "ms-python.python", + "ms-toolsai.jupyter-renderers", + "ms-toolsai.jupyter", + "ms-python.vscode-pylance", + "ms-azuretools.vscode-docker", + "ms-python.debugpy", + "ms-python.black-formatter", + "MS-vsliveshare.vsliveshare" + ] + } + }, + + // Use 'forwardPorts' to make a list of ports inside the container available locally. + // "forwardPorts": [], + + // Use 'preCreateCommand' to run commands before the container is created. + "preCreateCommand": "docker network prune -f && docker container prune -f", + + // Use 'postCreateCommand' to run commands after the container is created. + // "postCreateCommand": "pip3 install --user -r requirements.txt", // doing it instead in + // the dockerfile to cache + // requirements in container + // Install cai via "pip3 install -e ." + // Extensions, NOTE: only one postCreateCommand is allowed, so we need to install all extensions here + "postCreateCommand": "pip3 install -e /workspace/", + + // MSF setup, and RAG setup + "postStartCommand": [ + "nohup", "msfrpcd", "-P", "cai", "&", + "&&", + "python3", "cai/ins/rag/agent_helper.py" + ], + // // Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root. + // "remoteUser": "vscode" + + "runArgs": [ + "--privileged", + "-e", "DISPLAY=host.docker.internal:0", + "-v", "/tmp/.X11-unix:/tmp/.X11-unix" + ] +} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml new file mode 100644 index 00000000..92b14fb9 --- /dev/null +++ b/.devcontainer/docker-compose.yml @@ -0,0 +1,78 @@ +version: '3' + +################# +# SERVICES +################# +services: + # Developer environment + devenv: + build: + context: .. + dockerfile: .devcontainer/Dockerfile + # platform: linux/amd64 # TODO: uncomment this when we have a x86_64 machine + volumes: + # Mount the root folder that contains .git + - ..:/workspace:cached + # Mount Docker socket to enable Docker-in-Docker + - /var/run/docker.sock:/var/run/docker.sock + command: /bin/sh -c "while sleep 10; do :; done" + networks: + cainet: + ipv4_address: 192.168.2.5 + cap_add: + - NET_ADMIN + - SYS_ADMIN + privileged: true + + # # Vulnhub vulnerable machine (Hackable II) SUGGESTION: Comment this and use PentestPerf + # hackableii: + # image: vmayoral/vulnhub:hackableii + # command: | + # /bin/bash -c " + # rm -rf /var/lock && \ + # mkdir -p /var/lock && \ + # chmod 755 /var/lock && \ + # /etc/init.d/apache2 start && \ + # /etc/init.d/ssh start && \ + # /etc/init.d/runproftpd.sh && \ + # /etc/init.d/php7.0-fpm start && \ + # while sleep 1; do :; done + # " + # networks: + # cainet: + # ipv4_address: 192.168.2.11 + # mac_address: 08:00:27:85:55:86 + + # # Vulnhub vulnerable machine (Bob) SUGGESTION: Comment this and use PentestPerf + # bob: + # image: vmayoral/vulnhub:bob + # command: | + # /bin/bash -c "rm -r /var/lock; mkdir -p /var/lock; chmod 755 /var/lock; /etc/init.d/apache2 start; /etc/init.d/ssh start; while sleep 10; do :; done" + # ports: # map port in the container to the host systems + # - "8080:80" + # networks: + # cainet: + # ipv4_address: 192.168.2.12 + # mac_address: 08:00:27:cb:07:d4 + + # Qdrant vector database + qdrant: + image: qdrant/qdrant + ports: + - "6333:6333" + - "6334:6334" + volumes: + - ./qdrant_storage:/qdrant/storage:z + networks: + cainet: + ipv4_address: 192.168.2.13 + +################# +# NETWORKS +################# +networks: + cainet: + ipam: + driver: default + config: + - subnet: 192.168.2.0/24 diff --git a/.devcontainer/qdrant_storage/aliases/data.json b/.devcontainer/qdrant_storage/aliases/data.json new file mode 100644 index 00000000..9e26dfee --- /dev/null +++ b/.devcontainer/qdrant_storage/aliases/data.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/.devcontainer/qdrant_storage/raft_state.json b/.devcontainer/qdrant_storage/raft_state.json new file mode 100644 index 00000000..b7bcb749 --- /dev/null +++ b/.devcontainer/qdrant_storage/raft_state.json @@ -0,0 +1 @@ +{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6120741261401965],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6120741261401965,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6120741261401965} \ No newline at end of file diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt new file mode 100644 index 00000000..2409daf2 --- /dev/null +++ b/.devcontainer/requirements.txt @@ -0,0 +1,6 @@ +pymetasploit3 +networkx==2.5 +requests +wasabi +xmltodict +pydot==1.4.2 diff --git a/src/cai/cli.py b/src/cai/cli.py index 51255117..e59ff7f7 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -340,7 +340,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), get_toolbar_with_refresh, current_text ) + messages_ctf - else: + else: user_input = messages_ctf idle_time += time.time() - idle_start_time From 3a203cee5fea499b7a5a0238bdb76951d609b3a3 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:56:30 +0200 Subject: [PATCH 20/25] fix --- src/cai/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index e59ff7f7..51255117 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -340,7 +340,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), get_toolbar_with_refresh, current_text ) + messages_ctf - else: + else: user_input = messages_ctf idle_time += time.time() - idle_start_time From 521c33a6b11d646a91b8b4afdddb764be26a54e2 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 20 May 2025 16:05:20 +0200 Subject: [PATCH 21/25] cker virtualization for ctf_inside false --- .devcontainer/qdrant_storage/raft_state.json | 2 +- src/cai/util.py | 37 ++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/.devcontainer/qdrant_storage/raft_state.json b/.devcontainer/qdrant_storage/raft_state.json index b7bcb749..23b070d6 100644 --- a/.devcontainer/qdrant_storage/raft_state.json +++ b/.devcontainer/qdrant_storage/raft_state.json @@ -1 +1 @@ -{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6120741261401965],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6120741261401965,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6120741261401965} \ No newline at end of file +{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6120741261401965],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6120741261401965,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6120741261401965} diff --git a/src/cai/util.py b/src/cai/util.py index b7223fad..ff62608b 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -2997,3 +2997,40 @@ def setup_ctf(): bg="blue")) return ctf, messages + + +def build_docker_container(devcontainer_path: str) -> str: + dockerfile_path = os.path.join(devcontainer_path, "Dockerfile") + image_name = "cai-container3" + container_name = "cai-container3" + + # Try to remove existing container with the same name + try: + subprocess.run(["docker", "rm", "-f", container_name], check=True) + except subprocess.CalledProcessError: + pass # If it doesn't exist, continue normally + + # Build the image + subprocess.run([ + "docker", "build", + "-t", image_name, + "-f", dockerfile_path, + devcontainer_path + ], check=True) + + # Run the container + result = subprocess.run([ + "docker", "run", "-d", "--name", container_name, image_name, "tail", "-f", "/dev/null" + ], stdout=subprocess.PIPE, check=True, text=True) + + container_id = result.stdout.strip() + return container_id + +def remove_container(name_or_id: str): + try: + # Force removal (stops if running, then removes) + subprocess.run(["docker", "rm", "-f", name_or_id], check=True) + print(f"Container '{name_or_id}' removed successfully.") + except subprocess.CalledProcessError as e: + print(f"Error removing container {name_or_id}:", e) + From 3007ac49edba70e1cd7679b5a7c7495c6ca3712b Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Wed, 21 May 2025 10:07:58 +0200 Subject: [PATCH 22/25] remove devcontainer --- .devcontainer/Dockerfile | 108 ------------------ .devcontainer/devcontainer.json | 86 -------------- .devcontainer/docker-compose.yml | 78 ------------- .../qdrant_storage/aliases/data.json | 1 - .devcontainer/requirements.txt | 6 - src/cai/util.py | 37 ------ 6 files changed, 316 deletions(-) delete mode 100644 .devcontainer/Dockerfile delete mode 100644 .devcontainer/devcontainer.json delete mode 100644 .devcontainer/docker-compose.yml delete mode 100644 .devcontainer/qdrant_storage/aliases/data.json delete mode 100644 .devcontainer/requirements.txt diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile deleted file mode 100644 index 43fe72ae..00000000 --- a/.devcontainer/Dockerfile +++ /dev/null @@ -1,108 +0,0 @@ -FROM kalilinux/kali-rolling - -# Set environment variable to non-interactive -ENV DEBIAN_FRONTEND=noninteractive - -# [Optional] Uncomment this section to install additional OS packages. -RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ - && apt-get -y install --no-install-recommends \ - net-tools python3 python3-pip python3-venv \ - curl gnupg nmap iputils-ping ssh git \ - graphviz pkg-config libhdf5-dev \ - build-essential python3-dev \ - asciinema dnsutils \ - apt-transport-https ca-certificates \ - wget dirb gobuster whatweb gfortran \ - cmake libopenblas-dev sshpass seclists \ - golang-go - -# # Install Metasploit - path 1 -# RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > /tmp/msfinstall \ -# && chmod 755 /tmp/msfinstall \ -# && /tmp/msfinstall - -# # Install Metasploit - path 2 -# RUN curl https://apt.metasploit.com/metasploit-framework.gpg.key | apt-key add - -# RUN echo "deb https://apt.metasploit.com/ lucid main" | tee /etc/apt/sources.list.d/metasploit-framework.list -# RUN apt-get update && apt-get install metasploit-framework - -# Install Metasploit - path 3 -RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstall && chmod 755 msfinstall && ./msfinstall - -# Update Kali Linux repositories -RUN apt-get update \ - && apt-get upgrade -y - -# # Create a virtual environment -# ENV VIRTUAL_ENV=/opt/venv -# RUN python3 -m venv $VIRTUAL_ENV -# ENV PATH="$VIRTUAL_ENV/bin:$PATH" - -# # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image. -# # alternatively, this also helps avoid having to pull requirements from the internet every single time -# COPY requirements.txt /tmp/pip-tmp/ -# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ -# && rm -rf /tmp/pip-tmp - -# FIX for: -# -# Python packages system-wide, try apt install -# python3-xyz, where xyz is the package you are trying to -# install. -# -# see https://github.com/raspiblitz/raspiblitz/issues/4170 -RUN mkdir -p /root/.pip && touch /root/.pip/pip.conf -RUN echo "[global]" > /root/.pip/pip.conf && \ - echo "break-system-packages = true" >> /root/.pip/pip.conf - -RUN pip3 install --upgrade setuptools -RUN pip3 install pandas \ - opentelemetry-sdk opentelemetry-exporter-otlp \ - mem0ai PyPDF2 sentence_transformers tf-keras \ - pytest-repeat parameterized pytest-rerunfailures \ - pytest-clarity tiktoken blobfile pytest-timeout \ - invoke fabric - -# Install rust compiler - dependency of outlines-core -RUN apt-get update && \ - apt-get install -y curl && \ - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \ - . $HOME/.cargo/env && \ - rustup default stable && \ - rustup component add rustfmt clippy && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* - -ENV PATH="/root/.cargo/bin:${PATH}" -RUN wget https://raw.githubusercontent.com/adfoster-r7/metasploit-info/master/info/module_metadata.json -O /tmp/module_metadata.json - -# NOTE: align dev env in Linux -# work -# RUN echo "192.168.2.1 host.docker.internal" >> /etc/hosts -# Update and upgrade packages -# Preconfigure console-setup and install kali-linux-headless fixers -RUN apt-get update && \ - echo "console-setup console-setup/variant select Latin1 and Latin5 - western Europe and Turkic languages" | debconf-set-selections && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends kali-linux-headless - -# Update pip again -> python3 -m pip install --upgrade pip causes crashes -# Install spacy deps - en_core_web_sm -# RUN python3 -m spacy download en_core_web_sm - -# Install docker inside of dev env -RUN apt-get update && \ - apt-get install -y apt-transport-https ca-certificates curl gnupg && \ - curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg && \ - echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian bullseye stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && \ - apt-get update && \ - apt-get install -y docker-ce docker-ce-cli containerd.io cloc - -# # Activate virtual environment by default -# RUN echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc - -RUN cargo install --git https://github.com/asciinema/agg - -# Remove system's python3 libs to avoid conflicts -RUN apt-get remove -y python3-jsonschema && \ - apt-get remove -y python3-numpy && \ - apt-get autoremove -y diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json deleted file mode 100644 index 22e1192e..00000000 --- a/.devcontainer/devcontainer.json +++ /dev/null @@ -1,86 +0,0 @@ -// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at: -// https://github.com/microsoft/vscode-dev-containers/tree/v0.241.1/containers/python-3 -{ - "name": "cai_devenv", - - // "build": { - // "dockerfile": "Dockerfile", - // "context": "..", - // "args": { - // // Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6 - // // Append -bullseye or -buster to pin to an OS version. - // // Use -bullseye variants on local on arm64/Apple Silicon. - // "VARIANT": "3.10-bullseye", - // // Options - // "NODE_VERSION": "lts/*" - // } - // }, - - "dockerComposeFile": ["./docker-compose.yml"], - "service": "devenv", - // "shutdownAction": "none", // don't shut down container when vscode is closed - "workspaceFolder": "/workspace", - - // Configure tool-specific properties. - "customizations": { - // Configure properties specific to VS Code. - "vscode": { - // Set *default* container specific settings.json values on container create. - "settings": { - "python.defaultInterpreterPath": "/usr/local/bin/python", - "python.linting.enabled": true, - "python.linting.pylintEnabled": true, - "python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8", - "python.formatting.blackPath": "/usr/local/py-utils/bin/black", - "python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf", - "python.linting.banditPath": "/usr/local/py-utils/bin/bandit", - "python.linting.flake8Path": "/usr/local/py-utils/bin/flake8", - "python.linting.mypyPath": "/usr/local/py-utils/bin/mypy", - "python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle", - "python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", - "python.linting.pylintPath": "/usr/local/py-utils/bin/pylint" - }, - - // Add the IDs of extensions you want installed when the container is created. - "extensions": [ - "ms-python.python", - "ms-toolsai.jupyter-renderers", - "ms-toolsai.jupyter", - "ms-python.vscode-pylance", - "ms-azuretools.vscode-docker", - "ms-python.debugpy", - "ms-python.black-formatter", - "MS-vsliveshare.vsliveshare" - ] - } - }, - - // Use 'forwardPorts' to make a list of ports inside the container available locally. - // "forwardPorts": [], - - // Use 'preCreateCommand' to run commands before the container is created. - "preCreateCommand": "docker network prune -f && docker container prune -f", - - // Use 'postCreateCommand' to run commands after the container is created. - // "postCreateCommand": "pip3 install --user -r requirements.txt", // doing it instead in - // the dockerfile to cache - // requirements in container - // Install cai via "pip3 install -e ." - // Extensions, NOTE: only one postCreateCommand is allowed, so we need to install all extensions here - "postCreateCommand": "pip3 install -e /workspace/", - - // MSF setup, and RAG setup - "postStartCommand": [ - "nohup", "msfrpcd", "-P", "cai", "&", - "&&", - "python3", "cai/ins/rag/agent_helper.py" - ], - // // Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root. - // "remoteUser": "vscode" - - "runArgs": [ - "--privileged", - "-e", "DISPLAY=host.docker.internal:0", - "-v", "/tmp/.X11-unix:/tmp/.X11-unix" - ] -} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml deleted file mode 100644 index 92b14fb9..00000000 --- a/.devcontainer/docker-compose.yml +++ /dev/null @@ -1,78 +0,0 @@ -version: '3' - -################# -# SERVICES -################# -services: - # Developer environment - devenv: - build: - context: .. - dockerfile: .devcontainer/Dockerfile - # platform: linux/amd64 # TODO: uncomment this when we have a x86_64 machine - volumes: - # Mount the root folder that contains .git - - ..:/workspace:cached - # Mount Docker socket to enable Docker-in-Docker - - /var/run/docker.sock:/var/run/docker.sock - command: /bin/sh -c "while sleep 10; do :; done" - networks: - cainet: - ipv4_address: 192.168.2.5 - cap_add: - - NET_ADMIN - - SYS_ADMIN - privileged: true - - # # Vulnhub vulnerable machine (Hackable II) SUGGESTION: Comment this and use PentestPerf - # hackableii: - # image: vmayoral/vulnhub:hackableii - # command: | - # /bin/bash -c " - # rm -rf /var/lock && \ - # mkdir -p /var/lock && \ - # chmod 755 /var/lock && \ - # /etc/init.d/apache2 start && \ - # /etc/init.d/ssh start && \ - # /etc/init.d/runproftpd.sh && \ - # /etc/init.d/php7.0-fpm start && \ - # while sleep 1; do :; done - # " - # networks: - # cainet: - # ipv4_address: 192.168.2.11 - # mac_address: 08:00:27:85:55:86 - - # # Vulnhub vulnerable machine (Bob) SUGGESTION: Comment this and use PentestPerf - # bob: - # image: vmayoral/vulnhub:bob - # command: | - # /bin/bash -c "rm -r /var/lock; mkdir -p /var/lock; chmod 755 /var/lock; /etc/init.d/apache2 start; /etc/init.d/ssh start; while sleep 10; do :; done" - # ports: # map port in the container to the host systems - # - "8080:80" - # networks: - # cainet: - # ipv4_address: 192.168.2.12 - # mac_address: 08:00:27:cb:07:d4 - - # Qdrant vector database - qdrant: - image: qdrant/qdrant - ports: - - "6333:6333" - - "6334:6334" - volumes: - - ./qdrant_storage:/qdrant/storage:z - networks: - cainet: - ipv4_address: 192.168.2.13 - -################# -# NETWORKS -################# -networks: - cainet: - ipam: - driver: default - config: - - subnet: 192.168.2.0/24 diff --git a/.devcontainer/qdrant_storage/aliases/data.json b/.devcontainer/qdrant_storage/aliases/data.json deleted file mode 100644 index 9e26dfee..00000000 --- a/.devcontainer/qdrant_storage/aliases/data.json +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt deleted file mode 100644 index 2409daf2..00000000 --- a/.devcontainer/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -pymetasploit3 -networkx==2.5 -requests -wasabi -xmltodict -pydot==1.4.2 diff --git a/src/cai/util.py b/src/cai/util.py index ff62608b..b7223fad 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -2997,40 +2997,3 @@ def setup_ctf(): bg="blue")) return ctf, messages - - -def build_docker_container(devcontainer_path: str) -> str: - dockerfile_path = os.path.join(devcontainer_path, "Dockerfile") - image_name = "cai-container3" - container_name = "cai-container3" - - # Try to remove existing container with the same name - try: - subprocess.run(["docker", "rm", "-f", container_name], check=True) - except subprocess.CalledProcessError: - pass # If it doesn't exist, continue normally - - # Build the image - subprocess.run([ - "docker", "build", - "-t", image_name, - "-f", dockerfile_path, - devcontainer_path - ], check=True) - - # Run the container - result = subprocess.run([ - "docker", "run", "-d", "--name", container_name, image_name, "tail", "-f", "/dev/null" - ], stdout=subprocess.PIPE, check=True, text=True) - - container_id = result.stdout.strip() - return container_id - -def remove_container(name_or_id: str): - try: - # Force removal (stops if running, then removes) - subprocess.run(["docker", "rm", "-f", name_or_id], check=True) - print(f"Container '{name_or_id}' removed successfully.") - except subprocess.CalledProcessError as e: - print(f"Error removing container {name_or_id}:", e) - From 3229320af2ca37328c0a4d38e510f7a0e178a1a5 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 22 May 2025 11:47:38 +0200 Subject: [PATCH 23/25] Add qdrant --- .devcontainer/Dockerfile | 108 ++++++++++++++++++ .devcontainer/devcontainer.json | 86 ++++++++++++++ .devcontainer/docker-compose.yml | 78 +++++++++++++ .../qdrant_storage/aliases/data.json | 1 + .devcontainer/qdrant_storage/raft_state.json | 2 +- .devcontainer/requirements.txt | 6 + .gitignore | 2 +- 7 files changed, 281 insertions(+), 2 deletions(-) create mode 100644 .devcontainer/Dockerfile create mode 100644 .devcontainer/devcontainer.json create mode 100644 .devcontainer/docker-compose.yml create mode 100644 .devcontainer/qdrant_storage/aliases/data.json create mode 100644 .devcontainer/requirements.txt diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 00000000..43fe72ae --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,108 @@ +FROM kalilinux/kali-rolling + +# Set environment variable to non-interactive +ENV DEBIAN_FRONTEND=noninteractive + +# [Optional] Uncomment this section to install additional OS packages. +RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ + && apt-get -y install --no-install-recommends \ + net-tools python3 python3-pip python3-venv \ + curl gnupg nmap iputils-ping ssh git \ + graphviz pkg-config libhdf5-dev \ + build-essential python3-dev \ + asciinema dnsutils \ + apt-transport-https ca-certificates \ + wget dirb gobuster whatweb gfortran \ + cmake libopenblas-dev sshpass seclists \ + golang-go + +# # Install Metasploit - path 1 +# RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > /tmp/msfinstall \ +# && chmod 755 /tmp/msfinstall \ +# && /tmp/msfinstall + +# # Install Metasploit - path 2 +# RUN curl https://apt.metasploit.com/metasploit-framework.gpg.key | apt-key add - +# RUN echo "deb https://apt.metasploit.com/ lucid main" | tee /etc/apt/sources.list.d/metasploit-framework.list +# RUN apt-get update && apt-get install metasploit-framework + +# Install Metasploit - path 3 +RUN curl https://raw.githubusercontent.com/rapid7/metasploit-omnibus/master/config/templates/metasploit-framework-wrappers/msfupdate.erb > msfinstall && chmod 755 msfinstall && ./msfinstall + +# Update Kali Linux repositories +RUN apt-get update \ + && apt-get upgrade -y + +# # Create a virtual environment +# ENV VIRTUAL_ENV=/opt/venv +# RUN python3 -m venv $VIRTUAL_ENV +# ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +# # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image. +# # alternatively, this also helps avoid having to pull requirements from the internet every single time +# COPY requirements.txt /tmp/pip-tmp/ +# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ +# && rm -rf /tmp/pip-tmp + +# FIX for: +# +# Python packages system-wide, try apt install +# python3-xyz, where xyz is the package you are trying to +# install. +# +# see https://github.com/raspiblitz/raspiblitz/issues/4170 +RUN mkdir -p /root/.pip && touch /root/.pip/pip.conf +RUN echo "[global]" > /root/.pip/pip.conf && \ + echo "break-system-packages = true" >> /root/.pip/pip.conf + +RUN pip3 install --upgrade setuptools +RUN pip3 install pandas \ + opentelemetry-sdk opentelemetry-exporter-otlp \ + mem0ai PyPDF2 sentence_transformers tf-keras \ + pytest-repeat parameterized pytest-rerunfailures \ + pytest-clarity tiktoken blobfile pytest-timeout \ + invoke fabric + +# Install rust compiler - dependency of outlines-core +RUN apt-get update && \ + apt-get install -y curl && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \ + . $HOME/.cargo/env && \ + rustup default stable && \ + rustup component add rustfmt clippy && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +ENV PATH="/root/.cargo/bin:${PATH}" +RUN wget https://raw.githubusercontent.com/adfoster-r7/metasploit-info/master/info/module_metadata.json -O /tmp/module_metadata.json + +# NOTE: align dev env in Linux +# work +# RUN echo "192.168.2.1 host.docker.internal" >> /etc/hosts +# Update and upgrade packages +# Preconfigure console-setup and install kali-linux-headless fixers +RUN apt-get update && \ + echo "console-setup console-setup/variant select Latin1 and Latin5 - western Europe and Turkic languages" | debconf-set-selections && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends kali-linux-headless + +# Update pip again -> python3 -m pip install --upgrade pip causes crashes +# Install spacy deps - en_core_web_sm +# RUN python3 -m spacy download en_core_web_sm + +# Install docker inside of dev env +RUN apt-get update && \ + apt-get install -y apt-transport-https ca-certificates curl gnupg && \ + curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg && \ + echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian bullseye stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && \ + apt-get update && \ + apt-get install -y docker-ce docker-ce-cli containerd.io cloc + +# # Activate virtual environment by default +# RUN echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc + +RUN cargo install --git https://github.com/asciinema/agg + +# Remove system's python3 libs to avoid conflicts +RUN apt-get remove -y python3-jsonschema && \ + apt-get remove -y python3-numpy && \ + apt-get autoremove -y diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 00000000..22e1192e --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,86 @@ +// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at: +// https://github.com/microsoft/vscode-dev-containers/tree/v0.241.1/containers/python-3 +{ + "name": "cai_devenv", + + // "build": { + // "dockerfile": "Dockerfile", + // "context": "..", + // "args": { + // // Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6 + // // Append -bullseye or -buster to pin to an OS version. + // // Use -bullseye variants on local on arm64/Apple Silicon. + // "VARIANT": "3.10-bullseye", + // // Options + // "NODE_VERSION": "lts/*" + // } + // }, + + "dockerComposeFile": ["./docker-compose.yml"], + "service": "devenv", + // "shutdownAction": "none", // don't shut down container when vscode is closed + "workspaceFolder": "/workspace", + + // Configure tool-specific properties. + "customizations": { + // Configure properties specific to VS Code. + "vscode": { + // Set *default* container specific settings.json values on container create. + "settings": { + "python.defaultInterpreterPath": "/usr/local/bin/python", + "python.linting.enabled": true, + "python.linting.pylintEnabled": true, + "python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8", + "python.formatting.blackPath": "/usr/local/py-utils/bin/black", + "python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf", + "python.linting.banditPath": "/usr/local/py-utils/bin/bandit", + "python.linting.flake8Path": "/usr/local/py-utils/bin/flake8", + "python.linting.mypyPath": "/usr/local/py-utils/bin/mypy", + "python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle", + "python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", + "python.linting.pylintPath": "/usr/local/py-utils/bin/pylint" + }, + + // Add the IDs of extensions you want installed when the container is created. + "extensions": [ + "ms-python.python", + "ms-toolsai.jupyter-renderers", + "ms-toolsai.jupyter", + "ms-python.vscode-pylance", + "ms-azuretools.vscode-docker", + "ms-python.debugpy", + "ms-python.black-formatter", + "MS-vsliveshare.vsliveshare" + ] + } + }, + + // Use 'forwardPorts' to make a list of ports inside the container available locally. + // "forwardPorts": [], + + // Use 'preCreateCommand' to run commands before the container is created. + "preCreateCommand": "docker network prune -f && docker container prune -f", + + // Use 'postCreateCommand' to run commands after the container is created. + // "postCreateCommand": "pip3 install --user -r requirements.txt", // doing it instead in + // the dockerfile to cache + // requirements in container + // Install cai via "pip3 install -e ." + // Extensions, NOTE: only one postCreateCommand is allowed, so we need to install all extensions here + "postCreateCommand": "pip3 install -e /workspace/", + + // MSF setup, and RAG setup + "postStartCommand": [ + "nohup", "msfrpcd", "-P", "cai", "&", + "&&", + "python3", "cai/ins/rag/agent_helper.py" + ], + // // Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root. + // "remoteUser": "vscode" + + "runArgs": [ + "--privileged", + "-e", "DISPLAY=host.docker.internal:0", + "-v", "/tmp/.X11-unix:/tmp/.X11-unix" + ] +} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml new file mode 100644 index 00000000..92b14fb9 --- /dev/null +++ b/.devcontainer/docker-compose.yml @@ -0,0 +1,78 @@ +version: '3' + +################# +# SERVICES +################# +services: + # Developer environment + devenv: + build: + context: .. + dockerfile: .devcontainer/Dockerfile + # platform: linux/amd64 # TODO: uncomment this when we have a x86_64 machine + volumes: + # Mount the root folder that contains .git + - ..:/workspace:cached + # Mount Docker socket to enable Docker-in-Docker + - /var/run/docker.sock:/var/run/docker.sock + command: /bin/sh -c "while sleep 10; do :; done" + networks: + cainet: + ipv4_address: 192.168.2.5 + cap_add: + - NET_ADMIN + - SYS_ADMIN + privileged: true + + # # Vulnhub vulnerable machine (Hackable II) SUGGESTION: Comment this and use PentestPerf + # hackableii: + # image: vmayoral/vulnhub:hackableii + # command: | + # /bin/bash -c " + # rm -rf /var/lock && \ + # mkdir -p /var/lock && \ + # chmod 755 /var/lock && \ + # /etc/init.d/apache2 start && \ + # /etc/init.d/ssh start && \ + # /etc/init.d/runproftpd.sh && \ + # /etc/init.d/php7.0-fpm start && \ + # while sleep 1; do :; done + # " + # networks: + # cainet: + # ipv4_address: 192.168.2.11 + # mac_address: 08:00:27:85:55:86 + + # # Vulnhub vulnerable machine (Bob) SUGGESTION: Comment this and use PentestPerf + # bob: + # image: vmayoral/vulnhub:bob + # command: | + # /bin/bash -c "rm -r /var/lock; mkdir -p /var/lock; chmod 755 /var/lock; /etc/init.d/apache2 start; /etc/init.d/ssh start; while sleep 10; do :; done" + # ports: # map port in the container to the host systems + # - "8080:80" + # networks: + # cainet: + # ipv4_address: 192.168.2.12 + # mac_address: 08:00:27:cb:07:d4 + + # Qdrant vector database + qdrant: + image: qdrant/qdrant + ports: + - "6333:6333" + - "6334:6334" + volumes: + - ./qdrant_storage:/qdrant/storage:z + networks: + cainet: + ipv4_address: 192.168.2.13 + +################# +# NETWORKS +################# +networks: + cainet: + ipam: + driver: default + config: + - subnet: 192.168.2.0/24 diff --git a/.devcontainer/qdrant_storage/aliases/data.json b/.devcontainer/qdrant_storage/aliases/data.json new file mode 100644 index 00000000..9e26dfee --- /dev/null +++ b/.devcontainer/qdrant_storage/aliases/data.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/.devcontainer/qdrant_storage/raft_state.json b/.devcontainer/qdrant_storage/raft_state.json index 23b070d6..b7bcb749 100644 --- a/.devcontainer/qdrant_storage/raft_state.json +++ b/.devcontainer/qdrant_storage/raft_state.json @@ -1 +1 @@ -{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6120741261401965],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6120741261401965,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6120741261401965} +{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6120741261401965],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6120741261401965,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6120741261401965} \ No newline at end of file diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt new file mode 100644 index 00000000..2409daf2 --- /dev/null +++ b/.devcontainer/requirements.txt @@ -0,0 +1,6 @@ +pymetasploit3 +networkx==2.5 +requests +wasabi +xmltodict +pydot==1.4.2 diff --git a/.gitignore b/.gitignore index 5101cdeb..f23dc111 100644 --- a/.gitignore +++ b/.gitignore @@ -30,7 +30,7 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST - +qdrant_storage/ # PyInstaller *.manifest *.spec From 5c6880a336f5b861ce2b4f99320c0a0992b9fa53 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 22 May 2025 11:48:02 +0200 Subject: [PATCH 24/25] Add qdrant --- .devcontainer/qdrant_storage/aliases/data.json | 1 - .devcontainer/qdrant_storage/raft_state.json | 1 - .gitignore | 2 +- 3 files changed, 1 insertion(+), 3 deletions(-) delete mode 100644 .devcontainer/qdrant_storage/aliases/data.json delete mode 100644 .devcontainer/qdrant_storage/raft_state.json diff --git a/.devcontainer/qdrant_storage/aliases/data.json b/.devcontainer/qdrant_storage/aliases/data.json deleted file mode 100644 index 9e26dfee..00000000 --- a/.devcontainer/qdrant_storage/aliases/data.json +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/.devcontainer/qdrant_storage/raft_state.json b/.devcontainer/qdrant_storage/raft_state.json deleted file mode 100644 index b7bcb749..00000000 --- a/.devcontainer/qdrant_storage/raft_state.json +++ /dev/null @@ -1 +0,0 @@ -{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[6120741261401965],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":6120741261401965,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":6120741261401965} \ No newline at end of file diff --git a/.gitignore b/.gitignore index f23dc111..593a7ca6 100644 --- a/.gitignore +++ b/.gitignore @@ -30,7 +30,7 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST -qdrant_storage/ +.devcontainer/qdrant_storage/ # PyInstaller *.manifest *.spec From 7e32792a2446e44847da0b5eda2f29b99c1817c5 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 22 May 2025 11:17:42 +0000 Subject: [PATCH 25/25] Improve prompting --- src/cai/cli.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index 51255117..b6cbe6e2 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -289,7 +289,6 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), console.print("[yellow]You must increase the limit using the /config command: /config CAI_MAX_TURNS=[/yellow]") console.print("[yellow]Only CLI commands (starting with '/') will be processed until the limit is increased.[/yellow]") ->>>>>>> src/cai/cli.py try: # Start measuring user idle time start_idle_timer() @@ -331,7 +330,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), - if not force_until_flag: + if not force_until_flag and messages_ctf == "": # Get user input with command completion and history user_input = get_user_input( command_completer, @@ -341,7 +340,9 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), current_text ) + messages_ctf else: - user_input = messages_ctf + user_input = messages_ctf + if os.getenv('CTF_INSIDE', 'True').lower() == 'false': + user_input += ctf_global.get_ip() idle_time += time.time() - idle_start_time # Stop measuring user idle time and start measuring active time