251 lines
8.6 KiB
C++
251 lines
8.6 KiB
C++
#include <sys/types.h>
|
|
#include <sys/wait.h>
|
|
#include <unistd.h>
|
|
#include <vector>
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <cstdlib>
|
|
#include <sstream>
|
|
#include <libassert/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(localfile::bb64()); // 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
|
|
// ----------------------------------------------------------------------------------------------------------
|
|
|
|
class fancypinter
|
|
{
|
|
public:
|
|
fancypinter(sColour startColour) : startColour_(startColour), currentColour_(startColour) {}
|
|
|
|
void print_chunk(std::string chunk)
|
|
{
|
|
if (chunk.empty())
|
|
return;
|
|
|
|
if (newline_)
|
|
{
|
|
// sniff the mode... if the string starts with warning or warning: then set mode to WARNING. etc.
|
|
if (chunk.find("warning") == 0)
|
|
currentColour_ = sColour::WARNING;
|
|
else if (chunk.find("error") == 0)
|
|
currentColour_ = sColour::ERROR;
|
|
else if (chunk.find("debug") == 0)
|
|
currentColour_ = sColour::DEBUG;
|
|
else if (chunk.find("info") == 0)
|
|
currentColour_ = sColour::INFO;
|
|
else
|
|
currentColour_ = startColour_;
|
|
}
|
|
colourstream(currentColour_) << chunk;
|
|
newline_ = (chunk[chunk.size() - 1] == '\n');
|
|
}
|
|
|
|
void print(const std::string &buffer)
|
|
{
|
|
size_t start = 0;
|
|
while (start < buffer.size())
|
|
{
|
|
size_t newline_pos = buffer.find('\n', start);
|
|
if (newline_pos == std::string::npos)
|
|
{
|
|
if (start < buffer.size())
|
|
{
|
|
print_chunk(buffer.substr(start));
|
|
}
|
|
break;
|
|
}
|
|
print_chunk(buffer.substr(start, newline_pos - start + 1)); // include the newline
|
|
start = newline_pos + 1;
|
|
}
|
|
}
|
|
|
|
private:
|
|
bool newline_ = true;
|
|
sColour startColour_;
|
|
sColour currentColour_;
|
|
};
|
|
|
|
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(output == nullptr, "Interactive mode and an output string cannot be used together");
|
|
|
|
return execute_local_command_interactive(command);
|
|
}
|
|
|
|
if (command.get_command_to_run().empty())
|
|
return false;
|
|
|
|
bool silent = hasFlag(mode, cMode::Silent);
|
|
|
|
std::string full_cmd;
|
|
if (!hasFlag(mode, cMode::NoBB64))
|
|
full_cmd = command.construct_cmd(localfile::bb64());
|
|
else
|
|
full_cmd = command.construct_cmd("");
|
|
|
|
if (output != nullptr)
|
|
full_cmd += " 2>&1"; // capture both stdout and stderr
|
|
|
|
FILE *pipe = popen(full_cmd.c_str(), "r");
|
|
if (!pipe)
|
|
{
|
|
return false;
|
|
}
|
|
char buffer[128];
|
|
fancypinter fancyprint(sColour::DEBUG);
|
|
while (fgets(buffer, sizeof(buffer), pipe) != nullptr)
|
|
{
|
|
if (output != nullptr)
|
|
(*output) += buffer;
|
|
|
|
if (!silent)
|
|
fancyprint.print(buffer);
|
|
}
|
|
int ret = pclose(pipe);
|
|
return EXITSTATUSCHECK(ret);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------------------------
|
|
// 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;
|
|
|
|
std::stringstream ssh_cmd;
|
|
ssh_cmd << "ssh -p " << ssh_info.get_port() << " " << (hasFlag(mode, cMode::Interactive) ? "-tt " : "")
|
|
<< ssh_info.get_user() << "@" << ssh_info.get_host();
|
|
|
|
std::string remote_bb64_path;
|
|
|
|
if (!hasFlag(mode, cMode::NoBB64))
|
|
remote_bb64_path = remotepath(ssh_info.get_server_ID(), ssh_info.get_user()).agent() + "/bb64";
|
|
|
|
bool rval = execute_local_command(
|
|
"", // local directory to run in
|
|
ssh_cmd.str() + " " + remote_command.construct_cmd(remote_bb64_path), // local command to run
|
|
{}, // environment variables
|
|
output, // output string
|
|
mode // mode
|
|
);
|
|
|
|
if (!rval && !hasFlag(mode, cMode::Silent))
|
|
{
|
|
error << "Error: Failed to execute ssh command:" << std::endl;
|
|
debug << ssh_cmd.str() + " " + remote_command.construct_cmd(remote_bb64_path) << std::endl;
|
|
}
|
|
return rval;
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------------------------
|
|
// makesafecmd
|
|
// ----------------------------------------------------------------------------------------------------------
|
|
std::string sCommand::makesafecmd(std::string bb64path, const std::string &command) const
|
|
{
|
|
if (command.empty())
|
|
return "";
|
|
std::string encoded = base64_encode(dequote(trim(command)));
|
|
std::string commandstr = bb64path + " " + encoded;
|
|
return commandstr;
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------------------------
|
|
// construct_cmd
|
|
// ----------------------------------------------------------------------------------------------------------
|
|
std::string sCommand::construct_cmd(std::string bb64path) const
|
|
{
|
|
if (mCmd.empty())
|
|
return "";
|
|
|
|
// need to construct to change directory and set environment variables
|
|
std::string cmdstr;
|
|
|
|
if (!bb64path.empty())
|
|
{
|
|
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;
|
|
|
|
cmdstr = makesafecmd(bb64path, cmdstr);
|
|
}
|
|
else
|
|
{ // raw! bootstrapping only.
|
|
ASSERT(mVars.empty(), "Bootstrapping command must not have environment variables");
|
|
if (!mDir.empty())
|
|
cmdstr += mDir + "/" + mCmd;
|
|
else
|
|
cmdstr += mCmd;
|
|
}
|
|
|
|
return cmdstr;
|
|
}
|
|
|
|
bool sSSHInfo::valid() const
|
|
{
|
|
if (host.empty() || user.empty() || port.empty() || server_ID.empty() || user_dir.empty())
|
|
return false;
|
|
|
|
if (atoi(port.c_str()) == 0)
|
|
return false;
|
|
|
|
return true;
|
|
}
|
|
|
|
} // namespace dropshell
|