Ticket #131: ssh_api.h

/*
 * ssh_api.h
 *
 *  Created on: 2 juil. 2015
 *      Author: cresson
 */

#ifndef SSH_API_H_
#define SSH_API_H_

#include <libssh/libssh.h>
#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <iostream>
#include <sstream>

// For remote files upload
#include <fstream>      // std::ifstream
#include <libssh/sftp.h>
#include <sys/stat.h>
#include <fcntl.h>

// Remote constants TODO: get them from cfg 
const std::string REMOTE_PROCESS_GET_STATUS = "/gpfs-dell/data/work/GEOSUD/drmaa/test/drmaa_operator/bin/get_status ";
const std::string REMOTE_PROCESS_SUBMIT_COMMAND_ASYNC = "/gpfs-dell/data/work/GEOSUD/drmaa/test/drmaa_operator/bin/submit_command_async ";
const std::string REMOTE_PROCESS_SUBMIT_COMMAND_SYNC = "/gpfs-dell/data/work/GEOSUD/drmaa/test/drmaa_operator/bin/submit_command_sync ";
const std::string REMOTE_SCRATCH_DIRECTORY = "/gpfs-dell/scratch/cressonr/";
const std::string REMOTE_SERVER_NAME = "<your server name here (frontal)>";
const int REMOTE_SHELL_TIMEOUT_MS = 10000; // 10 sec

// Good chunk size
#define MAX_XFER_BUF_SIZE 16384

/*
 * Authenticating the server
 */
int verify_knownhost(ssh_session &session)
{
        int state, hlen;
        unsigned char *hash = NULL;
        char *hexa;
        char buf[10];
        state = ssh_is_server_known(session);
        ssh_key srv_pubkey;
        if (ssh_get_publickey(session, &srv_pubkey) < 0)
        {
                fprintf(stderr, "Public key not found:\n");
                return SSH_ERROR;
        }
        switch (state)
        {
        case SSH_SERVER_KNOWN_OK:
                break; /* ok */
        case SSH_SERVER_KNOWN_CHANGED:
                fprintf(stderr, "Host key for server changed: it is now:\n");
                ssh_print_hexa("Public key hash", hash, hlen);
                fprintf(stderr, "For security reasons, connection will be stopped\n");
                free(hash);
                return SSH_ERROR;
        case SSH_SERVER_FOUND_OTHER:
                fprintf(stderr, "The host key for this server was not found but an other"
                                "type of key exists.\n");
                fprintf(stderr, "An attacker might change the default server key to"
                                "confuse your client into thinking the key does not exist\n");
                free(hash);
                return SSH_ERROR;
        case SSH_SERVER_FILE_NOT_FOUND:
                fprintf(stderr, "Could not find known host file.\n");
                fprintf(stderr, "If you accept the host key here, the file will be"
                                "automatically created.\n");
                return SSH_ERROR;
        case SSH_SERVER_NOT_KNOWN:
                hexa = ssh_get_hexa(hash, hlen);
                fprintf(stderr,"The server is unknown. Do you trust the host key?\n");
                fprintf(stderr, "Public key hash: %s\n", hexa);
                free(hexa);
                if (fgets(buf, sizeof(buf), stdin) == NULL)
                {
                        free(hash);
                        return SSH_ERROR;
                }
                if (strncasecmp(buf, "yes", 3) != 0)
                {
                        free(hash);
                        return SSH_ERROR;
                }
                if (ssh_write_knownhost(session) < 0)
                {
                        fprintf(stderr, "Error %s\n", strerror(errno));
                        free(hash);
                        return SSH_ERROR;
                }
                break;
        case SSH_SERVER_ERROR:
                fprintf(stderr, "Error %s\n", ssh_get_error(session));
                free(hash);
                return SSH_ERROR;
        }
        free(hash);
        return SSH_OK;
}

/*
 * Authenticate ourselves
 */
int authenticate_pubkey(ssh_session &session)
{
        int rc;
        rc = ssh_userauth_publickey_auto(session, "cressonr", NULL);
        if (rc == SSH_AUTH_ERROR)
        {
                fprintf(stderr, "Authentication failed: %s\n",
                                ssh_get_error(session));
                return SSH_ERROR;
        }
        return SSH_OK;
}

