/*
 * NNChat - Custom chat client for NewbieNudes.com chatrooms
 * Written by Matti 'ccr' Hämäläinen
 * (C) Copyright 2008-2013 Tecnic Software productions (TNSP)
 */
#include "network.h"
#include <errno.h>

static BOOL nn_network_inited = FALSE;

static const char *nn_proxy_types[] =
{
    "none",
    "SOCKS 4",
    "SOCKS 4a",
    NULL
};


#ifdef __WIN32
const char *hstrerror(int err)
{
    static char buf[64];
    snprintf(buf, sizeof(buf), "Error #%d", err);
    return buf;
}


int nn_get_socket_errno(void)
{
    return WSAGetLastError();
}


const char *nn_get_socket_errstr(int err)
{
    static char buf[64];
    switch (err)
    {
        case WSAEADDRINUSE:      return "Address already in use";
        case WSAECONNABORTED:    return "Software caused connection abort";
        case WSAECONNREFUSED:    return "Connection refused";
        case WSAECONNRESET:      return "Connection reset by peer";
        case WSAEHOSTUNREACH:    return "No route to host";
        case WSAENETDOWN:        return "Network is down";
        case WSAETIMEDOUT:       return "Connection timed out";
        case WSAHOST_NOT_FOUND:  return "Host not found";
        case WSAVERNOTSUPPORTED: return "Wrong WinSock DLL version";
        default:
            snprintf(buf, sizeof(buf), "Error #%d", err);
            return buf;
            break;
    }
}
#else
int nn_get_socket_errno(void)
{
    return errno;
}


const char *nn_get_socket_errstr(int err)
{
    return strerror(err);
}
#endif


void nn_conn_err(nn_conn_t *conn, const char *fmt, ...)
{
    if (conn->errfunc != NULL)
    {
        char *msg;
        va_list ap;
        va_start(ap, fmt);
        msg = th_strdup_vprintf(fmt, ap);
        va_end(ap);

        conn->errfunc(conn, msg);
        th_free(msg);
    }
}


static void nn_conn_msg(nn_conn_t *conn, const char *fmt, ...)
{
    if (conn->msgfunc != NULL)
    {
        char *msg;
        va_list ap;
        va_start(ap, fmt);
        msg = th_strdup_vprintf(fmt, ap);
        va_end(ap);

        conn->msgfunc(conn, msg);
        th_free(msg);
    }
}


struct hostent *nn_resolve_host(nn_conn_t *conn, const char *name)
{
    struct hostent *res = gethostbyname(name);

    if (res == NULL)
        nn_conn_err(conn, "Could not resolve hostname '%s': %s\n", name, strerror(h_errno));
    else
        nn_conn_msg(conn, "True hostname for %s is %s\n", name, res->h_name);

    return res;
}


nn_conn_t * nn_conn_new(
    void (*errfunc)(nn_conn_t *conn, const char *msg),
    void (*msgfunc)(nn_conn_t *conn, const char *msg))
{
    nn_conn_t *conn = th_calloc(1, sizeof(nn_conn_t));

    if (conn == NULL)
        return NULL;

    conn->errfunc = errfunc;
    conn->msgfunc = msgfunc;

    return conn;
}


static BOOL nn_get_addr(struct in_addr *addr, struct hostent *hst)
{
    if (hst != NULL)
    {
        *addr = *(struct in_addr *) (hst->h_addr);
        return TRUE;
    }
    else
    {
        addr->s_addr = 0;
        return FALSE;
    }
}


int nn_conn_set_proxy(nn_conn_t *conn, int type, int port, const char *host, const char *userid, const char *passwd)
{
    if (conn == NULL)
        return -1;

    conn->proxy.type = type;
    conn->proxy.port = port;
    conn->proxy.host = th_strdup(host);
    conn->proxy.userid = th_strdup(userid);
    conn->proxy.passwd = th_strdup(passwd);

    if (host != NULL)
    {
        conn->proxy.hst = nn_resolve_host(conn, host);
        nn_get_addr(&(conn->proxy.addr), conn->proxy.hst);
    }
    else
        return -2;

    return 0;
}


