149 lines
4.6 KiB
Python
149 lines
4.6 KiB
Python
from core.models.Result import Result, ResultError
|
|
from core.errors.logger import logger
|
|
|
|
import subprocess
|
|
import shlex
|
|
import select
|
|
import time
|
|
import fcntl
|
|
import os
|
|
|
|
|
|
# Global
|
|
_bash_session = None
|
|
|
|
def init_terminal():
|
|
"""Initialize the global bash terminal session."""
|
|
global _bash_session
|
|
env = {**os.environ, "SUDO_ASKPASS": "/bin/false"}
|
|
|
|
_bash_session = subprocess.Popen(
|
|
['bash', '--noprofile', '--norc'],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
bufsize=1,
|
|
env=env
|
|
)
|
|
|
|
# Make stdout non-blocking
|
|
flags = fcntl.fcntl(_bash_session.stdout, fcntl.F_GETFL)
|
|
fcntl.fcntl(_bash_session.stdout, fcntl.F_SETFL, flags | os.O_NONBLOCK)
|
|
|
|
logger.info("Bash terminal initialized")
|
|
|
|
|
|
def get_terminal():
|
|
"""Return the global bash terminal or raise RuntimeError."""
|
|
if _bash_session is None:
|
|
raise RuntimeError("Terminal not initialized. Call init_terminal() first.")
|
|
return _bash_session
|
|
|
|
def close_terminal():
|
|
"""Close the global bash terminal session."""
|
|
global _bash_session
|
|
if _bash_session is not None:
|
|
_bash_session.terminate()
|
|
_bash_session.wait()
|
|
_bash_session = None
|
|
logger.info("Bash terminal closed")
|
|
|
|
|
|
|
|
def extract_returncode(line: str, marker: str) -> tuple[int, str]:
|
|
"""Extract returncode and text before marker from a completion line."""
|
|
before, after = line.split(marker, 1)
|
|
returncode = int(after.strip())
|
|
return returncode, before
|
|
|
|
|
|
def _run_command_via_terminal(
|
|
command: list,
|
|
timeout: float | None = None
|
|
) -> tuple[int, str]:
|
|
"""
|
|
Execute a command via persistent bash session.
|
|
Returns (returncode, stdout, stderr)
|
|
"""
|
|
try:
|
|
terminal = get_terminal()
|
|
except RuntimeError:
|
|
init_terminal()
|
|
terminal = get_terminal()
|
|
|
|
# Convert command list to properly quoted bash string
|
|
cmd_str = ' '.join(shlex.quote(arg) for arg in command)
|
|
|
|
# Use a marker to detect command completion
|
|
# Redirect stderr to stdout and append marker with exit code
|
|
marker = "::__CMD_DONE__::"
|
|
# wrapped_cmd = f"{cmd_str} 2>&1; echo \"{marker}$?\"\n"
|
|
wrapped_cmd = f"{cmd_str} 2>&1; echo {marker}$?\n"
|
|
|
|
|
|
terminal.stdin.write(wrapped_cmd)
|
|
terminal.stdin.flush()
|
|
|
|
output_lines = []
|
|
returncode = 1
|
|
|
|
timeout = timeout or float('inf')
|
|
start_time = time.time()
|
|
|
|
try:
|
|
while time.time() - start_time < timeout:
|
|
line = terminal.stdout.readline() # non-blocking due to init fcntl
|
|
if not line:
|
|
time.sleep(0.02)
|
|
continue
|
|
|
|
if marker in line:
|
|
returncode, before = extract_returncode(line, marker)
|
|
if before:
|
|
output_lines.append(before)
|
|
return returncode, ''.join(output_lines)
|
|
|
|
output_lines.append(line)
|
|
|
|
# If we exit the loop without returning first, timeout occurred
|
|
logger.error(f"Command timed out after {timeout}s")
|
|
close_terminal()
|
|
raise TimeoutError(f"Command exceeded {timeout}s limit")
|
|
except Exception as e:
|
|
logger.error(f"We hit an error in the loop of reading, so we'll close the terminal, then the caller can deal with {e}")
|
|
close_terminal()
|
|
raise
|
|
|
|
|
|
def run_generic_command(
|
|
command: list,
|
|
human_readable_goal: str,
|
|
timeout: float | None = None
|
|
) -> Result:
|
|
|
|
output_data = None
|
|
|
|
# Add -n flag to sudo for non-interactive mode (no prompts)
|
|
if command and command[0] == "sudo":
|
|
if "-n" not in command:
|
|
command = [command[0], "-n"] + command[1:]
|
|
|
|
try:
|
|
returncode, stdout = _run_command_via_terminal(command, timeout)
|
|
output_data = stdout
|
|
|
|
if returncode == 0:
|
|
logger.info(f"{human_readable_goal} was successful")
|
|
return Result(valid=True, data=output_data)
|
|
else:
|
|
# Try stderr first, fallback to stdout for error parsing
|
|
error_output = (stderr or stdout).strip()
|
|
error_enum = parse_errors(error_output)
|
|
logger.error(f"{human_readable_goal} Failed, error type is {error_enum}.")
|
|
return Result(valid=False, error_type=error_enum, goal=human_readable_goal, message=error_output)
|
|
|
|
except TimeoutError as e:
|
|
return Result(valid=False, error_type=ResultError.TIMEOUT, goal=human_readable_goal, data=output_data, message=f"Command timed out {e}")
|
|
except Exception as e:
|
|
return Result(valid=False, error_type=ResultError.UNKNOWN, goal=human_readable_goal, data=output_data, message=str(e))
|