#include <sys/socket.h>
#include <arpa/inet.h>

#include <fcntl.h>
#include <wolfssh/ssh.h>
#include <errno.h>
#include <sys/epoll.h>

#define SSH_SERVER_ADDR "192.168.188.69"
#define SSH_USER "myUsername"
#define SSH_PASS "myPassword"

int sock = 0;
WOLFSSH_CTX* ctx = NULL;
WOLFSSH* ssh = NULL;

struct epoll_event events[5];

enum states
{
    CONNECTING,
    LOGIN,
    DONE
};

int currState = CONNECTING;

int setup_socket();

int connect_sock();

int get_socket_error();

int init_ssh();

int authCb(byte authType, WS_UserAuthData* authData, void* userData);

int login();

int main()
{
    // create nonblocking socket
    if(setup_socket() < 0)
    {
        perror("Setting up socket failed");
        if(sock > 0)
            close(sock);
        exit(0);
    }

    // call non-blocking connect
    if(connect_sock() < 0)
    {
        perror("Connecting failed");
        close(sock);
        exit(0);
    }

    // initialize wolfssh
    if(init_ssh() < 0)
    {
        printf("ssh error: %d", wolfSSH_get_error(ssh));
        if(ctx)
            wolfSSH_CTX_free(ctx);
        if(ssh)
            wolfSSH_free(ssh);
        close(sock);
        exit(0);
    }

    // setup epoll
    int epoll = epoll_create(1);

    struct epoll_event event;
    event.events = EPOLLIN | EPOLLOUT;
    event.data.fd = sock;
    if(epoll_ctl(epoll, EPOLL_CTL_ADD, sock, &event) < 0)
    {
        perror("epoll_ctl");
        wolfSSH_CTX_free(ctx);
        wolfSSH_free(ssh);
        close(sock);
        exit(0);
    }

    int nfds;
    int notDone = 1;
    while(notDone)
    {
        nfds = epoll_wait(epoll, events, sizeof(events), 3);
        if(nfds == -1)
        {
            perror("epoll_wait");
            break;
        }

        int ret;
        for(int n = 0; n < nfds; ++n)
        {
            if(events[n].data.fd == sock && (events[n].events & (EPOLLIN | EPOLLOUT)))
            {
                switch(currState)
                {
                    case CONNECTING:
                        ret = get_socket_error();
                        if(ret != 0 && ret != EINPROGRESS)
                        {
                            printf("Failed to connect, %s\n", strerror(ret));
                            currState = DONE;
                        }
                        else if(ret != EINPROGRESS)
                        {
                            printf("Connected\n");
                            currState = LOGIN;
                        }
                        else
                        {
                            currState = CONNECTING;
                        }
                        break;
                    case LOGIN:
                        ret = login();
                        if(ret == WS_SUCCESS)
                        {
                            printf("Login success");
                            currState = DONE;
                        }
                        else if(ret == WS_WANT_READ || ret == WS_WANT_WRITE)
                        {
                            // repeat calling login()
                            currState = LOGIN;
                        }
                        else
                        {
                            printf("Login failed, code %d", wolfSSH_get_error(ssh));
                            currState = DONE;
                        }
                        break;
                    case DONE:
                        notDone = 0;
                        break;
                }
            }
        }
    }

    wolfSSH_shutdown(ssh);
    wolfSSH_CTX_free(ctx);
    wolfSSH_free(ssh);
    close(sock);

    return 0;
}

int setup_socket()
{
    sock = socket(AF_INET, SOCK_STREAM, 0);
    if(sock < 0)
        return -1;

    // set non-blocking
    if(fcntl(sock, F_SETFL, O_NONBLOCK | fcntl(sock, F_GETFL, 0)) < 0)
        return -1;

    return 0;
}

int connect_sock()
{
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(22);
    inet_pton(AF_INET, SSH_SERVER_ADDR, &(addr.sin_addr));

    if(connect(sock, (struct sockaddr*)(&addr), sizeof(addr)) < 0 && errno != EINPROGRESS)
        return -1;

    return 0;
}

int get_socket_error()
{
    int err;
    socklen_t size = sizeof(err);
    int ret = getsockopt(sock, SOL_SOCKET, SO_ERROR, &err, &size);
    if(ret < 0)
        return ret;

    return err;
}

int init_ssh()
{
    wolfSSH_Init();

    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
    if(!ctx)
        return -1;

    ssh = wolfSSH_new(ctx);
    if(!ssh)
        return -1;

    int ret = wolfSSH_set_fd(ssh, sock);
    if(ret != WS_SUCCESS)
        return -1;

    // set channel to shell
    ret = wolfSSH_SetChannelType(ssh, WOLFSSH_SESSION_SHELL, NULL, 0);
    if(ret != WS_SUCCESS)
        return -1;

    // set auth callback
    wolfSSH_SetUserAuth(ctx, authCb);
    // set callback data
    wolfSSH_SetUserAuthCtx(ssh, NULL);
    // set username
    wolfSSH_SetUsername(ssh, SSH_USER);

    return 0;
}

int authCb(byte authType, WS_UserAuthData* authData, void* userData)
{

    if(authType == WOLFSSH_USERAUTH_PASSWORD)
    {
        authData->sf.password.password = (byte*)(SSH_PASS);
        authData->sf.password.passwordSz = strlen(SSH_PASS);
    }
    else
    {
        return WOLFSSH_USERAUTH_REJECTED;
    }

    return WOLFSSH_USERAUTH_SUCCESS;
}

int login()
{
    int ret = wolfSSH_connect(ssh);

    // Here comes th odd part. If you remove the call to usleep(),
    // WS_OVERFLOW_E is returned by wolfssh.
    // With usleep it works perfectly.
    usleep(10000);

    if(ret != WS_SUCCESS)
    {
        return wolfSSH_get_error(ssh);
    }

    return WS_SUCCESS;
}