/*
 * Show remote process output
 */
int show_remote_processes(ssh_session &session, const std::string &command, std::string &output)
{
        output.clear();
        ssh_channel channel;
        int rc;
        char buffer[256];
        unsigned int nbytes;
        channel = ssh_channel_new(session);
        fprintf(stderr, "Opening ssh session\n" );
        if (channel == NULL)
        {
                fprintf(stderr, "ssh_channel_new(session) failed %s\n", ssh_get_error(session));
                output = "ssh_channel_new(session) failed";
                return SSH_ERROR;
        }

        rc = ssh_channel_open_session(channel);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error opening session\n");
                output = "Error opening session";
                ssh_channel_free(channel);
                return SSH_ERROR;
        }

        rc = ssh_channel_request_pty(channel);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error requesting pty\n");
                output = "Error requesting pty\n";
                ssh_channel_close(channel);
                ssh_channel_free(channel);

                return rc;
        }
        rc = ssh_channel_change_pty_size(channel, 80, 24);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error changing pty size\n");
                output = "Error changing pty size\n";
                ssh_channel_close(channel);
                ssh_channel_free(channel);

                return rc;
        }
        rc = ssh_channel_request_shell(channel);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error requesting shell\n");
                output = "Error requesting shell\n";
                ssh_channel_close(channel);
                ssh_channel_free(channel);

                return rc;
        }

        // Send the command to shell
        std::string tosend = command + " ; RETVAL=$? ; [ $RETVAL -eq 0 ] && echo Success ; [ $RETVAL -ne 0 ] && echo Failure ; exit\n";
        char * tosendBuffer = (char*)tosend.c_str();
        nbytes = tosend.size(); //read(0, tosendBuffer, sizeof(tosendBuffer));
        fprintf(stderr, "Sending number of bytes:%d\n", nbytes);
        int nwritten;
        if (nbytes < 0) return SSH_ERROR;
        if (nbytes > 0)
        {
                nwritten = ssh_channel_write(channel, tosendBuffer, nbytes);
                if (nwritten != nbytes) return SSH_ERROR;
        }

        // Receive the shell output
        while (ssh_channel_is_open(channel) &&
                        !ssh_channel_is_eof(channel))
        {
                nbytes = ssh_channel_read_timeout(channel, buffer, sizeof(buffer), 0, REMOTE_SHELL_TIMEOUT_MS);
                if (nbytes < 0)
                {
                        if (nbytes==SSH_AGAIN)
                        {
                                fprintf(stderr, "SSH Channel read error: Timeout");
                                output += "SSH Channel read error: Timeout";
                        }
                        return SSH_ERROR;
                }
                else if (nbytes>0)
                {
                        write(2, buffer, nbytes);

                        char subbuff[nbytes+1];
                        memcpy( subbuff, &buffer[0], nbytes);
                        subbuff[nbytes] = '\0';
                        std::string tmpStr(subbuff);
                        output += tmpStr;

                }
        }

        if (nbytes < 0)
        {
                fprintf(stderr, "Error receiving command result\n");
                output = "Error receiving command result";
                ssh_channel_close(channel);
                ssh_channel_free(channel);
                return SSH_ERROR;
        }
        ssh_channel_send_eof(channel);
        ssh_channel_close(channel);
        ssh_channel_free(channel);
        return SSH_OK;
}

/*
 * Open an SSH session
 */
int open_ssh_session(const std::string &hostname, ssh_session &my_ssh_session)
{
        int rc;
        char *password;

        // Open session and set options
        my_ssh_session = ssh_new();
        if (my_ssh_session == NULL)
        {
                fprintf(stderr, "ssh_new() failed: %s\n",
                                ssh_get_error(my_ssh_session));
                return SSH_ERROR;
        }
        ssh_options_set(my_ssh_session, SSH_OPTIONS_HOST, hostname.c_str());

        // Connect to server
        rc = ssh_connect(my_ssh_session);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error connecting to host: %s\n",
                                ssh_get_error(my_ssh_session));
                ssh_free(my_ssh_session);
                return rc;
        }

        // Verify the server's identity
        // For the source code of verify_knowhost(), check previous example
        if (verify_knownhost(my_ssh_session) < 0)
        {
                fprintf(stderr, "Error verifying the server's identity: %s\n",
                                ssh_get_error(my_ssh_session));
                ssh_disconnect(my_ssh_session);
                ssh_free(my_ssh_session);
                return SSH_ERROR;
        }

        // Authenticate ourselves
        rc = authenticate_pubkey(my_ssh_session);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error authenticating ourselves: %s\n",
                                ssh_get_error(my_ssh_session));
                ssh_free(my_ssh_session);
        }
        return rc;
}


