diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index fee806da..33e189cb 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -36,7 +36,8 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li agent, history, config, - flush + flush, + workspace, ) # Define helper functions diff --git a/src/cai/repl/commands/config.py b/src/cai/repl/commands/config.py index 3588cd19..06a72ebe 100644 --- a/src/cai/repl/commands/config.py +++ b/src/cai/repl/commands/config.py @@ -128,6 +128,11 @@ ENV_VARS = { "description": "Boolean to enable real-time, chunked responses instead of full messages.", "default": "True" }, + 23: { + "name": "CAI_WORKSPACE", + "description": "Name of the current workspace (affects log file naming)", + "default": None + }, } diff --git a/src/cai/repl/commands/workspace.py b/src/cai/repl/commands/workspace.py new file mode 100644 index 00000000..a00ff36c --- /dev/null +++ b/src/cai/repl/commands/workspace.py @@ -0,0 +1,687 @@ +""" +Virtualization command for CAI REPL. +This module provides commands for setting up and managing Docker virtualization +environments. +""" +# Standard library imports +import os +import json +import subprocess +import datetime +import time +from typing import List, Optional, Dict, Any, Tuple + +# Third-party imports +from rich.console import Console +from rich.table import Table +from rich.panel import Panel +from rich.markdown import Markdown +import rich.box + +# Local imports +from cai.repl.commands.base import Command, register_command + +console = Console() + +class WorkspaceCommand(Command): + """Command for workspace management within Docker containers or locally.""" + + def __init__(self): + """Initialize the workspace command.""" + super().__init__( + name="/workspace", + description=( + "Set or display the current workspace name and manage files." + " Affects log file naming and where files are stored." + ), + aliases=["/ws"] + ) + + # Add subcommands + self.add_subcommand( + "set", + "Set the current workspace name", + self.handle_set + ) + self.add_subcommand( + "get", + "Display the current workspace name", + self.handle_get + ) + self.add_subcommand( + "ls", + "List files in the workspace", + self.handle_ls_subcommand + ) + self.add_subcommand( + "exec", + "Execute a command in the workspace", + self.handle_exec_subcommand + ) + self.add_subcommand( + "copy", + "Copy files between host and container", + self.handle_copy_subcommand + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the workspace command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + # If there are subcommands, process them + if args and args[0] in self.subcommands: + return super().handle(args) + + # No arguments means show workspace info (same as get) + return self.handle_get() + + def handle_no_args(self) -> bool: + """Handle the command when no arguments are provided.""" + return self.handle_get() + + def handle_get(self, _: Optional[List[str]] = None) -> bool: + """Display the current workspace name and directory information.""" + # Get workspace info + workspace_name = os.getenv("CAI_WORKSPACE", None) + + # Check if a container is active + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # Determine environment (container or host) + if active_container: + try: + # Get container details + result = subprocess.run( + ["docker", "inspect", active_container], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + container_info = json.loads(result.stdout) + if container_info: + image = container_info[0].get("Config", {}).get("Image", "unknown") + env_type = "container" + env_name = f"Container ({image})" + + # For containers, if workspace is set, use container workspace path + # otherwise use root directory + if workspace_name: + # This will create the workspace in the container if it doesn't exist + workspace_dir = f"/workspace/workspaces/{workspace_name}" + # Ensure the directory exists in the container + subprocess.run( + ["docker", "exec", active_container, "mkdir", "-p", workspace_dir], + capture_output=True, + check=False + ) + else: + workspace_dir = "/" + else: + env_type = "host" + env_name = "Host System (container not running)" + # Use common._get_workspace_dir() for consistency + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + workspace_dir = get_common_workspace_dir() + except ImportError: + workspace_dir = os.getcwd() # Basic fallback + except Exception: + env_type = "host" + env_name = "Host System (error inspecting container)" + # Use common._get_workspace_dir() for consistency + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + workspace_dir = get_common_workspace_dir() + except ImportError: + workspace_dir = os.getcwd() # Basic fallback + else: + env_type = "host" + env_name = "Host System" + # Use common._get_workspace_dir() for consistency + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + workspace_dir = get_common_workspace_dir() + except ImportError: + workspace_dir = os.getcwd() # Basic fallback + + # Show workspace information + console.print( + Panel( + f"Current workspace: [bold green]{workspace_name or 'None'}[/bold green]\n" + f"Working in environment: [bold]{env_name}[/bold]\n" + f"Workspace directory: [bold]{workspace_dir}[/bold]", + title="Workspace Information", + border_style="green" + ) + ) + + # Show available workspace commands + console.print("\n[cyan]Workspace Commands:[/cyan]") + console.print( + " [bold]/workspace set [/bold] - " + "Set the current workspace name") + console.print( + " [bold]/workspace ls[/bold] - " + "List files in the workspace") + console.print( + " [bold]/workspace exec [/bold] - " + "Execute a command in the workspace") + + if active_container: + console.print( + " [bold]/workspace copy [/bold] - " + "Copy files between host and container") + + # List contents of the workspace + self._list_workspace_contents(env_type, workspace_dir) + + return True + + def handle_set(self, args: Optional[List[str]] = None) -> bool: + """Set the current workspace name """ + if not args or len(args) != 1: + console.print( + "[yellow]Usage: /workspace set [/yellow]" + ) + return False + + workspace_name = args[0] + # Allow alphanumeric, underscores, hyphens + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + console.print( + "[red]Invalid workspace name. " + "Use alphanumeric, underscores, or hyphens only.[/red]" + ) + return False + + # Import the necessary modules for setting environment variables + # And for getting workspace dir consistently + try: + from cai.repl.commands.config import set_env_var + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + from cai.tools.common import _get_container_workspace_path as get_common_container_path + + # Set the environment variable + if not set_env_var("CAI_WORKSPACE", workspace_name): + console.print( + "[red]Failed to set workspace environment variable.[/red]" + ) + return False + except ImportError: + # Fallback if import fails + os.environ["CAI_WORKSPACE"] = workspace_name + # Define basic fallbacks for path functions if import failed + def get_common_workspace_dir(): + base = os.getenv("CAI_WORKSPACE_DIR", ".") # Default to current dir base + name = os.getenv("CAI_WORKSPACE") + if name: + return os.path.abspath(os.path.join(base, name)) + return os.path.abspath(base) # Use base dir if no name + + def get_common_container_path(): + name = os.getenv("CAI_WORKSPACE") + if name: + return f"/workspace/workspaces/{name}" + return "/" # Default container path + + # Get the new workspace directory using the common function + new_workspace_dir = get_common_workspace_dir() + + # Create the directory if it doesn't exist on host + try: # Add try-except for robustness + os.makedirs(new_workspace_dir, exist_ok=True) + except OSError as e: + console.print(f"[red]Error creating host directory {new_workspace_dir}: {e}[/red]") + # Decide if this is fatal or just a warning + + # If container is active, also create the directory in the container + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + if active_container: + # Check if container is running + check_process = subprocess.run( + ["docker", "inspect", "--format", "{{.State.Running}}", active_container], + capture_output=True, + text=True, + check=False + ) + + if check_process.returncode == 0 and "true" in check_process.stdout.lower(): + # Get container workspace path using the common function + container_workspace_path = get_common_container_path() + try: + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path] + mkdir_result = subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + if mkdir_result.returncode == 0: + console.print( + f"[dim]Created workspace directory in container: {container_workspace_path}[/dim]" + ) + else: + console.print( + f"[yellow]Warning: Could not create workspace directory in container: {mkdir_result.stderr}[/yellow]" + ) + except Exception as e: + console.print( + f"[yellow]Warning: Failed to setup workspace in container: {str(e)}[/yellow]" + ) + + # Use a different panel style to indicate success + console.print( + Panel( + f"Workspace changed to: [bold green]{workspace_name}[/bold green]\n" + f"New workspace directory: [bold]{new_workspace_dir}[/bold]", + title="Workspace Updated", + border_style="green" + ) + ) + + return True + + def _get_workspace_dir(self) -> str: + """Get the host workspace directory using the common utility. + + Returns: + The host workspace directory path. + """ + try: + # Use the centralized function from common.py + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + return get_common_workspace_dir() + except ImportError: + # Provide a basic fallback if import fails, mirroring common.py logic + # without 'cai_default' + base_dir = os.getenv("CAI_WORKSPACE_DIR") + workspace_name = os.getenv("CAI_WORKSPACE") + + if base_dir and workspace_name: + # Basic validation + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(f"[yellow]Warning: Invalid CAI_WORKSPACE name '{workspace_name}' in fallback.[/yellow]") + # Fallback to base directory if name is invalid + return os.path.abspath(base_dir) + target_dir = os.path.join(base_dir, workspace_name) + return os.path.abspath(target_dir) + elif base_dir: + # If only base dir is set, use that + return os.path.abspath(base_dir) + else: + # Default to current working directory if nothing else is set + return os.getcwd() + + def _list_workspace_contents(self, env_type: str, workspace_dir: str) -> None: + """List the contents of the workspace. + + Args: + env_type: The environment type (container or host) + workspace_dir: The workspace directory + """ + console.print("\n[bold]Workspace Contents:[/bold]") + + if env_type == "container": + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # For containers, use the workspace path provided + # This should already be the correct path from handle_get + + # First ensure the workspace directory exists in the container + try: + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", workspace_dir] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + # Now list the contents + result = subprocess.run( + ["docker", "exec", active_container, "ls", "-la", workspace_dir], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + else: + console.print(f"[yellow]Error listing container files: {result.stderr}[/yellow]") + # Fallback to host + self._list_host_files(workspace_dir) + except Exception as e: + console.print(f"[yellow]Error accessing container: {str(e)}[/yellow]") + # Fallback to host + self._list_host_files(workspace_dir) + else: + # List files in host + self._list_host_files(workspace_dir) + + def _list_host_files(self, workspace_dir: str) -> None: + """List files in the host workspace. + + Args: + workspace_dir: The workspace directory + """ + # Ensure the directory exists + os.makedirs(workspace_dir, exist_ok=True) + + try: + result = subprocess.run( + ["ls", "-la", workspace_dir], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + else: + console.print(f"[yellow]Error listing files: {result.stderr}[/yellow]") + except Exception as e: + console.print(f"[yellow]Error: {str(e)}[/yellow]") + + def handle_ls_subcommand(self, args: Optional[List[str]] = None) -> bool: + """Handle the ls subcommand. + + Args: + args: Optional list of subcommand arguments + + Returns: + True if the subcommand was handled successfully, False otherwise + """ + # Get workspace info using common functions + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + from cai.tools.common import _get_container_workspace_path as get_common_container_path + except ImportError: + # Define basic fallbacks if import fails + def get_common_workspace_dir(): + base = os.getenv("CAI_WORKSPACE_DIR", ".") + name = os.getenv("CAI_WORKSPACE") + if name: return os.path.abspath(os.path.join(base, name)) + return os.path.abspath(base) + def get_common_container_path(): + name = os.getenv("CAI_WORKSPACE") + if name: return f"/workspace/workspaces/{name}" + return "/" + + host_workspace_dir = get_common_workspace_dir() + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # Execute command in the appropriate environment + if active_container: + # Use the container workspace path from common function + container_workspace_path = get_common_container_path() + + # Determine the target path within the container + target_path_in_container = container_workspace_path + if args: + # Ensure args[0] is treated as relative to the workspace + target_path_in_container = os.path.join(container_workspace_path, args[0]) + + # Ensure the base workspace directory exists in the container + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + # Try in container + result = subprocess.run( + ["docker", "exec", active_container, "ls", "-la", target_path_in_container], # Use target path + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + return True + + # If failed, try on host + console.print(f"[yellow]Failed to list files in container: {result.stderr}[/yellow]") + console.print("[yellow]Falling back to host system...[/yellow]") + + # List on host + # Determine target path on host relative to host workspace dir + target_path_on_host = host_workspace_dir + if args: + # Ensure args[0] is treated as relative to the workspace + target_path_on_host = os.path.join(host_workspace_dir, args[0]) + + # Ensure the target directory exists on host before listing + # Use os.path.dirname if target is potentially a file path + dir_to_ensure = os.path.dirname(target_path_on_host) if '.' in os.path.basename(target_path_on_host) else target_path_on_host + try: + os.makedirs(dir_to_ensure, exist_ok=True) + except OSError as e: + console.print(f"[red]Error creating directory {dir_to_ensure} on host: {e}[/red]") + # Potentially return False or handle error appropriately + + try: + result = subprocess.run( + ["ls", "-la", target_path_on_host], # Use target path + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + return True + else: + console.print(f"[red]Error listing files: {result.stderr}[/red]") + return False + except Exception as e: + console.print(f"[red]Error: {str(e)}[/red]") + return False + + return True + + def handle_exec_subcommand(self, args: Optional[List[str]] = None) -> bool: + """Handle the exec subcommand. + + Args: + args: Optional list of subcommand arguments + + Returns: + True if the subcommand was handled successfully, False otherwise + """ + if not args: + console.print("[yellow]Please specify a command to execute.[/yellow]") + return False + + command = " ".join(args) + # Get workspace info using common functions + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + from cai.tools.common import _get_container_workspace_path as get_common_container_path + except ImportError: + # Define basic fallbacks if import fails + def get_common_workspace_dir(): + base = os.getenv("CAI_WORKSPACE_DIR", ".") + name = os.getenv("CAI_WORKSPACE") + if name: return os.path.abspath(os.path.join(base, name)) + return os.path.abspath(base) + def get_common_container_path(): + name = os.getenv("CAI_WORKSPACE") + if name: return f"/workspace/workspaces/{name}" + return "/" + + host_workspace_dir = get_common_workspace_dir() + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # Execute in container if active + if active_container: + try: + # Use the container workspace path from common function + container_workspace_path = get_common_container_path() + + # First ensure the workspace directory exists in the container + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + # Execute the command in the container's workspace directory + result = subprocess.run( + ["docker", "exec", "-w", container_workspace_path, active_container, "sh", "-c", command], + capture_output=True, + text=True, + check=False + ) + + console.print(f"[dim]$ {command}[/dim]") + if result.stdout: + console.print(result.stdout) + + if result.stderr: + console.print(f"[yellow]{result.stderr}[/yellow]") + + if result.returncode != 0: + console.print("[yellow]Command failed in container. Trying on host...[/yellow]") + return self._exec_on_host(command, host_workspace_dir) # Pass host_workspace_dir + + return True + except Exception as e: + console.print(f"[yellow]Error executing in container: {str(e)}[/yellow]") + console.print("[yellow]Falling back to host execution...[/yellow]") + + # Execute on host + return self._exec_on_host(command, host_workspace_dir) # Pass host_workspace_dir + + def _exec_on_host(self, command: str, workspace_dir: str) -> bool: + """Execute a command on the host. + + Args: + command: The command to execute + workspace_dir: The workspace directory + + Returns: + True if the command was executed successfully, False otherwise + """ + # Ensure the directory exists + os.makedirs(workspace_dir, exist_ok=True) + + try: + result = subprocess.run( + command, + shell=True, # nosec B602 + capture_output=True, + text=True, + check=False, + cwd=workspace_dir + ) + + console.print(f"[dim]$ {command}[/dim]") + if result.stdout: + console.print(result.stdout) + + if result.stderr: + console.print(f"[yellow]{result.stderr}[/yellow]") + + return result.returncode == 0 + except Exception as e: + console.print(f"[red]Error executing command: {str(e)}[/red]") + return False + + def handle_copy_subcommand(self, args: Optional[List[str]] = None) -> bool: + """Handle the copy subcommand. + + Args: + args: Optional list of subcommand arguments + + Returns: + True if the subcommand was handled successfully, False otherwise + """ + if not args or len(args) < 2: + console.print("[yellow]Please specify source and destination for copy.[/yellow]") + console.print("Usage: /workspace copy ") + return False + + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + if not active_container: + console.print("[yellow]No active container. Copy only works with containers.[/yellow]") + return False + + source = args[0] + destination = args[1] + + # Check if copying from container to host or vice versa + if source.startswith("container:"): + # Copy from container to host + container_path = source[10:] # Remove "container:" prefix + host_path = destination + + if not container_path.startswith("/"): + container_path = f"/workspace/{container_path}" + + try: + result = subprocess.run( + ["docker", "cp", f"{active_container}:{container_path}", host_path], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(f"[green]Copied from container:{container_path} to {host_path}[/green]") + return True + else: + console.print(f"[red]Error copying from container: {result.stderr}[/red]") + return False + except Exception as e: + console.print(f"[red]Error: {str(e)}[/red]") + return False + elif destination.startswith("container:"): + # Copy from host to container + host_path = source + container_path = destination[10:] # Remove "container:" prefix + + if not container_path.startswith("/"): + container_path = f"/workspace/{container_path}" + + try: + result = subprocess.run( + ["docker", "cp", host_path, f"{active_container}:{container_path}"], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(f"[green]Copied from {host_path} to container:{container_path}[/green]") + return True + else: + console.print(f"[red]Error copying to container: {result.stderr}[/red]") + return False + except Exception as e: + console.print(f"[red]Error: {str(e)}[/red]") + return False + else: + # Ambiguous copy - show help + console.print("[yellow]Ambiguous copy direction. Please specify container: prefix.[/yellow]") + console.print("Examples:") + console.print(" /workspace copy file.txt container:file.txt # Host to container") + console.print(" /workspace copy container:file.txt file.txt # Container to host") + return False + + +# Register the commands +register_command(WorkspaceCommand()) \ No newline at end of file diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py index 44874ac5..95b84728 100644 --- a/src/cai/repl/ui/toolbar.py +++ b/src/cai/repl/ui/toolbar.py @@ -57,7 +57,22 @@ def update_toolbar_in_background(): ip_address = sys_info['ip_address'] os_name = sys_info['os_name'] os_version = sys_info['os_version'] + + # Get the current workspace and base directory + workspace_name = os.getenv("CAI_WORKSPACE") + base_dir = os.getenv("CAI_WORKSPACE_DIR", "workspaces") + # Construct the workspace path + standard_path = os.path.join(base_dir, workspace_name) if workspace_name else "" + workspace_path = "" + if workspace_name: + if os.path.isdir(standard_path): + workspace_path = standard_path + elif os.path.isdir(workspace_name): + workspace_path = os.path.abspath(workspace_name) + else: + workspace_path = standard_path + # Get Ollama information ollama_status = "unavailable" try: diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index e1e58778..e475d5f4 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -24,14 +24,70 @@ except ImportError: # Global dictionary to store active sessions ACTIVE_SESSIONS = {} +def _get_workspace_dir() -> str: + """Determines the target workspace directory based on env vars for host.""" + base_dir_env = os.getenv("CAI_WORKSPACE_DIR") + workspace_name = os.getenv("CAI_WORKSPACE") + + # Determine the base directory + if base_dir_env: + base_dir = os.path.abspath(base_dir_env) + else: # Default base directory is 'workspaces' + if workspace_name: + base_dir = os.path.join(os.getcwd(), "workspaces") + else: # If no workspace name is set, the workspace IS the CWD. + return os.getcwd() + + # If a workspace name is provided, append it to the base directory + if workspace_name: + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}'. " + f"Using directory '{base_dir}' instead.", fg="yellow")) + target_dir = base_dir + else: + target_dir = os.path.join(base_dir, workspace_name) + else: + target_dir = base_dir + + # Ensure the final target directory exists on the host + try: + abs_target_dir = os.path.abspath(target_dir) + os.makedirs(abs_target_dir, exist_ok=True) + return abs_target_dir + except OSError as e: + print(color(f"Error creating/accessing host workspace directory '{abs_target_dir}': {e}", + fg="red")) + print(color(f"Falling back to current directory: {os.getcwd()}", fg="yellow")) + return os.getcwd() + +def _get_container_workspace_path() -> str: + """Determines the target workspace path inside the container.""" + workspace_name = os.getenv("CAI_WORKSPACE") + if workspace_name: + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}' for container. " + f"Using '/workspace'.", fg="yellow")) + return "/" + # Standard path inside CAI containers + return f"/workspace/workspaces/{workspace_name}" + else: + return "/" class ShellSession: # pylint: disable=too-many-instance-attributes """Class to manage interactive shell sessions""" - def __init__(self, command, session_id=None, ctf=None): + def __init__(self, command, session_id=None, ctf=None, workspace_dir=None, container_id=None): # noqa E501 self.session_id = session_id or str(uuid.uuid4())[:8] - self.command = command + self.command = command self.ctf = ctf + self.container_id = container_id + # Determine workspace based on context (container, ctf or local host) + if self.container_id: + self.workspace_dir = _get_container_workspace_path() + elif self.ctf: + self.workspace_dir = workspace_dir or _get_workspace_dir() + else: + self.workspace_dir = _get_workspace_dir() self.process = None self.master = None self.slave = None @@ -40,9 +96,40 @@ class ShellSession: # pylint: disable=too-many-instance-attributes self.last_activity = time.time() def start(self): - """Start the shell session""" + """Start the shell session in the appropriate environment.""" + start_message_cmd = self.command + + # --- Start in Container --- + if self.container_id: + try: + self.master, self.slave = pty.openpty() + docker_cmd_list = [ + "docker", "exec", "-i", + "-w", self.workspace_dir, + self.container_id, + "sh", "-c", # Use shell to handle complex commands if needed + self.command # The actual command to run + ] + self.process = subprocess.Popen( + docker_cmd_list, + stdin=self.slave, + stdout=self.slave, + stderr=self.slave, + preexec_fn=os.setsid, + universal_newlines=True + ) + self.is_running = True + self.output_buffer.append( + f"[Session {self.session_id}] Started in container {self.container_id[:12]}: " + f"{start_message_cmd} in {self.workspace_dir}") + threading.Thread(target=self._read_output, daemon=True).start() + except Exception as e: + self.output_buffer.append(f"Error starting container session: {str(e)}") + self.is_running = False + return + + # --- Start in CTF --- if self.ctf: - # For CTF environments self.is_running = True self.output_buffer.append( f"[Session { @@ -52,81 +139,100 @@ class ShellSession: # pylint: disable=too-many-instance-attributes output = self.ctf.get_shell(self.command) self.output_buffer.append(output) except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error: {str(e)}") - self.is_running = False + self.output_buffer.append(f"Error executing CTF command: {str(e)}") + self.is_running = False return - # For local environment + # --- Start Locally (Host) --- try: - # Create a pseudo-terminal self.master, self.slave = pty.openpty() - - # Start the process self.process = subprocess.Popen( # pylint: disable=subprocess-popen-preexec-fn, consider-using-with # noqa: E501 - self.command, - shell=True, # nosec B602 + self.command, + shell=True, # nosec B602 stdin=self.slave, stdout=self.slave, stderr=self.slave, - preexec_fn=os.setsid, # Create a new process group + cwd=self.workspace_dir, + preexec_fn=os.setsid, universal_newlines=True ) - self.is_running = True self.output_buffer.append( f"[Session { self.session_id}] Started: { self.command}") - # Start a thread to read output threading.Thread(target=self._read_output, daemon=True).start() except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error starting session: {str(e)}") + self.output_buffer.append(f"Error starting local session: {str(e)}") self.is_running = False def _read_output(self): """Read output from the process""" try: - while self.is_running: + while self.is_running and self.master is not None: try: + # Check if process has exited before reading + if self.process and self.process.poll() is not None: + self.is_running = False + break + # Read the output output = os.read(self.master, 1024).decode() if output: self.output_buffer.append(output) self.last_activity = time.time() - except OSError: - # No data available or terminal closed - time.sleep(0.1) - if not self.is_process_running(): + else: self.is_running = False break - except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error reading output: {str(e)}") + except Exception as e: + self.output_buffer.append(f"Error reading output buffer: {str(read_err)}") + self.is_running = False + break + # Add a small sleep to prevent busy-waiting if no output + if is_process_running(self): + time.sleep(0.05) + except Exception as e: + self.output_buffer.append(f"Error in read_output loop: {str(e)}") self.is_running = False + def is_process_running(self): """Check if the process is still running""" + # For CTF or container + if self.container_id or self.ctf: + return self.is_running + # For local host if not self.process: return False return self.process.poll() is None def send_input(self, input_data): - """Send input to the process""" - if not self.is_running: - return "Session is not running" + """Send input to the process (local or container)""" + if not self.is_running: # For CTF or container + if self.process and self.process.poll() is None: + self.is_running = True + else: # For local host + return "Session is not running" try: + # --- Send to CTF --- if self.ctf: - # For CTF environments output = self.ctf.get_shell(input_data) self.output_buffer.append(output) return "Input sent to CTF session" - # For local environment - input_data = input_data.rstrip() + "\n" - os.write(self.master, input_data.encode()) - self.last_activity = time.time() - return "Input sent to session" + # --- Send to Local or Container PTY --- + if self.master is not None: + input_data_bytes = (input_data.rstrip() + "\n").encode() + bytes_written = os.write(self.master, input_data_bytes) + if bytes_written != len(input_data_bytes): + self.output_buffer.append(f"[Session {self.session_id}] Warning: Partial input write.") + self.last_activity = time.time() + return "Input sent to session" + else: + return "Session PTY not available for input" except Exception as e: # pylint: disable=broad-except + self.output_buffer.append(f"Error sending input: {str(e)}") return f"Error sending input: {str(e)}" def get_output(self, clear=True): @@ -138,8 +244,12 @@ class ShellSession: # pylint: disable=too-many-instance-attributes def terminate(self): """Terminate the session""" + session_id_short = self.session_id[:8] if not self.is_running: - return "Session already terminated" + if self.process and self.process.poll() is None: + pass # Process is running, proceed with termination + else: + return f"Session {session_id_short} already terminated or finished." try: self.is_running = False @@ -147,28 +257,64 @@ class ShellSession: # pylint: disable=too-many-instance-attributes if self.process: # Try to terminate the process group try: - os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) - except BaseException: # pylint: disable=bare-except,broad-except # noqa: E501 - # If that fails, try to terminate just the process - self.process.terminate() + os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) + except ProcessLookupError: + pass # Process already gone + except subprocess.TimeoutExpired: + print(color(f"Session {session_id_short} did not terminate gracefully, sending SIGKILL...", fg="yellow")) # noqa E501 + try: + if pgid: + os.killpg(pgid, signal.SIGKILL) # Force kill + else: + self.process.kill() + except ProcessLookupError: + pass # Already gone + except Exception as kill_err: + termination_message = f" (Error during SIGKILL: {kill_err})" + except Exception as term_err: # Catch other errors during SIGTERM + termination_message = f" (Error during SIGTERM: {term_err})" + try: + self.process.kill() + except Exception: pass # Ignore nested errors - # Clean up resources - if self.master: - os.close(self.master) - if self.slave: - os.close(self.slave) - return f"Session {self.session_id} terminated" + # Final check + if self.process.poll() is None: + print(color(f"Session {session_id_short} process {self.process.pid} may still be running after termination attempts.", fg="red")) # noqa E501 + termination_message += " (Warning: Process may still be running)" + + + # Clean up PTY resources if they exist + if self.master: + try: os.close(self.master) + except OSError: pass + self.master = None + if self.slave: + try: os.close(self.slave) + except OSError: pass + self.slave = None + + return termination_message or f"Session {self.session_id} terminated" except Exception as e: # pylint: disable=broad-except - return f"Error terminating session: {str(e)}" + return f"Error terminating session {session_id_short}: {str(e)}" -def create_shell_session(command, ctf=None): - """Create a new shell session""" - session = ShellSession(command, ctf=ctf) +def create_shell_session(command, ctf=None, container_id=None, **kwargs): + """Create a new shell session in the correct workspace/environment.""" + if container_id: + session = ShellSession(command, ctf=ctf, container_id=container_id) + else: + workspace_dir = _get_workspace_dir() + session = ShellSession(command, ctf=ctf, workspace_dir=workspace_dir) + session.start() - ACTIVE_SESSIONS[session.session_id] = session - return session.session_id + if session.is_running or (ctf and not session.is_running): + ACTIVE_SESSIONS[session.session_id] = session + return session.session_id + else: + error_msg = session.get_output(clear=True) + print(color(f"Failed to start session: {error_msg}", fg="red")) + return f"Failed to start session: {error_msg}" def list_shell_sessions(): @@ -212,47 +358,100 @@ def get_session_output(session_id, clear=True): def terminate_session(session_id): """Terminate a specific session""" if session_id not in ACTIVE_SESSIONS: - return f"Session {session_id} not found" + return f"Session {session_id} not found or already terminated." session = ACTIVE_SESSIONS[session_id] result = session.terminate() - del ACTIVE_SESSIONS[session_id] + if session_id in ACTIVE_SESSIONS: + del ACTIVE_SESSIONS[session_id] return result -def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None): +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}" + original_cmd_for_msg = command # For logging + context_msg = f"(ctf:{target_dir})" try: - # Ensure the command is executed in a shell that supports command - # chaining - output = ctf.get_shell(command, timeout=timeout) - # exploit_logger.log_ok() - + output = ctf.get_shell(full_command, timeout=timeout) if stdout: - print("\033[32m" + output + "\033[0m") - return output # output if output else result.stder + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + return output except Exception as e: # pylint: disable=broad-except - print(color(f"Error executing CTF command: {e}", fg="red")) - # exploit_logger.log_error(str(e)) - return f"Error executing CTF command: {str(e)}" + error_msg = f"Error executing CTF command '{original_cmd_for_msg}' in '{target_dir}': {e}" # noqa E501 + print(color(error_msg, fg="red")) + return error_msg + +def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): + """Runs command via SSH. Assumes SSH agent or passwordless setup unless sshpass is used externally.""" # noqa E501 + ssh_user = os.environ.get('SSH_USER') + ssh_host = os.environ.get('SSH_HOST') + ssh_pass = os.environ.get('SSH_PASS') + remote_command = command + original_cmd_for_msg = command + context_msg = f"({ssh_user}@{ssh_host})" + + # Construct base SSH command list + if ssh_pass: + ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", f"{ssh_user}@{ssh_host}"] # noqa E501 + else: + ssh_cmd_list = ["ssh", f"{ssh_user}@{ssh_host}"] + ssh_cmd_list.append(remote_command) + + try: + # Use subprocess.run with list of args for better security than shell=True + result = subprocess.run( + ssh_cmd_list, + capture_output=True, + text=True, + check=False, # Don't raise exception on non-zero exit code + timeout=timeout + ) + output = result.stdout if result.stdout else result.stderr + if stdout: + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + # Return combined output, potentially including errors + return output.strip() + except subprocess.TimeoutExpired as e: + error_output = e.stdout if e.stdout else str(e) + timeout_msg = f"Timeout executing SSH command: {error_output}" + if stdout: + print(f"\033[33m{context_msg} $ {original_cmd_for_msg}\nTIMEOUT\n{error_output}\033[0m") # noqa E501 + return timeout_msg + except FileNotFoundError: + # Handle case where ssh or sshpass isn't installed + error_msg = f"'sshpass' or 'ssh' command not found. Ensure they are installed and in PATH." # noqa E501 + print(color(error_msg, fg="red")) + return error_msg + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing SSH command '{original_cmd_for_msg}' on {ssh_host}: {e}" # noqa E501 + print(color(error_msg, fg="red")) + return error_msg -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): + """Runs command locally in the specified workspace_dir.""" # If streaming is enabled and we have a call_id if stream and call_id: - return _run_local_streamed(command, call_id, timeout, tool_name) + return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir) + target_dir = workspace_dir or _get_workspace_dir() + original_cmd_for_msg = command # For logging + context_msg = f"(local:{target_dir})" try: - # nosec B602 - shell=True is required for command chaining result = subprocess.run( command, shell=True, # nosec B602 capture_output=True, text=True, - check=False, - timeout=timeout) + check=False, + timeout=timeout, + cwd=target_dir + ) output = result.stdout if result.stdout else result.stderr if stdout: - print("\033[32m" + output + "\033[0m") + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 # Skip passing output to cli_print_tool_output when CAI_STREAM=true # This prevents duplicate output in streaming mode @@ -261,20 +460,21 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # Optional: Add cli_print_tool_output call here if needed for non-streaming pass - return output + return output.strip() except subprocess.TimeoutExpired as e: - error_output = e.stdout.decode() if e.stdout else str(e) + error_output = e.stdout if e.stdout else str(e) if stdout: print("\033[32m" + error_output + "\033[0m") - return error_output + return error_output except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing local command: {e}" - print(color(error_msg, fg="red")) - return error_msg + error_msg = f"Error executing local command: {e}" + print(color(error_msg, fg="red")) + return error_msg -def _run_local_streamed(command, call_id, timeout=100, tool_name=None): - """Run a local command with streaming output to the Tool output panel""" +def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None): + """Run a local command with streaming output to the Tool output panel.""" + target_dir = workspace_dir or _get_workspace_dir() try: # Try to import Rich for nice display try: @@ -298,7 +498,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, - bufsize=1 + bufsize=1, + cwd=target_dir # Set CWD for local process ) # If tool_name is not provided, derive it from the command @@ -328,7 +529,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): header.append(")", style="yellow") tool_time = 0 start_time = time.time() - total_time = time.time() - START_TIME + total_time = time.time() - START_TIME timing_info = [] if total_time: timing_info.append(f"Total: {format_time(total_time)}") @@ -505,7 +706,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg async_mode=False, session_id=None, timeout=100, stream=False, call_id=None, tool_name=None): """ - Run command either in CTF container or on the local attacker machine + Run command in the appropriate environment (Docker, CTF, SSH, Local) + and workspace. Args: command: The command to execute @@ -520,34 +722,168 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg If None, the tool name will be derived from the command. Returns: - str: Command output, status message, or session ID + str: Command output, status message, or session ID. """ # If session_id is provided, send command to that session if session_id: if session_id not in ACTIVE_SESSIONS: return f"Session {session_id} not found" - - result = send_to_session(session_id, command) + session = ACTIVE_SESSIONS[session_id] + result = session.send_input(command) # Send the raw command string if stdout: output = get_session_output(session_id, clear=False) - print("\033[32m" + output + "\033[0m") - return result - - # If async_mode, create a new session - if async_mode: - session_id = create_shell_session(command, ctf) - if stdout: - # Wait a moment for initial output - time.sleep(0.5) - output = get_session_output(session_id, clear=False) - print("\033[32m" + output + "\033[0m") - return f"Created session {session_id}. Use this ID to interact with the session." + env_type = "Local" + if session.container_id: + env_type = f"Container({session.container_id[:12]})" + elif session.ctf: + env_type = "CTF" + print(f"\033[32m(Session {session_id} in {env_type}:{session.workspace_dir}) >> {command}\n{output}\033[0m") # noqa E501 + return result # Return the result of sending input ("Input sent..." or error) # Generate a call_id if we're streaming and one wasn't provided if stream and not call_id: call_id = str(uuid.uuid4())[:8] - - # Otherwise, run command normally + + # 2. Determine Execution Environment (Container > CTF > SSH > Local) + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + 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: + container_id = active_container + container_workspace = _get_container_workspace_path() + context_msg = f"(docker:{container_id[:12]}:{container_workspace})" + + # Handle Async Session Creation in Container + if async_mode: + # Create a session specifically for the container environment + new_session_id = create_shell_session(command, container_id=container_id) # noqa E501 + if "Failed" in new_session_id: # Check if session creation failed + return new_session_id + if stdout: + # Wait a moment for initial output + time.sleep(0.2) + output = get_session_output(new_session_id, clear=False) + print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501 + return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # noqa E501 + + # Handle Streaming Container Execution - not yet implemented for containers + if stream: + # For now, display that streaming isn't supported for containers + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "container": container_id[:12]} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for container execution. Running normally...", + call_id=call_id + ) + + # Handle Synchronous Execution in Container + try: + # Ensure container workspace exists (best effort) + # Consider moving this to workspace set/container activation + mkdir_cmd = ["docker", "exec", container_id, "mkdir", "-p", container_workspace] # noqa E501 + subprocess.run(mkdir_cmd, capture_output=True, text=True, check=False, timeout=10) # noqa E501 + + # Construct the docker exec command with workspace context + cmd_list = [ + "docker", "exec", + "-w", container_workspace, # Set working directory + container_id, + "sh", "-c", command # Execute command via shell + ] + result = subprocess.run( + cmd_list, + capture_output=True, + text=True, + check=False, # Don't raise exception on non-zero exit + timeout=timeout + ) + + output = result.stdout if result.stdout else result.stderr + output = output.strip() # Clean trailing newline + + if stdout: + print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501 + + # Check if command failed specifically because container isn't running + if result.returncode != 0 and "is not running" in result.stderr: + print(color(f"{context_msg} Container is not running. Attempting execution on host instead.", fg="yellow")) # noqa E501 + # Fallback to local execution, preserving workspace context + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + return output # Return combined stdout/stderr + + except subprocess.TimeoutExpired: + timeout_msg = "Timeout executing command in container." + if stdout: + print(f"\033[33m{context_msg} $ {command}\nTIMEOUT\033[0m") # noqa E501 + print(color("Attempting execution on host instead.", fg="yellow")) + # Fallback to local execution on timeout + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing command in container: {str(e)}" + print(color(f"{context_msg} {error_msg}", fg="red")) + print(color("Attempting execution on host instead.", fg="yellow")) + # Fallback to local execution on other errors + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + # --- CTF Execution --- if ctf: - return _run_ctf(ctf, command, stdout, timeout, stream, call_id) - return _run_local(command, stdout, timeout, stream, call_id, tool_name) + # Handling streaming for CTF - not fully implemented yet + if stream: + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "ctf": True} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for CTF execution. Running normally...", + call_id=call_id + ) + + # _run_ctf handles workspace internally using _get_workspace_dir() default + return _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir + + # --- SSH Execution --- + if is_ssh_env: + # Async for SSH would require session management via SSH client features + if async_mode: + return "Async mode not fully supported for SSH environment via this function yet." + + # Handling streaming for SSH - not fully implemented yet + if stream: + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "ssh": True} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for SSH execution. Running normally...", + call_id=call_id + ) + + # _run_ssh handles command execution, workspace is relative to remote home + return _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + + # --- Local Execution (Default Fallback) --- + # Let _run_local handle determining the host workspace + # Handle Async Session Creation Locally + if async_mode: + # create_shell_session uses _get_workspace_dir() when container_id is None + new_session_id = create_shell_session(command) + if isinstance(new_session_id, str) and "Failed" in new_session_id: # Check failure + return new_session_id + # Retrieve the actual workspace dir the session is using + session = ACTIVE_SESSIONS.get(new_session_id) + actual_workspace = session.workspace_dir if session else "unknown" + if stdout: + time.sleep(0.2) # Allow session buffer to populate + output = get_session_output(new_session_id, clear=False) + print(f"\033[32m(Started Session {new_session_id} in local:{actual_workspace})\n{output}\033[0m") + return f"Started async session {new_session_id} locally. Use this ID to interact." + + # Handle Synchronous Execution Locally using _run_local default with streaming support + return _run_local(command, stdout, timeout, stream, call_id, tool_name, None)