sp-hydra-veil-core/core/utils/run_commands.py

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=stdout)
        else:
            # Try stderr first, fallback to stdout for error parsing
            error_output = 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))

PERMISSION_ERRORS = [
    "No permissions",
    "no permissions",
    "password is required",
    "sudo: a password is required"
]

def parse_errors(error_output: str) -> ResultError:
    logger.error(f"error_output is {error_output}")

    # Permission error
    for each_permission_error in PERMISSION_ERRORS:
        if each_permission_error in error_output:
            return ResultError.PERMISSION

    # Interface missing
    if "No such device" in error_output:
        return ResultError.INTERFACE

    elif "command not found" in error_output:
        return ResultError.MISSING_DEPENDENCY

    else:
        return ResultError.UNKNOWN