/*
 * Output the result of the given input command
 */
int get_command_output(const std::string &command, std::string &output)
{
        // Create a ssh session
        ssh_session my_ssh_session;
        fprintf(stderr,"Creating ssh session\n");
        int rc = open_ssh_session(REMOTE_SERVER_NAME, my_ssh_session);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error creating ssh session %s \n",
                                ssh_get_error(my_ssh_session));
                output = "Error opening ssh session";
                return SSH_ERROR;
        }

        // Execute remote process
        fprintf(stderr,"Executing remote process\n");
        rc = show_remote_processes(my_ssh_session, command, output);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error showing remote process %s : %s\n",
                                command.c_str(),
                                ssh_get_error(my_ssh_session));
                output = "Error showing remote process";
        }

        // Disconnect session
        ssh_disconnect(my_ssh_session);
        ssh_free(my_ssh_session);

        return rc;
}

/*
 * Get the status of a job
 */
int get_status(int job_id, std::string &output)
{
        std::string command (REMOTE_PROCESS_GET_STATUS);
        command += static_cast<std::ostringstream*>( &(std::ostringstream() << job_id) )->str();

        return get_command_output(command, output);
}

/*
 * Submit a command, asynchronous way (and returns its ID)
 */
int submit_command_async(const std::string &input_command, std::string &output)
{
        std::string command(REMOTE_PROCESS_SUBMIT_COMMAND_ASYNC);
        command += input_command;

        return get_command_output(command, output);
}

/*
 * Submit a job, synchronous way (and returns its status)
 */
int submit_command_sync(const std::string &input_command, std::string &output)
{

        std::string command(REMOTE_PROCESS_SUBMIT_COMMAND_SYNC);
        command += input_command;

        return get_command_output(command, output);
}

/*
 * sftp: upload a file
 */
int sftp_uploadfile(ssh_session &session, sftp_session &sftp)
{
        int access_type = O_WRONLY | O_CREAT | O_TRUNC;
        sftp_file file;
        const char *helloworld = "Hello, World!\n";
        int length = strlen(helloworld);
        int rc, nwritten;
        //...
        file = sftp_open(sftp, "helloworld.txt",
                        access_type, S_IRWXU);
        if (file == NULL)
        {
                fprintf(stderr, "Can't open file for writing: %s\n",
                                ssh_get_error(session));
                return SSH_ERROR;
        }
        nwritten = sftp_write(file, helloworld, length);
        if (nwritten != length)
        {
                fprintf(stderr, "Can't write data to file: %s\n",
                                ssh_get_error(session));
                sftp_close(file);
                return SSH_ERROR;
        }
        rc = sftp_close(file);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Can't close the written file: %s\n",
                                ssh_get_error(session));
                return rc;
        }
        return SSH_OK;
}

void display_sftp_error_message(sftp_session &sftp)
{
        int error = sftp_get_error(sftp);

        switch (error) {
        case SSH_FX_NO_SUCH_FILE:
                fprintf(stderr, "No such file\n");
                break;
        case SSH_FX_PERMISSION_DENIED:
                fprintf(stderr, "Permission denied\n");
                break;
        case SSH_FX_EOF:
                fprintf(stderr, "EOF encountered\n");
                break;
        case SSH_FX_FAILURE:
                fprintf(stderr, "Failure\n");
                break;
        case SSH_FX_BAD_MESSAGE:
                fprintf(stderr, "Strange error message received\n");
                break;
        case SSH_FX_NO_CONNECTION:
                fprintf(stderr, "No connection\n");
                break;
        case SSH_FX_CONNECTION_LOST:
                fprintf(stderr, "Conection lost\n");
                break;
        case SSH_FX_OP_UNSUPPORTED:
                fprintf(stderr, "Operation not supported\n");
                break;
        case SSH_FX_INVALID_HANDLE:
                fprintf(stderr, "Invalid handle\n");
                break;
        case SSH_FX_NO_SUCH_PATH:
                fprintf(stderr, "No such path\n");
                break;
        case SSH_FX_FILE_ALREADY_EXISTS:
                fprintf(stderr, "Already exists\n");
                break;
        case SSH_FX_WRITE_PROTECT:
                fprintf(stderr, "Write protected\n");
                break;
        case SSH_FX_NO_MEDIA:
                fprintf(stderr, "No media\n");
                break;
        }
}

