/* Self test, based on examples/ssh2.c. */

#include "libssh2_config.h"
#include <libssh2.h>
#include <libssh2_sftp.h>

#ifdef HAVE_WINDOWS_H
# include <windows.h>
#endif
#ifdef HAVE_WINSOCK2_H
# include <winsock2.h>
#endif
#ifdef HAVE_SYS_SOCKET_H
# include <sys/socket.h>
#endif
#ifdef HAVE_NETINET_IN_H
# include <netinet/in.h>
#endif
# ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif
# ifdef HAVE_ARPA_INET_H
#include <arpa/inet.h>
#endif

#include <sys/types.h>
#include <fcntl.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>

static int total_read = 0, total_written = 0;
static LIBSSH2_SESSION *session;

static void write_retry(LIBSSH2_CHANNEL *channel, const char *buffer, int len) {
    int current_len = 0;
    while (current_len < len) {
        fprintf(stderr, "writing %d bytes of %d\n", len - current_len, len); fflush(stderr);
        int count = libssh2_channel_write_ex(channel, 0, buffer + current_len, len - current_len);
        fprintf(stderr, "written count: %d\n", count); fflush(stderr);
        if (count <= 0) {
            char *err;
            libssh2_session_last_error(session, &err, 0, 0);
            fprintf(stderr, "write failed: %s", err);
            exit(1);
        }
        current_len += count;
        total_written += count;
        fprintf(stderr, "written current: %d\n\t\twritten total: %d read total: %d\n",
                current_len, total_written, total_read); fflush(stderr);
    }
}

static void read_retry(LIBSSH2_CHANNEL *channel, char *buffer, int len) {
    int current_len = 0;
    while (current_len < len) {
        fprintf(stderr, "reading %d bytes of %d\n", len - current_len, len); fflush(stderr);
        int count = libssh2_channel_read_ex(channel, 0, buffer + current_len, len - current_len);
        fprintf(stderr, "read count: %d\n", count); fflush(stderr);
        if (count <= 0) {
            char *err;
            libssh2_session_last_error(session, &err, 0, 0);
            fprintf(stderr, "read failed: %s", err);
            exit(1);
        }
        current_len += count;
        total_read += count;
        fprintf(stderr, "read current: %d\n\t\twritten total: %d read total: %d\n",
                current_len, total_written, total_read); fflush(stderr);
        if (total_read > LIBSSH2_CHANNEL_WINDOW_DEFAULT) {
            fprintf(stderr, "done\n");
            exit(0);
        }
    }
}

static void discard_packet(LIBSSH2_CHANNEL *channel) {
    char buffer[1024];
    read_retry(channel, buffer, 4);
    unsigned int size = (buffer[0] << 24 ) | (buffer[1] << 16) | (buffer[2] << 8) | buffer[3];
    if (size > 1024) {
        fprintf(stderr, "packet too big: %d\n", size);
        exit(1);
    }
    read_retry(channel, buffer, size);
}

