Shift things around

source/src/utils/execute.cpp

@@ -0,0 +1,179 @@
+#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"
+
+namespace dropshell
+{
+    bool EXITSTATUSCHECK(int ret)
+    {
+        return (ret != -1 && WIFEXITED(ret) && (WEXITSTATUS(ret) == 0)); // ret is -1 if the command failed to execute.
+    }
+
+    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);
+        }
+    }
+
+    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);
+    }
+
+    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 = system(full_cmd.c_str());
+
+        bool ok = EXITSTATUSCHECK(ret);
+        if (!ok && !silent)
+        {
+            std::cerr << "Error: Failed to execute command: " << std::endl;
+            std::cerr << full_cmd << std::endl;
+        }
+        return ok;
+    }
+
+    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;
+    }
+
+    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;
+    }
+
+    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