int nn_conn_open(nn_conn_t *conn, const int port, const char *host)
{
    struct sockaddr_in dest;

    if (conn == NULL)
        return -1;

    conn->port = port;
    if (host != NULL)
    {
        conn->host = th_strdup(host);
        conn->hst = nn_resolve_host(conn, host);
    }

    nn_get_addr(&(conn->addr), conn->hst);

    // Prepare for connection
    dest.sin_family = AF_INET;

    if (conn->proxy.type > NN_PROXY_NONE && conn->proxy.type < NN_PROXY_LAST)
    {
        dest.sin_port = htons(conn->proxy.port);
        dest.sin_addr = conn->proxy.addr;

        nn_conn_msg(conn, "Connecting to %s proxy %s:%d ...\n",
            nn_proxy_types[conn->proxy.type],
            inet_ntoa(conn->proxy.addr), conn->proxy.port);
    }
    else
    {
        dest.sin_port = htons(conn->port);
        dest.sin_addr = conn->addr;

        nn_conn_msg(conn, "Connecting to %s:%d ...\n",
            inet_ntoa(conn->addr), conn->port);
    }

    if ((conn->socket = socket(PF_INET, SOCK_STREAM, 0)) == -1)
    {
        conn->err = nn_get_socket_errno();
        nn_conn_err(conn, "Could not open socket: %s\n", nn_get_socket_errstr(conn->err));
        goto error;
    }

    if (connect(conn->socket, (struct sockaddr *) &dest, sizeof(dest)) == -1)
    {
        conn->err = nn_get_socket_errno();
        nn_conn_err(conn, "Could not connect: %s\n", nn_get_socket_errstr(conn->err));
        goto error;
    }

    FD_ZERO(&(conn->sockfds));
    FD_SET(conn->socket, &(conn->sockfds));

    // Proxy-specific setup
    if (conn->proxy.type == NN_PROXY_SOCKS4 || conn->proxy.type == NN_PROXY_SOCKS4A)
    {
        struct nn_socks_t *socksh;
        size_t bufsiz;
        char *ptr, *buf;
        int tries, status = -1;

        nn_conn_msg(conn, "Initializing proxy negotiation.\n");

        bufsiz = sizeof(struct nn_socks_t) + strlen(conn->proxy.userid) + 1;
        if (conn->proxy.type == NN_PROXY_SOCKS4A)
            bufsiz += strlen(conn->host) + 1;

        ptr = buf = th_malloc(bufsiz);
        if (buf == NULL)
        {
            conn->err = -1;
            nn_conn_err(conn, "Could not allocate memory for SOCKS negotiation buffer, %d bytes.\n", bufsiz);
            goto error;
        }

        // Create SOCKS 4/4A request
        socksh = (struct nn_socks_t *) buf;
        socksh->version = 4;
        socksh->command = SOCKS_CMD_CONNECT;
        socksh->port    = htons(port);
        socksh->addr    = (conn->proxy.type == NN_PROXY_SOCKS4A) ?  htonl(0x00000032) : conn->addr.s_addr;
        ptr += sizeof(struct nn_socks_t);

        strcpy(ptr, conn->proxy.userid);

        if (conn->proxy.type == NN_PROXY_SOCKS4A)
        {
            ptr += strlen(conn->proxy.userid) + 1;
            strcpy(ptr, conn->host);
        }

        // Send request
        nn_conn_reset(conn);
        if (!nn_conn_send_buf(conn, buf, bufsiz))
        {
            th_free(buf);
            nn_conn_err(conn, "Error sending SOCKS proxy request.\n");
            goto error;
        }
        th_free(buf);

        // Wait for SOCKS server to reply
        for (status = tries = 1; tries <= 20 && status > 0; tries++)
        {
#ifdef __WIN32
            Sleep(50);
#else
            usleep(50000);
#endif
            nn_conn_reset(conn);
            status = nn_conn_pull(conn);
        }

        // Check results
        if (status == 0)
        {
            struct nn_socks_res_t *res = (struct nn_socks_res_t *) &(conn->buf);
            if (res->nb != 0)
            {
                nn_conn_err(conn, "Invalid SOCKS server reply, does not begin with NUL byte (%d).\n", res->nb);
                goto error;
            }
            if (res->result != 0x5a)
            {
                char *s = NULL;
                switch (res->result)
                {
                    case 0x5b: s = "Request rejected or failed"; break;
                    case 0x5c: s = "Request failed because client is not running identd (or not reachable from the server)"; break;
                    case 0x5d: s = "Request failed because client's identd could not confirm the user ID string in the request"; break;
                    default: s = "Unknown SOCKS error response"; break;
                }
                nn_conn_err(conn, "SOCKS setup failed, 0x%02x: %s.\n", res->result, s);
                goto error;
            }
            nn_conn_msg(conn, "SOCKS connection established!\n");
        }
        else if (status < 0)
        {
            nn_conn_err(conn, "Proxy negotiation failed at try %d with network error: %d\n", tries, status);
            goto error;
        }
        else
        {
            nn_conn_err(conn, "Proxy negotiation timed out.\n");
            goto error;
        }
    }

    nn_conn_reset(conn);
    conn->status = NN_CONN_OPEN;
    return 0;

error:
    conn->status = NN_CONN_CLOSED;
    return -2;
}