int main(int argc, char *argv[])
{
    unsigned long hostaddr;
    int sock, i, auth_pw = 0;
    struct sockaddr_in sin;
    const char *fingerprint;
    char *userauthlist;
    LIBSSH2_CHANNEL *channel;
#ifdef WIN32
    WSADATA wsadata;

    WSAStartup(MAKEWORD(2,0), &wsadata);
#endif
    const char *pubkeyfile="etc/user.pub";
    const char *privkeyfile="etc/user";
    const char *username="password";
    const char *password="password";
    int ec = 1;

    (void)argc;
    (void)argv;

    if (getenv ("USER"))
      username = getenv ("USER");

    if (getenv ("PRIVKEY"))
      privkeyfile = getenv ("PRIVKEY");

    if (getenv ("PUBKEY"))
      pubkeyfile = getenv ("PUBKEY");

    hostaddr = htonl(0x7F000001);

    sock = socket(AF_INET, SOCK_STREAM, 0);
#ifndef WIN32
    fcntl(sock, F_SETFL, 0);
#endif
    sin.sin_family = AF_INET;
    sin.sin_port = htons(22);
    sin.sin_addr.s_addr = hostaddr;
    if (connect(sock, (struct sockaddr*)(&sin),
                sizeof(struct sockaddr_in)) != 0) {
        fprintf(stderr, "failed to connect!\n");
        return -1;
    }

    /* Create a session instance and start it up
     * This will trade welcome banners, exchange keys, and setup crypto, compression, and MAC layers
     */
    session = libssh2_session_init();
    if (libssh2_session_startup(session, sock)) {
        fprintf(stderr, "Failure establishing SSH session\n");
        return -1;
    }

    /* At this point we havn't authenticated,
     * The first thing to do is check the hostkey's fingerprint against our known hosts
     * Your app may have it hard coded, may go to a file, may present it to the user, that's your call
     */
    fingerprint = libssh2_hostkey_hash(session, LIBSSH2_HOSTKEY_HASH_SHA1);
    printf("Fingerprint: ");
    for(i = 0; i < 20; i++) {
        printf("%02X ", (unsigned char)fingerprint[i]);
    }
    printf("\n");

    /* check what authentication methods are available */
    userauthlist = libssh2_userauth_list(session, username, strlen(username));
    printf("Authentication methods: %s\n", userauthlist);
    if (strstr(userauthlist, "password") != NULL) {
        auth_pw |= 1;
    }
    if (strstr(userauthlist, "keyboard-interactive") != NULL) {
        auth_pw |= 2;
    }
    if (strstr(userauthlist, "publickey") != NULL) {
        auth_pw |= 4;
    }

    if (auth_pw & 4) {
        /* Authenticate by public key */
        if (libssh2_userauth_publickey_fromfile(session, username, pubkeyfile, privkeyfile, password)) {
            printf("\tAuthentication by public key failed!\n");
            goto shutdown;
        } else {
            printf("\tAuthentication by public key succeeded.\n");
        }
    } else {
        printf("No supported authentication methods found!\n");
        goto shutdown;
    }

    /* Request a shell */
    if (!(channel = libssh2_channel_open_ex(session, "session", 7,
                                            LIBSSH2_CHANNEL_WINDOW_DEFAULT,
                                            LIBSSH2_CHANNEL_PACKET_DEFAULT,
                                            NULL, 0))) {
        fprintf(stderr, "Unable to open a session\n");
        goto shutdown;
    }

    libssh2_channel_handle_extended_data(channel, LIBSSH2_CHANNEL_EXTENDED_DATA_IGNORE);

    if (libssh2_channel_process_startup(channel, "subsystem", 9, "sftp", 4) < 0) {
        fprintf(stderr, "Unable to start subsystem\n");
        goto shutdown;
    }

    while (1) {
        switch(getchar()) {
        case 'w':
        {
            unsigned char buffer[65000];
            if (fread(buffer, 1, 4, stdin) != 4) {
                fprintf(stderr, "unable to read dump");
                exit(1);
            }
            int size = (buffer[0] << 24) | (buffer[1] << 16) | (buffer[2] << 8) | buffer[3];
            if (size > sizeof(buffer)) {
                fprintf(stderr, "saved packet is bigger than buffer: %d", size);
                exit(1);
            }
            if (fread(buffer, 1, size, stdin) != size) {
                fprintf(stderr, "unable to read dump");
                exit(1);
            }
            write_retry(channel, buffer, size);
            break;
        }
        case 'r':
            discard_packet(channel);
            break;
        case EOF:
            exit (0);
        default:
            exit (1);
        }
    }

  shutdown:

    libssh2_session_disconnect(session, "Normal Shutdown");
    libssh2_session_free(session);

#ifdef WIN32
    Sleep(1000);
    closesocket(sock);
#else
    sleep(1);
    close(sock);
#endif

    return ec;
}