/*
 * sftp: mkdir
 */
int sftp_mkdir(ssh_session &session, sftp_session &sftp,
                const std::string &directory)
{
        int rc;
        rc = sftp_mkdir(sftp, directory.c_str(), S_IRWXU);
        if (rc != SSH_OK)
        {
                display_sftp_error_message(sftp);
        }
        return rc;
}

/*
 * init. a sftp session
 */
int initialize_sftp_session(ssh_session &session, sftp_session &sftp)
{
        // Connect to server
        fprintf(stderr,"Open ssh session\n");
        int rc = open_ssh_session(REMOTE_SERVER_NAME, session);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "open_ssh_session failed: %s\n",
                                ssh_get_error(session));
                return SSH_ERROR;
        }

        // init. sftp session
        fprintf(stderr,"Init. sftp session\n");
        sftp = sftp_new(session);
        if (sftp == NULL)
        {
                fprintf(stderr, "sftp_new(session) failed: %s\n",
                                ssh_get_error(session));
                return SSH_ERROR;
        }
        rc = sftp_init(sftp);
        if (rc != SSH_OK)
        {
                fprintf(stderr, "Error initializing SFTP session: %s\n",
                                sftp_get_error(sftp));
        }
        return rc;
}

/*
 * create a directory on the host
 */
int create_remote_dir(const std::string &directory_name)
{
        std::string dir = directory_name;
        fprintf(stderr,"Creating directory: %s\n", dir.c_str() );

        int rc;

        // init. session
        ssh_session session;
        sftp_session sftp;
        rc = initialize_sftp_session(session, sftp);
        if (rc != SSH_OK)
        {
                return rc;
        }

        // create a directory
        rc = sftp_mkdir(session, sftp, dir.c_str());
        if (rc == SSH_OK)
        {
                fprintf(stderr,"Directory successfuly created\n");
        }

        // close sftp connection
        sftp_free(sftp);

        // Disconnect session
        ssh_disconnect(session);
        ssh_free(session);

        return SSH_OK;
}

/*
 * download a file
 */
int sftp_download_file(ssh_session &session, sftp_session &sftp,
                const std::string &local_filename, const std::string &remote_filename)
{
        std::string true_remote_filename = remote_filename;
        fprintf(stderr,"Creating local file: %s\n", local_filename.c_str() );

        int access_type;
        sftp_file file;
        char buffer[MAX_XFER_BUF_SIZE];
        int nbytes, nwritten, rc;
        int fd;
        access_type = O_RDONLY;
        file = sftp_open(sftp, true_remote_filename.c_str(),
                        access_type, 0);
        if (file == NULL) {
                fprintf(stderr, "Can't open file for reading: %s\n",
                                ssh_get_error(session));
                return SSH_ERROR;
        }
        mode_t mode = S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH;
        //      fd = open(local_filename.c_str(), O_RDWR|O_CREAT);
        fd = open(local_filename.c_str(), O_WRONLY | O_CREAT | O_TRUNC, mode);
        if (fd < 0) {
                fprintf(stderr, "Can't open file for writing: %s\n",
                                strerror(errno));
                return SSH_ERROR;
        }
        for (;;) {
                nbytes = sftp_read(file, buffer, sizeof(buffer));
                if (nbytes == 0) {
                        break; // EOF
                } else if (nbytes < 0) {
                        fprintf(stderr, "Error while reading file: %s\n",
                                        ssh_get_error(session));
                        sftp_close(file);
                        close(fd);
                        return SSH_ERROR;
                }
                nwritten = write(fd, buffer, nbytes);
                if (nwritten != nbytes) {
                        fprintf(stderr, "Error writing: %s\n",
                                        strerror(errno));
                        sftp_close(file);
                        close(fd);
                        return SSH_ERROR;
                }
        }
        rc = sftp_close(file);
        if (rc != SSH_OK) {
                fprintf(stderr, "Can't close the read file: %s\n",
                                ssh_get_error(session));
        }
        close(fd);
        return rc;
}

