// SimpleChat - simplechat.c
// Created by Arsham Skrenes on 2014-02-06.
// An example of a chat client and server using UDP over IPv4 or IPv6
// To compile in Linux: gcc -o simplechat simplechat.c -Wall -std=gnu99
// To compile in OS X:  gcc -o simplechat simplechat.c -Wall

#include <arpa/inet.h>
#include <netdb.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define PROG_NAME    "SimpleChat"
#define PROG_VERSION "v1.0"

#define FAMILY AF_UNSPEC // use IPv4 (AF_INET) or IPv6 (AF_INET6)
#define BUFFER 1500      // must be a margin larger than MAX_ALIAS + MAX_MESSG

#define MAX_ALIAS 20     // longest chat-alias (including null character)
#define MAX_MESSG 1000   // longest chat-message (including null character)
#define MAX_IDLE  10     // seconds of idle before warning (boot after 2x this)

// client structure (linked list element)
struct client {
    char alias[MAX_ALIAS];
    struct sockaddr_storage address;
    time_t lastActivity;
    time_t lastWarning;
    struct client *next;
};

// global variables
int g_family = FAMILY;
int g_signal = 1;
int g_sock   = -1;

// signal handler to terminate loop and then exit process
void signalHandler(int signal)
{
    if (g_signal)
    {
        g_signal = 0;  // setting signal to zero breaks out of loops
        close(g_sock); // closing the socket breaks blocked calls
        g_sock = -1;   // any 0 or greater value is a valid file descriptor
    }
}

// function prototypes
int checkArgs(int argc, const char * argv[]);
int setupSocket(const char *address, const char *port, int ai_socktype);
void getAddress(struct sockaddr_storage *addr, char *ip, in_port_t *port);
int runClient(const char *address, const char *port, const char *alias);
int runServer(const char *port);
struct client* findClient(struct client *list, struct sockaddr_storage *addr);

int main(int argc, const char * argv[])
{
    // check arguments
    if (checkArgs(argc, argv) == EXIT_FAILURE)
        return EXIT_FAILURE;

    // setup signal handler to cleanly terminate application
    signal(SIGTERM, signalHandler);
    signal(SIGINT, signalHandler);

    // run the appropriate role
    if (argc == 5)
        return runClient(argv[2], argv[3], argv[4]);
    else
        return runServer(argv[2]);
}

void printPortHelp()
{
    fprintf(stderr, "Bad port argument. Valid range is 1024-65535.\n");
}

