dropshell/source/src/commands/ssh.cpp

#include "command_registry.hpp"
//#include "config.hpp"
#include "ordered_env.hpp"
#include "utils/utils.hpp"
#include "utils/directories.hpp"
#include "utils/execute.hpp"
#include "shared_commands.hpp"
#include "servers.hpp"
#include "services.hpp"
#include "servers.hpp"
#include "templates.hpp"
#include <libassert/assert.hpp>

namespace dropshell
{

    int ssh_handler(const CommandContext &ctx);

    static std::vector<std::string> ssh_name_list = {"ssh"};

    // Static registration
    struct SSHCommandRegister
    {
        SSHCommandRegister()
        {
            CommandRegistry::instance().register_command({ssh_name_list,
                                                          ssh_handler,
                                                          shared_commands::std_autocomplete,
                                                          false, // hidden
                                                          true,  // requires_config
                                                          true,  // requires_install
                                                          1,     // min_args (after command)
                                                          2,     // max_args (after command)
                                                          "ssh SERVER",
                                                          "SSH into a server, or into a docker container for a service.",
                                                          R"(
                
                ssh SERVER SERVICE     SSH into a docker container for a service.
                ssh SERVER             SSH into a server.
                ssh USER@SERVER        SSH into a server as a specific user.
                )"});
        }
    } ssh_command_register;

    bool ssh_into_server(const std::string &server, std::string user, const ordered_env_vars & env_vars, std::string remote_path)
    {
        ServerConfig server_env(server);
        if (!server_env.is_valid())
        {
            error << "Server " << server << " is not valid" << std::endl;
            return false;
        }
        std::string bash_cmd = "ls --color && exec bash --rcfile <(cat ~/.bashrc 2>/dev/null; echo 'cd " + quote(remote_path) + "')";

        info << "SSHing into " << server << ":" << remote_path << " as user " << user << std::endl;

        execute_ssh_command(server_env.get_SSH_INFO(user), sCommand(remote_path, bash_cmd, env_vars), cMode::Interactive);
        return true;
    }

    bool ssh_into_service(const std::string &server, const std::string &service)
    {
        ServerConfig server_env(server);
        if (!server_env.is_valid())
        {
            error << "Server " << server << " is not valid" << std::endl;
            return false;
        }

        if (!legal_service_name(service))
        {
            error << "Service name contains illegal characters: " << service << std::endl;
            return false;
        }

        LocalServiceInfo sinfo = get_service_info(server, service);
        if (!SIvalid(sinfo))
        {
            error << "Service " << service << " is not valid" << std::endl;
            return false;
        }

        if (!gTemplateManager().has_template(sinfo.template_name))
        {
            error << "Template " << sinfo.template_name << " is not valid" << std::endl;
            return false;
        }

        if (gTemplateManager().template_command_exists(sinfo.template_name, "ssh"))
        {
            return server_env.run_remote_template_command(service, "ssh", {}, false, {}); // explicitly supports interactive ssh!
        }

        // ssh in without as ssh.sh script.
        ordered_env_vars env_vars;
        if (!get_all_service_env_vars(server, service, env_vars))
        {
            error << "Failed to get all service env vars for " << service << std::endl;
            return false;
        }
        set_var(env_vars, "HOST_NAME", server_env.get_SSH_HOST());
        std::string service_dir = remotepath(server, sinfo.user).service_template(service);

        return ssh_into_server(server, sinfo.user, env_vars, service_dir);
    }

    int ssh_handler(const CommandContext &ctx)
    {
        if (ctx.args.size() < 1)
        {
            error << "Server name is required" << std::endl;
            return 1;
        }

        // ssh into the server
        if (ctx.args.size() < 2)
        {
            std::string arg1 = safearg(ctx.args, 0);
            std::string server, user;
            ordered_env_vars env_vars;

            // parse either user@server or server
            if (arg1.find("@") != std::string::npos)
            {
                user = arg1.substr(0, arg1.find("@"));
                server = arg1.substr(arg1.find("@") + 1);
            }
            else
            {
                server = arg1;
            }

            // get the first user from the server.env file, and ssh in as that user.
            ServerConfig server_env(server);
            if (!server_env.is_valid())
            {
                error << "Server " << server << " is not valid" << std::endl;
                return 1;
            }

            if (user.empty())
            {
                ASSERT(server_env.get_users().size() > 0, "Server " + server + " has no users");
                user = server_env.get_users()[0].user;
            }

            set_var(env_vars, "HOST_NAME", server_env.get_SSH_HOST());
            std::string dropshell_dir = remotepath(server, user).DROPSHELL_DIR();
            return ssh_into_server(server, user, env_vars, dropshell_dir) ? 0 : 1;
        }
        else
        { // ssh into a service on the server.
            std::string server = safearg(ctx.args, 0);
            std::string service = safearg(ctx.args, 1);
            return ssh_into_service(server, service) ? 0 : 1;
        }
    }

} // namespace dropshell