/*
 * upload a file
 */
int sftp_upload_file(ssh_session &session, sftp_session &sftp,
                const std::string &local_filename, const std::string &remote_filename)
{
        std::string true_remote_filename = remote_filename;
        fprintf(stderr, "Creating remote file: %s\n", true_remote_filename.c_str() );

        int access_type;
        sftp_file file;
        char buffer[MAX_XFER_BUF_SIZE];
        int nbytes, nwritten, rc;
        int fd;
        //      access_type = O_WRONLY;
        access_type = O_WRONLY | O_CREAT | O_TRUNC;
        file = sftp_open(sftp, true_remote_filename.c_str(),
                        access_type, S_IRWXU /* 0 */);
        if (file == NULL) {
                fprintf(stderr, "Can't open file for writing: %s\n",
                                ssh_get_error(session));
                return SSH_ERROR;
        }

        fd = open(local_filename.c_str(), O_RDONLY);
        if (fd < 0) {
                fprintf(stderr, "Can't open file for reading: %s\n",
                                strerror(errno));
                return SSH_ERROR;
        }

        for (;;) {
                nbytes = read(fd, buffer, sizeof(buffer));
                if (nbytes == 0) {
                        break; // EOF
                } else if (nbytes < 0) {
                        fprintf(stderr, "Error while reading file: %s\n",
                                        ssh_get_error(session));
                        sftp_close(file);
                        return SSH_ERROR;
                }
                nwritten = sftp_write(file, buffer, nbytes);
                if (nwritten != nbytes) {
                        fprintf(stderr, "Error writing: %s\n",
                                        strerror(errno));
                        sftp_close(file);
                        return SSH_ERROR;
                }
        }
        rc = sftp_close(file);
        if (rc != SSH_OK) {
                fprintf(stderr, "Can't close the write file: %s\n",
                                ssh_get_error(session));
                rc = SSH_ERROR;
        }
        return rc;
}

/*
 * Upload a file to the remote host
 */
int upload_file(const std::string &local_filename, const std::string &remote_filename)
{

        fprintf(stderr, "Upload file %s (local) to %s (remote)\n", local_filename.c_str(), remote_filename.c_str() );

        int rc;

        // init. session
        ssh_session session;
        sftp_session sftp;
        rc = initialize_sftp_session(session, sftp);
        if (rc != SSH_OK)
        {
                return rc;
        }

        // transfert file
        rc = sftp_upload_file(session, sftp, local_filename, remote_filename);
        if (rc == SSH_OK)
        {
                fprintf(stderr,"Remote file successfuly uploaded\n");
        }

        // close sftp connection
        sftp_free(sftp);

        // Disconnect session
        ssh_disconnect(session);
        ssh_free(session);

        return rc;
}

/*
 * Download a file to the remote host
 */
int download_file(const std::string &local_filename, const std::string &remote_filename)
{

        fprintf(stderr,"Download file %s (remote) to %s (local)\n", remote_filename.c_str(), local_filename.c_str() );

        int rc;

        // init. session
        ssh_session session;
        sftp_session sftp;
        rc = initialize_sftp_session(session, sftp);
        if (rc != SSH_OK)
        {
                return rc;
        }

        // transfert file
        rc = sftp_download_file(session, sftp, local_filename, remote_filename);
        if (rc == SSH_OK)
        {
                fprintf(stderr,"Remote file successfuly downloaded\n");
        }

        // close sftp connection
        sftp_free(sftp);

        // Disconnect session
        ssh_disconnect(session);
        ssh_free(session);

        return rc;
}


#endif /* SSH_API_H_ */