/*
 * Copyright (c) 2014, Franck Cuny <franckcuny@gmail.com>
 *
 * Permission to use, copy, modify, and/or distribute this software for any purpose
 * with or without fee is hereby granted, provided that the above copyright notice
 * and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
 * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
 * FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
 * OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
 * TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF
 * THIS SOFTWARE.
*/

#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <getopt.h>

struct settings {
    int verbose;
    int localport;
    int daemon;
    char *remoteport;
    char *tag;
    char *remotehost;
    char *template;
};

struct peersock {
    int socket;
    socklen_t salen;
    struct sockaddr *sa;
};

struct settings settings;

static void settings_init() {
    settings.verbose = 0;
    settings.daemon = 0;
    settings.localport = 8126;
    settings.remoteport = "8125";
    settings.remotehost = NULL;
    settings.tag = "PROD";
    settings.template = "%s.hosts.%s.";
}

static void usage(char *prog) {
    printf("usage: %s [option]\n\n", prog);
    printf("-p <num>         local UDP port to listen on (default: %d)\n", settings.localport);
    printf("-P <num>         remote UDP port to send to (default: %s)\n", settings.remoteport);
    printf("-H <string>      remote host to send to (required)\n");
    printf("-t               tag to use (default: %s)\n", settings.tag);
    printf("-T               template to apply to the received line (default :%s)\n", settings.template);
    printf("-h               print this help and exit\n");
    printf("-d               run as a daemon\n");
    printf("--verbose        verbose\n");
    return;
}

static int localsock(int port) {
    int sock;
    struct sockaddr_in servaddr;

    if ((sock = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("socket");
        return -1;
    }

    memset((char *)&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
    servaddr.sin_port = htons(port);

    if(bind(sock, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0) {
        perror("bind");
        return -1;
    }
    return sock;
}

static int get_peersock(struct peersock *remote, char *host, char *port) {
    int sock, err;
    struct addrinfo hints, *result, *rp;
    struct sockaddr *sa;
    socklen_t salen;

    memset(&hints, 0, sizeof(struct addrinfo));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;
    hints.ai_protocol = IPPROTO_UDP;

    if ((err = getaddrinfo(host, port, &hints, &result)) != 0) {
        perror("getaddrinfo");
        freeaddrinfo(result);
        return -1;
    }

    for (rp = result; rp != NULL; rp = rp->ai_next) {
        sock = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if (sock >= 0) {
            break;
        }
    }

    sa = malloc(result->ai_addrlen);
    memcpy(sa, result->ai_addr, result->ai_addrlen);
    salen = result->ai_addrlen;

    remote->socket = sock;
    remote->salen = salen;
    remote->sa = sa;

    freeaddrinfo(result);
    return 0;
}

static int daemonize() {
    switch(fork()) {
        case -1:
            return -1;
        case 0:
            break;
        default:
            exit(EXIT_SUCCESS);
    }

    if (settings.verbose == 1) {
        fprintf(stdout, "Daemon is starting ...\n");
    }

    if (setsid() == -1) {
        return -1;
    }

    if(chdir("/") != 0) {
        perror("chdir");
        return -1;
    }

    int fd;
    if ((fd = open("/dev/null", O_RDWR, 0)) != -1) {
        if(dup2(fd, STDIN_FILENO) < 0) {
            perror("dup2 stdin");
            return (-1);
        }
        if(dup2(fd, STDOUT_FILENO) < 0) {
            perror("dup2 stdout");
            return (-1);
        }
        if(dup2(fd, STDERR_FILENO) < 0) {
            perror("dup2 stderr");
            return (-1);
        }

        if (fd > STDERR_FILENO) {
            if(close(fd) < 0) {
                perror("close");
                return (-1);
            }
        }
    }

    return 0;
}

static char *getlinecontent(char mesg[], char template[], int ii, int start) {
    char *line;
    size_t len;

    len = strlen(template) + (ii - start) + 1;
    line = (char *)malloc(len);
    memcpy(line, template, strlen(template) + 1);
    strncat(line, &mesg[start], (ii - start));
    line[len - 1] = '\0';

    return (char *)line;
}

static void proxy_line(struct peersock *p_socket, char mesg[], char template[], int ii, int start) {
    char *line;
    line = getlinecontent(mesg, template, ii, start);
    if (settings.verbose == 1) {
        fprintf(stdout, "Forwarding %s\n", line);
    }
    sendto(p_socket->socket, line, strlen(line), 0, p_socket->sa, p_socket->salen);
    free(line);
    return;
}

int main(int argc, char *argv[]) {

    static struct option long_options[] = {
        {"local-port",  required_argument, NULL, 'p'},
        {"remote-port", required_argument, NULL, 'P'},
        {"remote-host", required_argument, NULL, 'H'},
        {"tag",         required_argument, NULL, 't'},
        {"template",    required_argument, NULL, 'T'},
        {"help",        no_argument,       NULL, 'h'},
        {"verbose",     no_argument,       NULL, 'V'},
        {"daemon",      no_argument,       NULL, 'd'},
        { NULL,         0,                 NULL,  0 }
    };

    settings_init();

    while (1) {
        int option_index = 0;
        int c = getopt_long(argc, argv, "p:P:H:t:T:hvd", long_options, &option_index);

        if (c == -1)
            break;

        switch(c) {
            case 'h':
                usage(argv[0]);
                exit(EXIT_SUCCESS);
            case 'd':
                settings.daemon = 1;
                break;
            case 'p':
                settings.localport = atoi(optarg);
                break;
            case 't':
                settings.tag = optarg;
                break;
            case 'H':
                settings.remotehost = optarg;
                break;
            case 'P':
                settings.remoteport = optarg;
                break;
            case 'T':
                settings.template = optarg;
                break;
            case 'V':
                settings.verbose = 1;
                break;
            default:
                exit(EXIT_FAILURE);
        }
    }

    if (settings.remotehost == NULL) {
        usage(argv[0]);
        exit(EXIT_FAILURE);
    }

    char hostname[255];
    int l_socket;
    struct peersock *p_socket = malloc(sizeof(struct peersock));

    if (gethostname(hostname, sizeof(hostname) - 1) == -1) {
        if (settings.verbose) {
            fprintf(stderr, "Error while discovering hostname.\n");
        }
        hostname[0] = '\0';
    }

    if ((l_socket = localsock(settings.localport)) == -1) {
        exit(EXIT_FAILURE);
    }

    if ((get_peersock(p_socket, settings.remotehost, settings.remoteport)) == -1) {
        fprintf(stderr, "Failed to created a socket to remote host.\n");
        exit(EXIT_FAILURE);
    }

    if (settings.daemon) {
        if ((daemonize()) == -1) {
            fprintf(stderr, "Failed to daemonize()\n");
            exit(EXIT_FAILURE);
        }
    }

    char mesg[64 * 1024];
    char template[255];
    int n;
    sprintf(template, settings.template, settings.tag, hostname);

    while((n = recvfrom(l_socket, mesg, sizeof(mesg), 0, NULL, NULL))) {
        mesg[n] = '\0';
        size_t start = 0;
        int ii;
        for (ii = 0; ii < n; ii++) {
            if (mesg[ii] == '\n') {
                proxy_line(p_socket, mesg, template, ii, start);
                start = ii + 1;
            }
        }
        if ((ii - start) > 0) {
            proxy_line(p_socket, mesg, template, ii, start);
        }
    }

    close(l_socket);
    close(p_socket->socket);
    free(p_socket);

    return EXIT_SUCCESS;
}