dropshell/source/src/utils/execute.cpp

#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <vector>
#include <iostream>
#include <string>
#include <cstdlib>
#include <sstream>
#include "utils/assert.hpp"

#include "execute.hpp"
#include "utils/utils.hpp"
#include "utils/b64ed.hpp"
#include "config.hpp"
#include "utils/directories.hpp"
#include "utils/output.hpp"

namespace dropshell
{
    bool EXITSTATUSCHECK(int ret)
    {
        return (ret != -1 && WIFEXITED(ret) && (WEXITSTATUS(ret) == 0)); // ret is -1 if the command failed to execute.
    }


    // ----------------------------------------------------------------------------------------------------------
    // execute_local_command_interactive
    // ----------------------------------------------------------------------------------------------------------
    bool execute_local_command_interactive(const sCommand &command)
    {
        if (command.get_command_to_run().empty())
            return false;
        std::string full_command = command.construct_cmd(localpath::agent()); // Get the command string

        pid_t pid = fork();

        if (pid == -1)
        {
            // Fork failed
            perror("fork failed");
            return false;
        }
        else if (pid == 0)
        {
            int rval = system(full_command.c_str());
            exit(rval);
        }
        else
        {
            // Parent process
            int ret;
            // Wait for the child process to complete
            waitpid(pid, &ret, 0);

            return EXITSTATUSCHECK(ret);
        }
    }

    // ----------------------------------------------------------------------------------------------------------
    // execute_local_command_and_capture_output
    // ----------------------------------------------------------------------------------------------------------
    bool execute_local_command_and_capture_output(const sCommand &command, std::string *output)
    {
        ASSERT(output != nullptr, "Output string must be provided");
        if (command.get_command_to_run().empty())
            return false;
        std::string full_cmd = command.construct_cmd(localpath::agent()) + " 2>&1";
        FILE *pipe = popen(full_cmd.c_str(), "r");
        if (!pipe)
        {
            return false;
        }
        char buffer[128];
        while (fgets(buffer, sizeof(buffer), pipe) != nullptr)
        {
            (*output) += buffer;
        }
        int ret = pclose(pipe);
        return EXITSTATUSCHECK(ret);
    }

    // ----------------------------------------------------------------------------------------------------------
    // execute_local_command
    // ----------------------------------------------------------------------------------------------------------
    bool execute_local_command(std::string command, std::string *output, cMode mode)
    {
        return execute_local_command("", command, {}, output, mode);
    }

    bool execute_local_command(std::string directory_to_run_in, std::string command_to_run, const std::map<std::string, std::string> &env_vars, std::string *output, cMode mode)
    {
        sCommand command(directory_to_run_in, command_to_run, env_vars);

        if (hasFlag(mode, cMode::Interactive))
        {
            ASSERT(!hasFlag(mode, cMode::CaptureOutput), "Interactive mode and capture output mode cannot be used together");
            ASSERT(output == nullptr, "Interactive mode and an output string cannot be used together");

            return execute_local_command_interactive(command);
        }

        if (hasFlag(mode, cMode::CaptureOutput))
        {
            ASSERT(output != nullptr, "Capture output mode requires an output string to be provided");
            ASSERT(!hasFlag(mode, cMode::Silent), "Silent mode is not allowed with capture output mode");

            return execute_local_command_and_capture_output(command, output);
        }

        if (command.get_command_to_run().empty())
            return false;
        bool silent = hasFlag(mode, cMode::Silent);
        std::string full_cmd = command.construct_cmd(localpath::agent()) + " 2>&1" + (silent ? " > /dev/null" : "");

        int ret=0;
        {
            SwitchColour sc(sColour::DEBUG, std::cerr);
            ret = system(full_cmd.c_str());
        }

        bool ok = EXITSTATUSCHECK(ret);
        if (!ok && !silent)
        {
            PrintError("Error: Failed to execute command: ");
            PrintError(full_cmd);
        }
        return ok;
    }

    // ----------------------------------------------------------------------------------------------------------
    // execute_ssh_command
    // ----------------------------------------------------------------------------------------------------------
    bool execute_ssh_command(const sSSHInfo &ssh_info, const sCommand &remote_command, cMode mode, std::string *output)
    {
        if (remote_command.get_command_to_run().empty())
            return false;

        ASSERT(!(hasFlag(mode, cMode::CaptureOutput) && output == nullptr), "Capture output mode must be used with an output string");

        std::stringstream ssh_cmd;
        ssh_cmd << "ssh -p " << ssh_info.port << " " << (hasFlag(mode, cMode::Interactive) ? "-tt " : "")
                << ssh_info.user << "@" << ssh_info.host;

        std::string remote_agent_path = remotepath::agent(ssh_info.server_ID);

        bool rval = execute_local_command(
            "",                                                                          // directory to run in
            ssh_cmd.str() + " " + remote_command.construct_cmd(remote_agent_path),       // local command to run
            {},                                                                          // environment variables
            output,                                                                      // output string
            mode                                                                         // mode
        );

        if (!rval && !hasFlag(mode, cMode::Silent))
        {
            std::cerr << std::endl
                      << std::endl;
            std::cerr << "Error: Failed to execute ssh command:" << std::endl;
            std::cerr << "\033[90m" << ssh_cmd.str() + " " + remote_command.construct_cmd(remote_agent_path) << "\033[0m" << std::endl;
            std::cerr << std::endl
                      << std::endl;
        }
        return rval;
    }

    // ----------------------------------------------------------------------------------------------------------
    // makesafecmd
    // ----------------------------------------------------------------------------------------------------------
    std::string sCommand::makesafecmd(std::string agent_path, const std::string &command) const
    {
        if (command.empty())
            return "";
        std::string encoded = base64_encode(dequote(trim(command)));
        std::string commandstr = agent_path + "/bb64 " + encoded;
        return commandstr;
    }

    // ----------------------------------------------------------------------------------------------------------
    // construct_cmd
    // ----------------------------------------------------------------------------------------------------------
    std::string sCommand::construct_cmd(std::string agent_path) const
    {
        if (mCmd.empty())
            return "";

        // need to construct to change directory and set environment variables
        std::string cmdstr;

        if (!mDir.empty())
            cmdstr += "cd " + quote(mDir) + " && ";

        if (!mVars.empty())
            for (const auto &env_var : mVars)
                cmdstr += env_var.first + "=" + quote(dequote(trim(env_var.second))) + " ";

        cmdstr += mCmd;

        if (!agent_path.empty())
        cmdstr = makesafecmd(agent_path, cmdstr);

        return cmdstr;
    }

} // namespace dropshell