dropshell/runner/runner.cpp
Your Name 34f2763ef4
Some checks failed
Dropshell Test / Build_and_Test (push) Failing after 20s
Broken
2025-05-10 14:27:24 +12:00

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