void nn_conn_close(nn_conn_t *conn)
{
    if (conn == NULL)
        return;

    if (conn->socket >= 0)
    {
#ifdef __WIN32
        closesocket(conn->socket);
#else
        close(conn->socket);
#endif
        conn->socket = -1;
    }

    th_free(conn->host);
    th_free(conn->proxy.host);

    conn->status = NN_CONN_CLOSED;

    th_free(conn);
}


BOOL nn_conn_send_buf(nn_conn_t *conn, const char *buf, const size_t len)
{
    size_t bufLeft = len;
    const char *bufPtr = buf;

    while (bufLeft > 0)
    {
        ssize_t bufSent;
        bufSent = send(conn->socket, bufPtr, bufLeft, 0);
        if (bufSent < 0)
        {
            conn->err = nn_get_socket_errno();
            nn_conn_err(conn, "nn_conn_send_buf() failed: %s", nn_get_socket_errstr(conn->err));
            return FALSE;
        }
        bufLeft -= bufSent;
        bufPtr += bufSent;
    }

    return TRUE;
}

void nn_conn_reset(nn_conn_t *conn)
{
    if (conn != NULL)
    {
        conn->ptr = conn->in_ptr = conn->buf;
        conn->got_bytes = conn->total_bytes = 0;
    }
}


int nn_conn_pull(nn_conn_t *conn)
{
    int result;
    struct timeval socktv;
    fd_set tmpfds;

    if (conn == NULL)
        return -10;

    // Shift the input buffer
    if (conn->ptr > conn->buf)
    {
        size_t left = conn->in_ptr - conn->ptr;
        if (left > 0)
        {
            size_t moved = conn->ptr - conn->buf;
            memmove(conn->buf, conn->ptr, left);
            conn->ptr = conn->buf;
            conn->in_ptr -= moved;
            conn->total_bytes -= moved;
        }
        else
            nn_conn_reset(conn);        
    }

    // Check for incoming data
    socktv.tv_sec = 0;
    socktv.tv_usec = NN_DELAY_USEC;
    tmpfds = conn->sockfds;

    if ((result = select(conn->socket + 1, &tmpfds, NULL, NULL, &socktv)) == -1)
    {
        int err = nn_get_socket_errno();
        if (err != EINTR)
        {
            conn->err = err;
            nn_conn_err(conn, "Error occured in select(%d, sockfds): %d, %s\n",
                socket, conn->err, nn_get_socket_errstr(conn->err));
            return -1;
        }
    }
    else if (FD_ISSET(conn->socket, &tmpfds))
    {
        conn->got_bytes = recv(conn->socket, conn->in_ptr, NN_CONNBUF_SIZE - conn->total_bytes, 0);
        if (conn->got_bytes < 0)
        {
            conn->err = nn_get_socket_errno();
            nn_conn_err(conn, "Error in recv: %d, %s\n", conn->err, nn_get_socket_errstr(conn->err));
            return -2;
        }
        else if (conn->got_bytes == 0)
        {
            nn_conn_err(conn, "Server closed connection.\n");
            conn->status = NN_CONN_CLOSED;
            return -3;
        }
        else
        {
            conn->total_bytes += conn->got_bytes;
            conn->in_ptr += conn->got_bytes;
            return 0;
        }
    }

    return 1;
}


