Sleds/K2Client/main.cpp

/*
    Copyright (c)  2013 IOnU Security Inc. All rights reserved
    Created August 2013 by Kendrick Webster

    K2Client/main.c - Simple test application using the API exposed by the
    platform_binding module.
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifndef WIN32
#include <unistd.h>
#include <sys/select.h>
#include <termios.h>
#else
#include <Winsock2.h>
#include <conio.h>
#endif
#include <ctype.h>
#include "K2IPC.h"
#include "constants.h"
#include "util.h"
#include "platform_binding.h"

#ifdef WIN32
#define snprintf _snprintf
#endif

/* testing functions in platform_binding.c (unpublicized, not in header) */
void K2_Newsock(void);
void K2_Stop_NoLogout(void);

/* testing variables in K2Client.c (unpublicized, not in header) */
extern int heartbeat_interval_seconds;

static int urn_is_dynamic = 0;
static char* urn = "urn.client.test";
static const char* servers = "";
#ifndef WIN32
static struct termios orig_termios;
#endif
static int done = FALSE, server_test = FALSE, verbose = FALSE, server_test_result = 1;

#define OFFICE_URN_LIST_BASE "office.urn."
#define OFFICE_URN_LIST_ELEM(suffix) OFFICE_URN_LIST_BASE #suffix "::,"
#define OULE(s) OFFICE_URN_LIST_ELEM(s)
#define OFFICE_URN_LIST_SET(p) \
    OULE(p##A) OULE(p##B) OULE(p##C) OULE(p##D) OULE(p##E) OULE(p##F) OULE(p##G) OULE(p##H) OULE(p##I) \
    OULE(p##J) OULE(p##K) OULE(p##L) OULE(p##M) OULE(p##N) OULE(p##O) OULE(p##P) OULE(p##Q) OULE(p##R) \
    OULE(p##S) OULE(p##T) OULE(p##U) OULE(p##V) OULE(p##W) OULE(p##X) OULE(p##Y) OULE(p##Z)
#define OULS(p) OFFICE_URN_LIST_SET(p)
#define OFFICE_URN_BIG_LIST \
    OULS(A) OULS(B) OULS(C) OULS(D) OULS(E) OULS(F) OULS(G) OULS(H) OULS(I) \
    OULS(J) OULS(K) OULS(L) OULS(M) OULS(N) OULS(O) OULS(P) OULS(Q) OULS(R) \
    OULS(S) OULS(T) OULS(U) OULS(V) OULS(W) OULS(X) OULS(Y) OULS(Z)
#define OFFICE_URN_MEDIUM_LIST OULS(A) OULS(B) OULS(C) OULS(D) OULS(E) OULS(F)
#define OFFICE_URN_LIST OULS(_)

#ifndef WIN32
static void reset_terminal_mode()
{
    tcsetattr(0, TCSANOW, &orig_termios);
}

static void set_conio_terminal_mode()
{
    struct termios new_termios;
    tcgetattr(0, &orig_termios);
    memcpy(&new_termios, &orig_termios, sizeof(new_termios));
    atexit(reset_terminal_mode);
    new_termios.c_iflag &= ~(IGNBRK | BRKINT | PARMRK | ISTRIP
                    | INLCR | IGNCR | ICRNL | IXON);
    new_termios.c_lflag &= ~(ECHO | ECHONL | ICANON | ISIG | IEXTEN);
    new_termios.c_cflag &= ~(CSIZE | PARENB);
    new_termios.c_cflag |= CS8;
    tcsetattr(0, TCSANOW, &new_termios);
}

static int kbhit()
{
    struct timeval tv = { 0L, 0L };
    fd_set fds;
    FD_ZERO(&fds);
    FD_SET(0, &fds);
    return select(1, &fds, NULL, NULL, &tv);
}

static int getch()
{
    int r;
    unsigned char c;
    return (0 > (r = read(0, &c, sizeof(c)))) ? r : c;
}
#else
#define getch _getch
#define kbhit _kbhit
#endif

static void onMessage(const char* m)
{
    if (server_test)
    {
        if (0 == strcmp(m, "k2cli:logged in"))
        {
            server_test_result = 0;
            done = TRUE;
        }
        else if (0 == strcmp(m, "k2cli:ERROR: not connected"))
        {
            server_test_result = 1;
            done = TRUE;
        }
        else if (strstr(m, "k2cli:ERROR:"))
        {
            server_test_result = 2;
            done = TRUE;
            printf("[ERROR]%s\n", m + 12);
        }
    }
    else
    {
        printf("%s\n", m);
    }
}

static void showHelp(void)
{
    printf("urn     = \"%s\"\n", urn);
    printf("servers = \"%s\"\n", servers);
    printf("heartbeat interval = %d seconds\n", heartbeat_interval_seconds);
    printf("keyboard commands:\n");
    printf("   ?     show this help text\n");
    printf("   ESC   exit\n");
    printf("   Q     exit\n");
    printf("   L     log out\n");
    printf("   I     log in\n");
    printf("   D     trigger daemon database poll\n");
    printf("   S     call K2_Start() again\n");
    printf("   U     call K2_Start() with a new URN\n");
    printf("   O     call K2_StatusSelect() with a list of test offices\n");
    printf("   B     call K2_StatusSelect() with a big list of offices\n");
    printf("   M     call K2_StatusSelect() with a medium list of offices\n");
    printf("   C     call K2_StatusSelect() with an empty list of offices\n");
    printf("   N     new socket (new port number) to test daemon exception handling\n");
    printf("   H     stop thread without logging out to cause a heartbeat timeout\n");
    printf("   T     query daemon CPU and memory usage statistics\n");
}

static void onKeypress(int c)
{
    char* oldurn;
    char buf[K2IPC_MAX_STRING_LENGTH];
    switch (toupper(c))
    {
        case '?':
            printf("Displaying help:\n");
            showHelp();
            break;

        case 27:
        case 'Q':
            done = TRUE;
            break;

        case 'L':
            printf("Logging out\n");
            K2_Stop();
            break;

        case 'I':
            printf("Logging in\n");
            K2_Restart();
            break;

        case 'D':
            printf("Triggering daemon database poll\n");
            K2_Poll_DB();
            break;

        case 'S':
            printf("Calling K2_Start(...)\n");
            K2_Start(urn, servers, onMessage);
            break;

        case 'U':
            printf("Calling K2_Start(...) with new URN\n");
            oldurn = urn;
            c = (int)(strlen(oldurn) + 3);
            urn = (char *)malloc(c);
            snprintf(urn, c, "%s.n", oldurn);
            K2_Start(urn, servers, onMessage);
            if (urn_is_dynamic)
            {
                free(oldurn);
            }
            urn_is_dynamic = 1;
            break;

        case 'O':
            printf("Calling K2_StatusSelect(...) with small list\n");
            K2_StatusSelect(OFFICE_URN_LIST);
            break;

        case 'B':
            printf("Calling K2_StatusSelect(...) with big list\n");
            K2_StatusSelect(OFFICE_URN_BIG_LIST);
            break;

        case 'M':
            printf("Calling K2_StatusSelect(...) with medium list\n");
            K2_StatusSelect(OFFICE_URN_MEDIUM_LIST);
            break;

        case 'C':
            printf("Calling K2_StatusSelect(NULL) to clear subscriptions\n");
            K2_StatusSelect(NULL);
            break;

        case 'N':
            printf("Calling K2_Newsock()\n");
            K2_Newsock();
            break;

        case 'H':
            printf("Stopping thread without logging out\n");
            K2_Stop_NoLogout();
            break;

        case 'T':
            printf("Querying daemon CPU and memory usage statistics ...\n");
            if (K2_GetStats(buf, sizeof(buf)))
            {
                printf("%s\n", buf);
            }
            else
            {
                printf("FAILED to get statistics\n");
            }
            break;
    }
}

int main(int argc, char** argv)
{
    int c;

    if (argc == 2)
    {
        server_test = TRUE;
        servers = argv[1];
    }
    else if (argc == 3)
    {
        if ((0 == strcmp(argv[1], "-health-check")) || (0 == strcmp(argv[1], "--health-check")))
        {
            server_test = TRUE;
            verbose = TRUE;
            servers = argv[2];
        }
        else
        {
            urn = argv[1];
            servers = argv[2];
        }
    }
    else if (argc == 4)
    {
        urn = argv[1];
        servers = argv[2];
        heartbeat_interval_seconds = atoi(argv[3]);
    }
    else
    {
        printf("--- Usage ---\n");
        printf("Interactive:          K2Client_test <urn> <servers> [<heartbeat_seconds>]\n");
        printf("Server check:         K2Client_test -health-check <server>\n");
        printf("Server check silent:  K2Client_test <server>\n");
        printf("\n");
        printf("The server check modes will attempt to connect to <server>, then exit with\n");
        printf("status code 0 if the attempt succeeds, or non-zero if the attempt fails.\n");
        printf("\n");
        printf("--- Interactive mode ---\n");
        printf("Interactive mode behaves like a normal connected client.  It logs status to\n");
        printf("stdout and allows the user to issue keyboard commands from a small menu.\n");
        printf("\n");
        printf("<urn> may be any string that doesn't collide with other connected clients.\n");
        printf("For example: 'test.urn.1' or 'urn.test.my_name.123'\n");
        printf("\n");
        printf("<servers> consists of one or more hosts delimited by '%c'.\n", K2_SERVERS_SUBSTRING_DELIMITER);
        printf("Each host may have an optional port number after a colon ':' delimiter.\n");
        printf("If the port number is omitted, it defaults to UDP port %u.\n", K2IPC_DEFAULT_UDP_PORT);
        printf("\n");
        printf("Hosts on the following port(s) are considered proxies (HTTP instead of UDP):\n");
#define X(port) printf("   %5u\n", port);
        K2_PROXY_PORTS
#undef X
        printf("\n");
        printf("Connection attempts are made to all non-proxy servers (selected randomly)\n");
        printf("before trying proxy servers (also selected randomly).  The order in which\n");
        printf("hosts are listed in the <servers> string is irrelevant.\n");
        printf("\n");
        printf("<heartbeat_seconds> (optional, default=%u) specifies the nominal heartbeat\n", K2CLI_HEARTBEAT_INTERVAL_SECONDS);
        printf("interval in seconds.  The test client will send heartbeat packets at this\n");
        printf("rate +/- a random variation of up to %u seconds (for de-synchronization).\n", K2CLI_HEARTBEAT_RANGE_SECONDS);
        printf("\n");
        printf("--- Usage examples ---\n");
        printf("Interactive mode:\n");
        printf("K2Client_Test test.urn.1 dev.sl.int\n");
        printf("\n");
        printf("Interactive mode with heartbeat interval:\n");
        printf("K2Client_Test test.heartbeat.35 dev.sl.int 35\n");
        printf("\n");
        printf("Interactive mode with mutliple servers:\n");
        printf("K2Client_Test test.urn.1 \"uat.ionu.nu|uat.ionu.nu:32001|proxy.ionu.nu:80\"\n");
        printf("\n");
        return 1;
    }

    if (!server_test)
    {
        showHelp();
    }

    K2_Start(urn, servers, onMessage);

#ifndef WIN32
    if (!server_test)
    {
        set_conio_terminal_mode();
    }
#endif

    while (!done)
    {
#ifndef WIN32
        usleep(100000);
#else
        Sleep(100);
#endif
        if (!server_test && kbhit())
        {
            c = getch();
            if (c < 0)
            {
                printf("getch() error %d\n", c);
            }
            else
            {
                onKeypress(c);
            }
        }
    }

#ifndef WIN32
    if (!server_test)
    {
        reset_terminal_mode();
    }
#endif

    if (!server_test)
    {
        printf("Exiting\n");
    }

    K2_Stop();

    if (urn_is_dynamic)
    {
        free(urn);
    }

    if (!server_test)
    {
        return 0;
    }
    else
    {
        if (verbose)
        {
            switch (server_test_result)
            {
                case 0:
                    printf("[OK] Daemon communication successful\n");
                    break;

                case 1:
                    printf("[ERROR] Daemon did not respond\n");
                    break;

                default:
                    /* detailed error message was already printed */
                    break;
            }
        }
        return server_test_result;
    }
}