281 lines
9.2 KiB
C++
281 lines
9.2 KiB
C++
#include "runner.h"
|
|
#include <cstdlib>
|
|
#include <sstream>
|
|
#include <iostream>
|
|
#include <unistd.h>
|
|
#include <sys/wait.h>
|
|
#include <fcntl.h>
|
|
#include <libssh/libssh.h>
|
|
#include <libssh/callbacks.h>
|
|
#include <termios.h>
|
|
#include <sys/select.h>
|
|
|
|
namespace runner {
|
|
|
|
namespace {
|
|
|
|
ssh_session ssh_connect_and_auth(const sSSHInfo* sshinfo, const std::map<std::string, std::string>& env, std::string* error) {
|
|
ssh_session session = ssh_new();
|
|
if (!session) {
|
|
if (error) *error = "Failed to create SSH session.";
|
|
return nullptr;
|
|
}
|
|
ssh_options_set(session, SSH_OPTIONS_HOST, sshinfo->host.c_str());
|
|
if (!sshinfo->port.empty()) {
|
|
int port = std::stoi(sshinfo->port);
|
|
ssh_options_set(session, SSH_OPTIONS_PORT, &port);
|
|
}
|
|
if (!sshinfo->user.empty()) {
|
|
ssh_options_set(session, SSH_OPTIONS_USER, sshinfo->user.c_str());
|
|
}
|
|
int rc = ssh_connect(session);
|
|
if (rc != SSH_OK) {
|
|
if (error) *error = std::string("SSH connection failed: ") + ssh_get_error(session);
|
|
ssh_free(session);
|
|
return nullptr;
|
|
}
|
|
rc = ssh_userauth_publickey_auto(session, nullptr, nullptr);
|
|
if (rc != SSH_AUTH_SUCCESS) {
|
|
auto it = env.find("SSHPASS");
|
|
if (it != env.end()) {
|
|
rc = ssh_userauth_password(session, nullptr, it->second.c_str());
|
|
}
|
|
}
|
|
if (rc != SSH_AUTH_SUCCESS) {
|
|
if (error) *error = std::string("SSH authentication failed: ") + ssh_get_error(session);
|
|
ssh_disconnect(session);
|
|
ssh_free(session);
|
|
return nullptr;
|
|
}
|
|
return session;
|
|
}
|
|
|
|
std::string ssh_build_remote_command(const std::string& command, const std::vector<std::string>& args, const std::string& working_dir, const std::map<std::string, std::string>& env) {
|
|
std::ostringstream remote_cmd;
|
|
for (const auto& kv : env) {
|
|
if (kv.first == "SSHPASS") continue;
|
|
remote_cmd << kv.first << "='" << kv.second << "' ";
|
|
}
|
|
if (!working_dir.empty()) {
|
|
remote_cmd << "cd '" << working_dir << "' && ";
|
|
}
|
|
remote_cmd << command;
|
|
for (const auto& arg : args) {
|
|
remote_cmd << " '" << arg << "'";
|
|
}
|
|
return remote_cmd.str();
|
|
}
|
|
|
|
int ssh_interactive_shell_session(ssh_session session, ssh_channel channel, const std::string& remote_cmd_str, const std::string& command, std::string* output) {
|
|
int rc = ssh_channel_request_pty(channel);
|
|
if (rc != SSH_OK) {
|
|
if (output) *output = std::string("Failed to request pty: ") + ssh_get_error(session);
|
|
return -1;
|
|
}
|
|
rc = ssh_channel_request_shell(channel);
|
|
if (rc != SSH_OK) {
|
|
if (output) *output = std::string("Failed to request shell: ") + ssh_get_error(session);
|
|
return -1;
|
|
}
|
|
struct termios orig_termios, raw_termios;
|
|
tcgetattr(STDIN_FILENO, &orig_termios);
|
|
raw_termios = orig_termios;
|
|
cfmakeraw(&raw_termios);
|
|
tcsetattr(STDIN_FILENO, TCSANOW, &raw_termios);
|
|
if (!command.empty()) {
|
|
ssh_channel_write(channel, remote_cmd_str.c_str(), remote_cmd_str.size());
|
|
ssh_channel_write(channel, "\n", 1);
|
|
}
|
|
int maxfd = STDIN_FILENO > STDOUT_FILENO ? STDIN_FILENO : STDOUT_FILENO;
|
|
maxfd = maxfd > ssh_get_fd(session) ? maxfd : ssh_get_fd(session);
|
|
char buffer[4096];
|
|
bool done = false;
|
|
while (!done) {
|
|
fd_set fds_read;
|
|
FD_ZERO(&fds_read);
|
|
FD_SET(STDIN_FILENO, &fds_read);
|
|
FD_SET(ssh_get_fd(session), &fds_read);
|
|
int ret = select(maxfd + 1, &fds_read, nullptr, nullptr, nullptr);
|
|
if (ret < 0) break;
|
|
if (FD_ISSET(STDIN_FILENO, &fds_read)) {
|
|
ssize_t n = read(STDIN_FILENO, buffer, sizeof(buffer));
|
|
if (n > 0) {
|
|
ssh_channel_write(channel, buffer, n);
|
|
} else {
|
|
ssh_channel_send_eof(channel);
|
|
done = true;
|
|
}
|
|
}
|
|
if (FD_ISSET(ssh_get_fd(session), &fds_read)) {
|
|
int n = ssh_channel_read(channel, buffer, sizeof(buffer), 0);
|
|
if (n > 0) {
|
|
write(STDOUT_FILENO, buffer, n);
|
|
} else if (n == 0) {
|
|
done = true;
|
|
}
|
|
}
|
|
if (ssh_channel_is_closed(channel) || ssh_channel_is_eof(channel)) {
|
|
done = true;
|
|
}
|
|
}
|
|
tcsetattr(STDIN_FILENO, TCSANOW, &orig_termios);
|
|
return 0;
|
|
}
|
|
|
|
int ssh_exec_command(ssh_session session, ssh_channel channel, const std::string& remote_cmd_str, bool silent, std::string* output, const std::map<std::string, std::string>& env) {
|
|
// Set environment variables using ssh_channel_request_env
|
|
for (const auto& kv : env) {
|
|
if (kv.first == "SSHPASS") continue;
|
|
ssh_channel_request_env(channel, kv.first.c_str(), kv.second.c_str());
|
|
}
|
|
int rc = ssh_channel_request_exec(channel, remote_cmd_str.c_str());
|
|
if (rc != SSH_OK) {
|
|
if (output) *output = std::string("Failed to exec remote command: ") + ssh_get_error(session);
|
|
return -1;
|
|
}
|
|
if (output) {
|
|
std::ostringstream oss;
|
|
char buffer[4096];
|
|
int nbytes;
|
|
while ((nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0)) > 0) {
|
|
oss.write(buffer, nbytes);
|
|
}
|
|
*output = oss.str();
|
|
} else if (!silent) {
|
|
char buffer[4096];
|
|
int nbytes;
|
|
while ((nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0)) > 0) {
|
|
write(1, buffer, nbytes);
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
int local_execute_cmd(
|
|
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
|
|
) {
|
|
int pipefd[2];
|
|
bool use_pipe = output && !interactive;
|
|
if (use_pipe && pipe(pipefd) == -1) {
|
|
perror("pipe");
|
|
return -1;
|
|
}
|
|
pid_t pid = fork();
|
|
if (pid == -1) {
|
|
perror("fork");
|
|
return -1;
|
|
}
|
|
if (pid == 0) {
|
|
if (!working_dir.empty()) {
|
|
if (chdir(working_dir.c_str()) != 0) {
|
|
perror("chdir");
|
|
exit(-1);
|
|
}
|
|
}
|
|
for (const auto& kv : env) {
|
|
setenv(kv.first.c_str(), kv.second.c_str(), 1);
|
|
}
|
|
if (use_pipe) {
|
|
close(pipefd[0]);
|
|
dup2(pipefd[1], STDOUT_FILENO);
|
|
dup2(pipefd[1], STDERR_FILENO);
|
|
close(pipefd[1]);
|
|
} else if (silent && !interactive) {
|
|
int devnull = open("/dev/null", O_WRONLY);
|
|
dup2(devnull, STDOUT_FILENO);
|
|
dup2(devnull, STDERR_FILENO);
|
|
close(devnull);
|
|
}
|
|
if (!interactive) {
|
|
setsid();
|
|
}
|
|
std::vector<char*> argv;
|
|
argv.push_back(const_cast<char*>(command.c_str()));
|
|
for (const auto& arg : args) {
|
|
argv.push_back(const_cast<char*>(arg.c_str()));
|
|
}
|
|
argv.push_back(nullptr);
|
|
execvp(command.c_str(), argv.data());
|
|
perror("execvp");
|
|
exit(-1);
|
|
} else {
|
|
if (use_pipe) {
|
|
close(pipefd[1]);
|
|
std::ostringstream oss;
|
|
char buf[4096];
|
|
ssize_t n;
|
|
while ((n = read(pipefd[0], buf, sizeof(buf))) > 0) {
|
|
oss.write(buf, n);
|
|
}
|
|
close(pipefd[0]);
|
|
*output = oss.str();
|
|
}
|
|
int status = 0;
|
|
waitpid(pid, &status, 0);
|
|
if (WIFEXITED(status)) {
|
|
return WEXITSTATUS(status);
|
|
} else {
|
|
return -1;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // anonymous namespace
|
|
|
|
int execute_cmd(
|
|
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,
|
|
sSSHInfo* sshinfo,
|
|
std::string* output
|
|
) {
|
|
if (sshinfo) {
|
|
std::string error;
|
|
ssh_session session = ssh_connect_and_auth(sshinfo, env, &error);
|
|
if (!session) {
|
|
if (output) *output = error;
|
|
return -1;
|
|
}
|
|
ssh_channel channel = ssh_channel_new(session);
|
|
if (!channel) {
|
|
if (output) *output = "Failed to create SSH channel.";
|
|
ssh_disconnect(session);
|
|
ssh_free(session);
|
|
return -1;
|
|
}
|
|
int rc = ssh_channel_open_session(channel);
|
|
if (rc != SSH_OK) {
|
|
if (output) *output = std::string("Failed to open SSH channel: ") + ssh_get_error(session);
|
|
ssh_channel_free(channel);
|
|
ssh_disconnect(session);
|
|
ssh_free(session);
|
|
return -1;
|
|
}
|
|
std::string remote_cmd_str = ssh_build_remote_command(command, args, working_dir, {}); // Don't prefix env
|
|
int ret = 0;
|
|
if (interactive) {
|
|
ret = ssh_interactive_shell_session(session, channel, remote_cmd_str, command, output);
|
|
} else {
|
|
ret = ssh_exec_command(session, channel, remote_cmd_str, silent, output, env);
|
|
}
|
|
ssh_channel_send_eof(channel);
|
|
ssh_channel_close(channel);
|
|
ssh_channel_free(channel);
|
|
ssh_disconnect(session);
|
|
ssh_free(session);
|
|
return ret;
|
|
} else {
|
|
return local_execute_cmd(command, args, working_dir, env, silent, interactive, output);
|
|
}
|
|
}
|
|
|
|
} // namespace runner
|