BOOL nn_conn_check(nn_conn_t *conn)
{
    if (conn == NULL)
        return FALSE;

    return conn->err == 0 && conn->status == NN_CONN_OPEN;
}


BOOL nn_network_init(void)
{
#ifdef __WIN32
    // Initialize WinSock, if needed
    WSADATA wsaData;
    int err = WSAStartup(0x0101, &wsaData);
    if (err != 0)
    {
        THERR("Could not initialize WinSock library (err=%d).\n", err);
        return FALSE;
    }
#endif

    nn_network_inited = TRUE;
    return TRUE;
}


void nn_network_close(void)
{
    if (nn_network_inited)
    {
#ifdef __WIN32
        WSACleanup();
#endif
    }

    nn_network_inited = FALSE;
}


BOOL nn_conn_buf_check(nn_conn_t *conn, size_t n)
{
    return conn && (conn->ptr + n <= conn->in_ptr);
}


BOOL nn_conn_buf_skip(nn_conn_t *conn, size_t n)
{
    if (nn_conn_buf_check(conn, n))
    {
        conn->ptr += n;
        return TRUE;
    }
    else
        return FALSE;
}


int nn_conn_buf_strncmp(nn_conn_t *conn, const char *str, const size_t n)
{
    int ret;
    if (!nn_conn_buf_check(conn, n))
        return -1;

    if ((ret = strncmp(conn->ptr, str, n)) == 0)
    {
        conn->ptr += n;
        return 0;
    }
    else
        return ret;
}


int nn_conn_buf_strcmp(nn_conn_t *conn, const char *str)
{
    return nn_conn_buf_strncmp(conn, str, strlen(str));
}


char *nn_conn_buf_strstr(nn_conn_t *conn, const char *str)
{
    char *pos;
    size_t n = strlen(str);

    if (nn_conn_buf_check(conn, n) && ((pos = strstr(conn->ptr, str)) != NULL))
    {
        conn->ptr = pos + n;
        return pos;
    }
    else
        return NULL;
}


BOOL nn_conn_send_msg(nn_conn_t *conn, const char *user, const char *str)
{
    char *msg;

    if (str == NULL)
        return FALSE;
    
    msg = th_strdup_printf("<USER>%s</USER><MESSAGE>%s</MESSAGE>", user, str);

    if (msg != NULL)
    {
        BOOL ret = nn_conn_send_buf(conn, msg, strlen(msg) + 1);
        th_free(msg);
        return ret;
    }
    else
        return FALSE;
}


BOOL nn_conn_send_msg_v(nn_conn_t *conn, const char *user, const char *fmt, ...)
{
    BOOL res;
    char *tmp;
    va_list ap;

    va_start(ap, fmt);
    tmp = th_strdup_vprintf(fmt, ap);
    va_end(ap);

    res = nn_conn_send_msg(conn, user, tmp);
    th_free(tmp);
    return res;
}


void nn_conn_dump_buffer(FILE *f, nn_conn_t *conn)
{
    char *p;
    size_t offs, left;
    
    fprintf(f,
    "\n--------------------------------------------------------------\n"
    "err=%d, status=%d, got_bytes=%d, total_bytes=%d\n"
    "buf=0x%p, in_ptr=0x%04x, ptr=0x%04x\n",
    conn->err, conn->status, conn->got_bytes, conn->total_bytes,
    conn->buf, conn->in_ptr - conn->buf, conn->ptr - conn->buf);
    
    // Dump buffer contents as a hexdump
    for (offs = 0, left = conn->total_bytes, p = conn->buf; p < conn->in_ptr;)
    {
        char buf[NN_DUMP_BYTES + 1];
        size_t bufoffs, amount = left < NN_DUMP_BYTES ? left : NN_DUMP_BYTES;
        left -= amount;

        // Dump offs | xx xx xx xx | and fill string
        fprintf(f, "%04x | ", offs);
        for (bufoffs = 0; bufoffs < amount; offs++, bufoffs++, p++)
        {
            fprintf(f, "%02x ", *p);
            buf[bufoffs] = th_isprint(*p) ? *p : '.';
        }
        buf[bufoffs] = 0;

        // Add padding
        for (; bufoffs < NN_DUMP_BYTES; bufoffs++)
            fprintf(f, "   ");

        // Print string
        fprintf(f, "| %s\n", buf);
    }
}