// Checks arguments and returns EXIT_SUCCESS or EXIT_FAILURE
int checkArgs(int argc, const char * argv[])
{
    // check that there are 3 or 4 arguments if run as server or
    // check that there are 5 arguments if run as client
    if (argc < 3 || argc > 5 ||
        (argc <= 4 && strcmp(argv[1], "server") != 0) ||
        (argc == 5 && strcmp(argv[1], "client") != 0))
    {
        fprintf(stderr,
                "Usage: %s client <server-address> <server-port> <chat-alias>\n"
                "       %s server <listen-port> [IPv4|IPv6]\n",
                PROG_NAME,
                PROG_NAME);
        return EXIT_FAILURE;
    }

    // check server/client specific arguments
    int port;
    char buf[BUFFER];
    if (argc == 5)
    {
        // check port in client argument
        if (sscanf(argv[3], "%d%s", &port, buf) != 1)
        {
            printPortHelp();
            return EXIT_FAILURE;
        }

        // check the length of the chat-alias
        if (strlen(argv[4]) >= MAX_ALIAS)
        {
            fprintf(stderr, "The chat-alias cannot exceed %d characters.\n",
                    MAX_ALIAS-1);
            return EXIT_FAILURE;
        }
    }
    else
    {
        // check port in server argument
        if (sscanf(argv[2], "%d%s", &port, buf) != 1)
        {
            printPortHelp();
            return EXIT_FAILURE;
        }

        // check for IP version argument
        if (argc == 4)
        {
            if (strcasecmp(argv[3], "IPv4") == 0)
                g_family = AF_INET;
            else if (strcasecmp(argv[3], "IPv6") == 0)
                g_family = AF_INET6;
            else
            {
                fprintf(stderr, "The optional IP version argument must be "
                        "either IPv4 or IPv6.\n");
                return EXIT_FAILURE;
            }
        }
    }

    // check port range
    if (port < 1024 || port > 65535)
    {
        printPortHelp();
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}

// runs the client role; returns EXIT_SUCCESS or EXIT_FAILURE
int runClient(const char *address, const char *port, const char *alias)
{
    printf("%s %s client (press CTRL+C to exit)\n", PROG_NAME, PROG_VERSION);

    // create socket
    g_sock = setupSocket(address, port, SOCK_DGRAM);
    if (g_sock == -1) return EXIT_FAILURE;

    // prepare to greet the server
    char buf[BUFFER];
    sprintf(buf, "HELLO %s %s %s", PROG_NAME, PROG_VERSION, alias);
    fd_set read_fds;         // create a file descriptor set

    // keep sending greeting to chat server until a reply is received
    for (int i = 1; g_signal; i++)
    {
        // greet the server
        printf("\rGreeting server (attempt %d)...", i);
        fflush(stdout);
        if (send(g_sock, buf, strlen(buf), 0) == -1)
        {
            perror("\nFailed to send a greeting to the server");
            close(g_sock);
            return EXIT_FAILURE;
        }

        // create and set time-out period for select; may change after each call
        struct timeval timeout;
        timeout.tv_sec  = 1;
        timeout.tv_usec = 0;

        // wait for a response from the server to come in on the socket
        FD_ZERO(&read_fds);      // zero out the set
        FD_SET(g_sock, &read_fds); // put socket in set
        select(g_sock+1, &read_fds, NULL, NULL, &timeout);

        // check if a packet arrived (or if it was a timeout)
        if (FD_ISSET(g_sock, &read_fds)) break;
    }
    printf("\n");

    // check if the user invoked the breakout of the loop
    if (!g_signal)
    {
        printf("\rExiting.\n");
        close(g_sock);
        return EXIT_SUCCESS;
    }

    // get the packet from the server
    ssize_t num_bytes;
    if ((num_bytes = recv(g_sock, buf, BUFFER-1, 0)) < 0)
    {
        perror("Failed to receive a response from the server");
        close(g_sock);
        return EXIT_FAILURE;
    }
    buf[num_bytes] = '\0'; // null-terminate buffer

    // check the packet method
    if (strncmp(buf, "REJCT", 5) == 0)
    {
        printf("%s server rejected the greeting: %s\n", PROG_NAME, buf+6);
        close(g_sock);
        return EXIT_SUCCESS;
    }
    else if (strncmp(buf, "HELLO", 5) != 0)
    {
        fprintf(stderr,
                "Failed to receive a proper response from server.\n"
                "Response: %s\n",
                buf);
        close(g_sock);
        return EXIT_FAILURE;
    }

    printf("Successfully joined %s server at %s:%s\n",
           PROG_NAME, address, port);

    setvbuf(stdin, NULL, _IONBF, 0); //turn off buffering

    while (g_signal)
    {
        // wait for a response from the server or keyboard
        FD_ZERO(&read_fds);              // zero out the set
        FD_SET(g_sock, &read_fds);         // put socket in set
        FD_SET(STDIN_FILENO, &read_fds); // put stdin (standard input) in set
        select(g_sock+1, &read_fds, NULL, NULL, NULL);
        if (!g_signal) continue;

        // check if the keyboard was activated
        if (FD_ISSET(STDIN_FILENO, &read_fds))
        {
            // get input from stdin
            sprintf(buf, "MESSG ");
            fgets(buf+strlen(buf), MAX_MESSG, stdin);
            buf[strlen(buf)-1] = '\0'; // remove the \n from fgets

            // remove the input from the console
            printf("\033[<1>A\r"); // move up one line, then carriage return
            for (int i = 0; i < strlen(buf)-6; i++)
                printf(" ");       // insert spaces for input
            printf("\r");          // carriage return again
            fflush(stdout);        // flush output to reflect changes

            // send the input to the chat server
            send(g_sock, buf, strlen(buf), 0);
        }

        // check if a packet arrived
        if (FD_ISSET(g_sock, &read_fds))
        {
            if ((num_bytes = recv(g_sock, buf, BUFFER-1, 0)) < 0)
            {
                perror("Failed to retrieve a packet from the server");
                close(g_sock);
                return EXIT_FAILURE;
            }
            buf[num_bytes] = '\0'; // null-terminate the response

            // check the packet method
            if (strncmp(buf, "REJCT", 5) == 0)
            {
                printf("%s server sent a rejection: %s\n", PROG_NAME, buf+6);
                close(g_sock);
                return EXIT_SUCCESS;
            }
            else if (strncmp(buf, "WARNG", 5) == 0)
            {
                // server sending an idle warning; send back a keep-alive
                sprintf(buf, "ALIVE");
                send(g_sock, buf, strlen(buf), 0);
            }
            else if (strncmp(buf, "MESSG", 5) == 0)
            {
                printf("\r%s\n", buf+6);
            }
            else
            {
                fprintf(stderr,
                        "Failed to receive a proper response from server.\n"
                        "Response: %s\n",
                        buf);
                close(g_sock);
                return EXIT_FAILURE;
            }
        }
    }

    // clean up and exit
    close(g_sock);
    printf("\rExiting.\n");
    return EXIT_SUCCESS;
}

int runServer(const char *port)
{
    printf("%s %s server (press CTRL+C to exit)\n", PROG_NAME, PROG_VERSION);

    // create socket
    g_sock = setupSocket(NULL, port, SOCK_DGRAM);
    if (g_sock == -1) return EXIT_FAILURE;

    // create variables and structures
    struct client *client_list = NULL, *last_client, *curr_client;
    time_t curr_time;
    ssize_t num_bytes;
    char buf[BUFFER], buf2[BUFFER];
    struct sockaddr_storage client_addr;
    socklen_t client_addr_len = sizeof client_addr;
    char client_ip[INET6_ADDRSTRLEN]; // long enough for IPv4 and IPv6
    unsigned short client_port;

    int ret = EXIT_SUCCESS;
    while (g_signal)
    {
        // check client activity and warnings
        curr_time = time(NULL);
        curr_client = client_list;
        while (curr_client)
        {
            // prepare warning message
            sprintf(buf, "WARNG");

            // first check warnings
            if (curr_client->lastWarning)
            {
                if (curr_client->lastWarning < curr_time - MAX_IDLE)
                {
                    // boot client; they've been absent for at least 2*MAX_IDLE
                    getAddress(&curr_client->address, client_ip, &client_port);
                    printf("Booting \"%s\" (%s:%u) for inactivity.\n",
                           curr_client->alias,
                           client_ip,
                           client_port);
                    if (curr_client == client_list)
                    {
                        client_list = client_list->next;
                        free(curr_client);
                        curr_client = client_list;
                        last_client = client_list;
                    }
                    else
                    {
                        last_client->next = curr_client->next;
                        free(curr_client);
                        curr_client = last_client->next;
                    }
                    continue;
                }
                else
                {
                    // send the client another warning
                    sendto(g_sock,
                           buf,
                           strlen(buf),
                           0,
                           (struct sockaddr *)&curr_client->address,
                           client_addr_len);
                }
            }
            else if (curr_client->lastActivity < curr_time - MAX_IDLE)
            {
                // give the client a warning
                curr_client->lastWarning = curr_time;
                sendto(g_sock,
                       buf,
                       strlen(buf),
                       0,
                       (struct sockaddr *)&curr_client->address,
                       client_addr_len);
            }

            last_client = curr_client;
            curr_client = curr_client->next;
        }

        // listen for incoming chat clients
        if ((num_bytes = recvfrom(g_sock,
                                  buf,
                                  BUFFER-1,
                                  0,
                                  (struct sockaddr *)&client_addr,
                                  &client_addr_len)) == -1)
        {
            if (!g_signal) continue;
            perror("Failed to receive data on the socket");
            ret = EXIT_FAILURE;
            break;
        }
        buf[num_bytes] = '\0'; // null-terminate buffer
        getAddress(&client_addr, client_ip, &client_port);

        // parse the message
        if (strncmp(buf, "HELLO " PROG_NAME " ", 6+strlen(PROG_NAME)+1) == 0)
        {
            // check version
            int offset = 6+strlen(PROG_NAME)+1;
            if (strncmp(buf+offset,
                        PROG_VERSION " ",
                        strlen(PROG_VERSION)+1) != 0)
            {
                printf("Rejecting client (%s:%u) with wrong version.\n",
                       client_ip,
                       client_port);
                sprintf(buf,
                        "REJCT Wrong %s version. Require %s.",
                        PROG_NAME,
                        PROG_VERSION);
                sendto(g_sock,
                       buf,
                       strlen(buf),
                       0,
                       (struct sockaddr *)&client_addr,
                       client_addr_len);
                continue;
            }
            offset += strlen(PROG_VERSION)+1; // offset to start of alias

            // check if alias is unique
            curr_client = client_list;
            while (curr_client)
            {
                if (strcmp(curr_client->alias, buf+offset) == 0)
                {
                    // get string ip and port from curr_client
                    char tmp_ip[INET6_ADDRSTRLEN];
                    unsigned short tmp_port;
                    getAddress(&curr_client->address, tmp_ip, &tmp_port);

                    // check that this is not the same client, which could
                    // happen if the former "HELLO Welcome" message did not make
                    // it back to the client
                    if (strcmp(client_ip, tmp_ip) == 0 &&
                        client_port == tmp_port)
                    {
                        printf("\"%s\" (%s:%u) just rejoined.\n",
                               curr_client->alias,
                               client_ip,
                               client_port);
                        curr_client = NULL;
                        offset = 0; // use this to indicate this situation
                        break;
                    }

                    printf("Rejecting client (%s:%u) with same chat-alias as "
                           "\"%s\" (%s:%u).\n",
                           client_ip,
                           client_port,
                           curr_client->alias,
                           tmp_ip,
                           tmp_port);
                    break;
                }
                curr_client = curr_client->next;
            }

            if (curr_client)
            {
                // alias is not unique; reject new client
                sprintf(buf, "REJCT Username already taken.");
                sendto(g_sock,
                       buf,
                       strlen(buf),
                       0,
                       (struct sockaddr *)&client_addr,
                       client_addr_len);
                continue;
            }

            if (offset) // if offset == 0, then just re-send "HELLO Welcome"
            {
                // add client to client_list
                curr_client = malloc(sizeof(struct client));
                strcpy(curr_client->alias, buf+offset);
                memcpy(&curr_client->address,&client_addr,sizeof(client_addr));
                curr_client->lastActivity = time(NULL);
                curr_client->lastWarning = 0;
                curr_client->next = client_list;
                client_list = curr_client;

                printf("\"%s\" (%s:%u) just joined.\n",
                       curr_client->alias,
                       client_ip,
                       client_port);
            }

            // welcome client
            sprintf(buf, "HELLO Welcome to the %s server!", PROG_NAME);
            sendto(g_sock,
                   buf,
                   strlen(buf),
                   0,
                   (struct sockaddr *)&client_addr,
                   client_addr_len);
            continue;
        }
        else if (strncmp(buf, "MESSG", 5) == 0)
        {
            // check if this is a valid chat client
            curr_client = findClient(client_list, &client_addr);
            if (!curr_client)
            {
                sprintf(buf, "REJCT You are not eligible to chat.");
                sendto(g_sock,
                       buf,
                       strlen(buf),
                       0,
                       (struct sockaddr *)&client_addr,
                       client_addr_len);
                continue;
            }

            // send the message to everyone (including the sender)
            sprintf(buf2, "MESSG %s: %s", curr_client->alias, buf+6);
            curr_client = client_list;
            while (curr_client)
            {
                sendto(g_sock,
                       buf2,
                       strlen(buf2),
                       0,
                       (struct sockaddr *)&curr_client->address,
                       client_addr_len);
                curr_client = curr_client->next;
            }
        }
        else if (strncmp(buf, "ALIVE", 5) == 0)
        {
            // update timers if they are a valid client
            curr_client = findClient(client_list, &client_addr);
            if (curr_client)
            {
                curr_client->lastActivity = time(NULL);
                curr_client->lastWarning = 0;
            }
            continue; // ignore if not/no-longer a client
        }
        else
        {
            printf("Ignoring a bad datagram from %s:%u.\nResponse: %s\n",
                   client_ip,
                   client_port,
                   buf);
            continue;
        }
    }

    // cleanup
    close(g_sock);
    while (client_list)
    {
        curr_client = client_list;
        client_list = client_list->next;
        free(curr_client);
    }

    printf("\rExiting.\n");
    return ret;
}

// creates a socket connected to address/port; if address is null it then
// creates a socket using local address and bound to port
// ai_socktype can be SOCK_STREAM or SOCK_DGRAM
// returns socket descriptor or -1 on failure; IPv4 & IPv6 capable
int setupSocket(const char *address, const char *port, int ai_socktype)
{
    int sock, ret;
    struct addrinfo hints, *server_info_list, *server_info;

    // set hints for getaddrinfo()
    memset(&hints, 0, sizeof hints);  // zero hints structure
    hints.ai_family    = g_family;    // can be IPv4 or IPv6
    hints.ai_socktype  = ai_socktype; // can be TCP or UDP
    if (!address)
        hints.ai_flags = AI_PASSIVE;  // use local IP address

    if ((ret = getaddrinfo(address, port, &hints, &server_info_list)) != 0)
    {
        fprintf(stderr, "Failed to get address information: %s\n",
                gai_strerror(ret));
        return -1;
    }

    // loop through all the results trying to make and bind the socket
    for(server_info = server_info_list;
        server_info != NULL;
        server_info = server_info->ai_next)
    {
        if ((sock = socket(server_info->ai_family,
                           server_info->ai_socktype,
                           server_info->ai_protocol)) == -1) continue;

        if (!address)
        {
            if (bind(sock, server_info->ai_addr, server_info->ai_addrlen) == -1)
            {
                close(sock);
                continue;
            }
        }

        break;
    }

    // check if the addresses were exhausted before a successful socket
    if (!server_info)
    {
        // check if the above loop would have an errno (would report last error)
        if (!server_info_list)
            fprintf(stderr, "Failed to create a socket\n");
        else
            perror("Failed to create a socket");
        freeaddrinfo(server_info_list);
        return -1;
    }

    if (address)
    {
        // connect to the server
        if (connect(sock, server_info->ai_addr, server_info->ai_addrlen) == -1)
        {
            perror("Failed to connect to the server");
            close(sock);
            freeaddrinfo(server_info_list);
            return -1;
        }
    }
    else
    {
        // report address and port for server
        char ip[INET6_ADDRSTRLEN];
        unsigned short port;
        getAddress((struct sockaddr_storage *)server_info->ai_addr, ip, &port);
        printf("Waiting for clients at %s:%u\n", ip, port);
    }

    freeaddrinfo(server_info_list);
    return sock;
}

// takes struct sockaddr_storage and sets ip and port;
// compatible with IPv4 and IPv6; sizeof ip >= INET6_ADDRSTRLEN
void getAddress(struct sockaddr_storage *addr, char *ip, in_port_t *port)
{
    if (((struct sockaddr*)addr)->sa_family == AF_INET)
    {
        // get IPv4 address
        inet_ntop(addr->ss_family,
                  &(((struct sockaddr_in*)addr)->sin_addr),
                  ip,
                  INET6_ADDRSTRLEN);

        // get port
        *port = ntohs(((struct sockaddr_in*)addr)->sin_port);
    }
    else
    {
        // get IPv6 address
        inet_ntop(addr->ss_family,
                  &(((struct sockaddr_in6*)addr)->sin6_addr),
                  ip,
                  INET6_ADDRSTRLEN);

        // get port
        *port = ntohs(((struct sockaddr_in6*)addr)->sin6_port);
    }
}

// finds the client associated with the address; returns NULL if not found
struct client* findClient(struct client *list, struct sockaddr_storage *addr)
{
    while (list)
    {
        if (memcmp(&list->address, addr, sizeof(struct sockaddr_storage)) == 0)
            break;
        list = list->next;
    }
    return list;
}