dropshell/runner/src/runner.cpp

#include "runner.h"
#include <iostream>
#include <sstream>
#include <cstdlib>
#include <array>
#include <cstring>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <libssh/libssh.h>
#include <libssh/sftp.h>
#include <pwd.h>

namespace dropshell {

bool Runner::run(const nlohmann::json& run_json) {
    std::string command;
    std::vector<std::string> args;
    std::string working_dir;
    std::map<std::string, std::string> env;
    bool silent, interactive;
    
    if (!parse_json(run_json, command, args, working_dir, env, silent, interactive)) {
        return false;
    }
    
    int exit_code;
    if (run_json.contains("ssh")) {
        exit_code = execute_ssh(run_json["ssh"], command, args, working_dir, env, silent, interactive);
    } else {
        exit_code = execute_local(command, args, working_dir, env, silent, interactive);
    }
    
    return exit_code == 0;
}

bool Runner::run(const nlohmann::json& run_json, std::string& output) {
    std::string command;
    std::vector<std::string> args;
    std::string working_dir;
    std::map<std::string, std::string> env;
    bool silent, interactive;
    
    if (!parse_json(run_json, command, args, working_dir, env, silent, interactive)) {
        return false;
    }
    
    int exit_code;
    if (run_json.contains("ssh")) {
        exit_code = execute_ssh(run_json["ssh"], command, args, working_dir, env, silent, interactive, &output, true);
    } else {
        exit_code = execute_local(command, args, working_dir, env, silent, interactive, &output, true);
    }
    
    return exit_code == 0;
}

bool Runner::parse_json(
    const nlohmann::json& run_json, 
    std::string& command, 
    std::vector<std::string>& args,
    std::string& working_dir,
    std::map<std::string, std::string>& env,
    bool& silent,
    bool& interactive
) {
    try {
        // Command is required
        if (!run_json.contains("command") || !run_json["command"].is_string()) {
            std::cerr << "Error: 'command' field is required and must be a string" << std::endl;
            return false;
        }
        command = run_json["command"];
        
        // Args are optional
        args.clear();
        if (run_json.contains("args")) {
            if (!run_json["args"].is_array()) {
                std::cerr << "Error: 'args' field must be an array" << std::endl;
                return false;
            }
            for (const auto& arg : run_json["args"]) {
                if (!arg.is_string()) {
                    std::cerr << "Error: All arguments must be strings" << std::endl;
                    return false;
                }
                args.push_back(arg);
            }
        }
        
        // Working directory is optional
        working_dir = "";
        if (run_json.contains("working_directory")) {
            if (!run_json["working_directory"].is_string()) {
                std::cerr << "Error: 'working_directory' field must be a string" << std::endl;
                return false;
            }
            working_dir = run_json["working_directory"];
        }
        
        // Environment variables are optional
        env.clear();
        if (run_json.contains("env")) {
            if (!run_json["env"].is_object()) {
                std::cerr << "Error: 'env' field must be an object" << std::endl;
                return false;
            }
            for (auto it = run_json["env"].begin(); it != run_json["env"].end(); ++it) {
                if (!it.value().is_string()) {
                    std::cerr << "Error: All environment variable values must be strings" << std::endl;
                    return false;
                }
                env[it.key()] = it.value();
            }
        }
        
        // Options are optional
        silent = false;
        interactive = false;
        if (run_json.contains("options")) {
            if (!run_json["options"].is_object()) {
                std::cerr << "Error: 'options' field must be an object" << std::endl;
                return false;
            }
            
            if (run_json["options"].contains("silent")) {
                if (!run_json["options"]["silent"].is_boolean()) {
                    std::cerr << "Error: 'silent' option must be a boolean" << std::endl;
                    return false;
                }
                silent = run_json["options"]["silent"];
            }
            
            if (run_json["options"].contains("interactive")) {
                if (!run_json["options"]["interactive"].is_boolean()) {
                    std::cerr << "Error: 'interactive' option must be a boolean" << std::endl;
                    return false;
                }
                interactive = run_json["options"]["interactive"];
            }
        }
        
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Error parsing JSON: " << e.what() << std::endl;
        return false;
    }
}

int Runner::execute_local(
    const std::string& command,
    const std::vector<std::string>& args,
    const std::string& working_dir,
    const std::map<std::string, std::string>& env,
    bool silent,
    bool interactive,
    std::string* output,
    bool capture_output
) {
    int pipefd[2];
    if (capture_output && pipe(pipefd) == -1) {
        std::cerr << "Error creating pipe: " << strerror(errno) << std::endl;
        return -1;
    }
    
    pid_t pid = fork();
    if (pid == -1) {
        std::cerr << "Fork failed: " << strerror(errno) << std::endl;
        if (capture_output) {
            close(pipefd[0]);
            close(pipefd[1]);
        }
        return -1;
    }
    
    if (pid == 0) {
        // Child process
        
        // Set up output redirection if needed
        if (capture_output) {
            close(pipefd[0]);  // Close read end
            dup2(pipefd[1], STDOUT_FILENO);
            dup2(pipefd[1], STDERR_FILENO);
            close(pipefd[1]);
        } else if (silent) {
            int devnull = open("/dev/null", O_WRONLY);
            if (devnull != -1) {
                dup2(devnull, STDOUT_FILENO);
                dup2(devnull, STDERR_FILENO);
                close(devnull);
            }
        }
        
        // Change to working directory if specified
        if (!working_dir.empty()) {
            if (chdir(working_dir.c_str()) != 0) {
                std::cerr << "Error changing to directory " << working_dir << ": " << strerror(errno) << std::endl;
                exit(1);
            }
        }
        
        // Set environment variables
        for (const auto& [key, value] : env) {
            setenv(key.c_str(), value.c_str(), 1);
        }
        
        // Prepare arguments
        std::vector<char*> c_args;
        c_args.push_back(const_cast<char*>(command.c_str()));
        for (const auto& arg : args) {
            c_args.push_back(const_cast<char*>(arg.c_str()));
        }
        c_args.push_back(nullptr);
        
        // Execute command
        execvp(command.c_str(), c_args.data());
        
        // If exec fails
        std::cerr << "exec failed: " << strerror(errno) << std::endl;
        exit(1);
    } else {
        // Parent process
        if (capture_output) {
            close(pipefd[1]);  // Close write end
            
            std::array<char, 4096> buffer;
            std::ostringstream oss;
            
            ssize_t bytes_read;
            while ((bytes_read = read(pipefd[0], buffer.data(), buffer.size() - 1)) > 0) {
                buffer[bytes_read] = '\0';
                oss << buffer.data();
                
                if (!silent) {
                    std::cout << buffer.data();
                }
            }
            
            close(pipefd[0]);
            
            if (output) {
                *output = oss.str();
            }
        }
        
        int status;
        waitpid(pid, &status, 0);
        
        if (WIFEXITED(status)) {
            return WEXITSTATUS(status);
        } else {
            return -1;
        }
    }
}

std::string find_ssh_key_for_user() {
    const char* home_dir = getenv("HOME");
    if (!home_dir) {
        struct passwd* pw = getpwuid(getuid());
        if (pw) {
            home_dir = pw->pw_dir;
        }
    }
    
    if (!home_dir) {
        return "";
    }
    
    // Common SSH key locations
    std::vector<std::string> key_paths = {
        std::string(home_dir) + "/.ssh/id_rsa",
        std::string(home_dir) + "/.ssh/id_ed25519",
        std::string(home_dir) + "/.ssh/id_ecdsa",
        std::string(home_dir) + "/.ssh/id_dsa"
    };
    
    for (const auto& path : key_paths) {
        if (access(path.c_str(), F_OK) == 0) {
            return path;
        }
    }
    
    return "";
}

int Runner::execute_ssh(
    const nlohmann::json& ssh_config,
    const std::string& command,
    const std::vector<std::string>& args,
    const std::string& working_dir,
    const std::map<std::string, std::string>& env,
    bool silent,
    bool interactive,
    std::string* output,
    bool capture_output
) {
    if (!ssh_config.contains("host") || !ssh_config["host"].is_string()) {
        std::cerr << "Error: SSH configuration requires 'host' field" << std::endl;
        return -1;
    }
    std::string host = ssh_config["host"];
    
    int port = 22;
    if (ssh_config.contains("port")) {
        if (ssh_config["port"].is_number_integer()) {
            port = ssh_config["port"];
        } else {
            std::cerr << "Error: SSH 'port' must be an integer" << std::endl;
            return -1;
        }
    }
    
    std::string user = "";
    if (ssh_config.contains("user") && ssh_config["user"].is_string()) {
        user = ssh_config["user"];
    } else {
        // Get current username as default
        char username[256];
        if (getlogin_r(username, sizeof(username)) == 0) {
            user = username;
        } else {
            struct passwd* pw = getpwuid(getuid());
            if (pw) {
                user = pw->pw_name;
            }
        }
    }
    
    std::string key_path = "";
    if (ssh_config.contains("key") && ssh_config["key"].is_string()) {
        std::string key = ssh_config["key"];
        if (key == "auto") {
            key_path = find_ssh_key_for_user();
            if (key_path.empty()) {
                std::cerr << "Error: Could not find SSH key automatically" << std::endl;
                return -1;
            }
        } else {
            key_path = key;
        }
    } else {
        key_path = find_ssh_key_for_user();
    }
    
    // Initialize SSH session
    ssh_session session = ssh_new();
    if (session == nullptr) {
        std::cerr << "Error: Failed to create SSH session" << std::endl;
        return -1;
    }
    
    // Set SSH options
    ssh_options_set(session, SSH_OPTIONS_HOST, host.c_str());
    ssh_options_set(session, SSH_OPTIONS_PORT, &port);
    ssh_options_set(session, SSH_OPTIONS_USER, user.c_str());
    
    // Connect to server
    int rc = ssh_connect(session);
    if (rc != SSH_OK) {
        std::cerr << "Error connecting to " << host << ": " << ssh_get_error(session) << std::endl;
        ssh_free(session);
        return -1;
    }
    
    // Authenticate with key
    if (!key_path.empty()) {
        rc = ssh_userauth_publickey_auto(session, nullptr, key_path.empty() ? nullptr : key_path.c_str());
        if (rc != SSH_AUTH_SUCCESS) {
            std::cerr << "Error authenticating with key: " << ssh_get_error(session) << std::endl;
            ssh_disconnect(session);
            ssh_free(session);
            return -1;
        }
    } else {
        // Try default authentication methods
        rc = ssh_userauth_publickey_auto(session, nullptr, nullptr);
        if (rc != SSH_AUTH_SUCCESS) {
            std::cerr << "Error authenticating: " << ssh_get_error(session) << std::endl;
            ssh_disconnect(session);
            ssh_free(session);
            return -1;
        }
    }
    
    // Prepare command
    std::ostringstream cmd_stream;
    
    // Add environment variables
    for (const auto& [key, value] : env) {
        cmd_stream << "export " << key << "=\"" << value << "\"; ";
    }
    
    // Add cd command if working_directory is specified
    if (!working_dir.empty()) {
        cmd_stream << "cd \"" << working_dir << "\" && ";
    }
    
    // Add command and args
    cmd_stream << command;
    for (const auto& arg : args) {
        cmd_stream << " " << arg;
    }
    
    std::string full_command = cmd_stream.str();
    
    int exit_status = -1;
    ssh_channel channel = ssh_channel_new(session);
    if (channel == nullptr) {
        std::cerr << "Error creating SSH channel: " << ssh_get_error(session) << std::endl;
        ssh_disconnect(session);
        ssh_free(session);
        return -1;
    }
    
    rc = ssh_channel_open_session(channel);
    if (rc != SSH_OK) {
        std::cerr << "Error opening SSH session: " << ssh_get_error(session) << std::endl;
        ssh_channel_free(channel);
        ssh_disconnect(session);
        ssh_free(session);
        return -1;
    }
    
    if (interactive) {
        // Request a pseudo-terminal for interactive commands with specific term type
        const char* term_type = "xterm-256color";
        rc = ssh_channel_request_pty_size(channel, term_type, 80, 24);
        if (rc != SSH_OK) {
            // Fallback to basic PTY request
            std::cerr << "Warning: Could not set PTY with size, falling back to basic PTY" << std::endl;
            rc = ssh_channel_request_pty(channel);
            if (rc != SSH_OK) {
                std::cerr << "Error requesting PTY: " << ssh_get_error(session) << std::endl;
                ssh_channel_close(channel);
                ssh_channel_free(channel);
                ssh_disconnect(session);
                ssh_free(session);
                return -1;
            }
        }
    }
    
    // Execute the command
    rc = ssh_channel_request_exec(channel, full_command.c_str());
    if (rc != SSH_OK) {
        std::cerr << "Error executing command: " << ssh_get_error(session) << std::endl;
        ssh_channel_close(channel);
        ssh_channel_free(channel);
        ssh_disconnect(session);
        ssh_free(session);
        return -1;
    }
    
    // Read command output
    char buffer[4096];
    int nbytes;
    std::ostringstream oss;
    
    // Read from stdout
    while ((nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0)) > 0) {
        if (capture_output) {
            oss.write(buffer, nbytes);
        }
        
        if (!silent) {
            std::cout.write(buffer, nbytes);
        }
    }
    
    // Read from stderr if needed
    while ((nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 1)) > 0) {
        if (capture_output) {
            oss.write(buffer, nbytes);
        }
        
        if (!silent) {
            std::cerr.write(buffer, nbytes);
        }
    }
    
    if (capture_output && output) {
        *output = oss.str();
    }
    
    // Get exit status
    ssh_channel_send_eof(channel);
    exit_status = ssh_channel_get_exit_status(channel);
    
    // Clean up
    ssh_channel_close(channel);
    ssh_channel_free(channel);
    ssh_disconnect(session);
    ssh_free(session);
    
    return exit_status;
}

} // namespace dropshell