mirror of https://github.com/aliasrobotics/cai.git
add workspaces + some virt in common
This commit is contained in:
Mery-Sanz
parent
0cfcf2467d
commit
ed68c58d4c
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <name>[/bold] - "
|
||||
"Set the current workspace name")
|
||||
console.print(
|
||||
" [bold]/workspace ls[/bold] - "
|
||||
"List files in the workspace")
|
||||
console.print(
|
||||
" [bold]/workspace exec <cmd>[/bold] - "
|
||||
"Execute a command in the workspace")
|
||||
|
||||
if active_container:
|
||||
console.print(
|
||||
" [bold]/workspace copy <src> <dst>[/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 <workspace_name>[/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 <source> <destination>")
|
||||
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())
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue