// This is a personal academic project. Dear PVS-Studio, please check it.
// PVS-Studio Static Code Analyzer for C, C++ and C#: http://www.viva64.com
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <limits.h>
#include <errno.h>
#include <signal.h>
#include <stdbool.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <fcntl.h>
#include <arpa/inet.h>

#include "args.h"
#include "selector.h"
#include "socks5nio.h"
#include "confignio.h"
#include "buffer.h"
#include "server.h"

#define N(x) (sizeof(x)/sizeof((x)[0]))

static bool done = false;

static void sigterm_handler(const int signal) {
    printf("Signal %d, Cleaning up and exiting\n", signal);
    done = true;
}

static struct socks5args * args;
uint64_t config_token;

int main(int argc, char ** argv) {
    args = malloc(sizeof(struct socks5args));
    parse_args(argc, argv, args);

    const char *err_msg = NULL;
    selector_status ss = SELECTOR_SUCCESS;
    fd_selector selector = NULL;
    struct config * config_ret = NULL;
    struct config * config_ret_ipv6 = NULL;
    int server = -1;
    int server_ipv6 = -1;
    int udp_server = -1;
    int udp_server_ipv6 = -1;

    char * token = getenv("BPROXY_TOKEN");
    if (token != NULL) {
        config_token = strtoul(token, NULL, 16);
    }
    else {
        err_msg = "No token found. Please see README.md";
        goto finally;
    }

    close(STDIN_FILENO);

    unsigned char buf[sizeof(struct in6_addr)];
    int domain = -1;
    if (args->socks_addr != NULL) {
        if (inet_pton(AF_INET, args->socks_addr, buf)) { 
            domain = AF_INET;
        } else if (inet_pton(AF_INET6, args->socks_addr, buf)) {
            domain = AF_INET6;
        } else {
            err_msg = "Incorrect network address";
            goto finally;
        }
    }

    struct sockaddr_in addr;
    if (domain == AF_INET || domain == -1) {
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        if (domain == AF_INET) {
            if (inet_pton(AF_INET, args->socks_addr, &addr.sin_addr) <= 0) {
                err_msg = "Incorrect network address";
                goto finally;
            }
        } else {
            addr.sin_addr.s_addr = htonl(INADDR_ANY);
        }
        addr.sin_port = htons(args->socks_port);

        server = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
        if (server < 0) {
            err_msg = "Unable to create TCP socket";
            goto finally;
        }

        if (args->socks_addr != NULL)
            fprintf(stdout, "Proxy SOCKS listening on %s:%u TCP\n", args->socks_addr, args->socks_port);
        else
            fprintf(stdout, "Proxy SOCKS listening on 0.0.0.0:%u TCP\n", args->socks_port);

        setsockopt(server, SOL_SOCKET, SO_REUSEADDR, &(int){ 1 }, sizeof(int));

        if (bind(server, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
            err_msg = "Unable to bind TCP socket";
            goto finally;
        }

        if (listen(server, 20) < 0) {
            err_msg = "Unable to listen";
            goto finally;
        }
    }

    struct sockaddr_in6 tcp_addr_ipv6;
    if (domain == AF_INET6 || domain == -1) {
        memset(&tcp_addr_ipv6, 0, sizeof(tcp_addr_ipv6));
        tcp_addr_ipv6.sin6_family = AF_INET6;
        if (domain == AF_INET6) {
            if (inet_pton(AF_INET6, args->socks_addr, &tcp_addr_ipv6.sin6_addr) <= 0) {
                err_msg = "Incorrect network address";
                goto finally;
            }
        } else {
            tcp_addr_ipv6.sin6_addr = in6addr_any;
        }
        tcp_addr_ipv6.sin6_port = htons(args->socks_port);

        server_ipv6 = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
        if (server_ipv6 < 0) {
            err_msg = "Unable to create TCP ipv6 socket";
            goto finally;
        }

        if (args->socks_addr != NULL)
            fprintf(stdout, "Proxy SOCKS listening on %s:%u TCP\n", args->socks_addr, args->socks_port);
        else
            fprintf(stdout, "Proxy SOCKS listening on :::%u TCP\n", args->socks_port);

        setsockopt(server_ipv6, SOL_SOCKET, SO_REUSEADDR, &(int){ 1 }, sizeof(int));
        setsockopt(server_ipv6, SOL_IPV6, IPV6_V6ONLY, &(int) { 1 }, sizeof(int));

        if (bind(server_ipv6, (struct sockaddr *) &tcp_addr_ipv6, sizeof(tcp_addr_ipv6)) < 0) {
            err_msg = "Unable to bind TCP ipv6 socket";
            goto finally;
        }

        if (listen(server_ipv6, 20) < 0) {
            err_msg = "Unable to listen";
            goto finally;
        }
    }

    int domain_udp = -1;
    memset(&buf, 0, sizeof(buf));
    if (args->mng_addr != NULL) {
        if (inet_pton(AF_INET, args->mng_addr, buf)) { 
            domain_udp = AF_INET;
        } else if (inet_pton(AF_INET6, args->mng_addr, buf)) {
            domain_udp = AF_INET6;
        } else {
            err_msg = "Incorrect network address";
            goto finally;
        }
    }

    struct sockaddr_in udp_addr;
    if (domain_udp == AF_INET || domain_udp == -1) {
        memset(&udp_addr, 0, sizeof(udp_addr));
        udp_addr.sin_family = AF_INET;
        if (domain_udp == AF_INET) {
            if (inet_pton(AF_INET, args->mng_addr, &udp_addr.sin_addr) <= 0) {
                err_msg = "Incorrect network address";
                goto finally;
            } 
        } else {
            udp_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        }
        udp_addr.sin_port = htons(args->mng_port);

        udp_server = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
        if (udp_server < 0) {
            err_msg = "Unable to create UDP socket";
            goto finally;
        }

        if (args->mng_addr != NULL)
            fprintf(stdout, "Management service listening on %s:%u UDP\n", args->mng_addr, args->mng_port);
        else
            fprintf(stdout, "Management service listening on 127.0.0.1:%u UDP\n", args->mng_port);

        setsockopt(udp_server, SOL_SOCKET, SO_REUSEADDR, &(int){ 1 }, sizeof(int));

        if (bind(udp_server, (struct sockaddr *) &udp_addr, sizeof(udp_addr)) < 0) {
            err_msg = "Unable to bind UDP socket";
            goto finally;
        }

        if (fcntl(udp_server, F_SETFL, O_NONBLOCK) < 0) {
            err_msg = "Unable to put UDP socket into non-blocking mode";
            goto finally;
        }

        config_ret = malloc(sizeof(struct config));
        if (config_ret == NULL) {
            err_msg = "Malloc failed";
            goto finally;
        }
        memset(config_ret, 0x00, sizeof(*config_ret));
        buffer_init(&config_ret->read_buffer, N(config_ret->raw_buff_a), config_ret->raw_buff_a);
        buffer_init(&config_ret->write_buffer, N(config_ret->raw_buff_b), config_ret->raw_buff_b);
    }

    struct sockaddr_in6 udp_addr_ipv6;
    if (domain_udp == AF_INET6 || domain_udp == -1) {
        memset(&udp_addr_ipv6, 0, sizeof(udp_addr_ipv6));
        udp_addr_ipv6.sin6_family = AF_INET6;
        if (domain_udp == AF_INET6) {
            if (inet_pton(AF_INET6, args->mng_addr, &udp_addr_ipv6.sin6_addr) <= 0) {
                err_msg = "Incorrect network address";
                goto finally;
            }
        } else {
            udp_addr_ipv6.sin6_addr = in6addr_loopback;
        }
        udp_addr_ipv6.sin6_port = htons(args->mng_port);

        udp_server_ipv6 = socket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP);
        if (udp_server_ipv6 < 0) {
            err_msg = "Unable to create UDP socket";
            goto finally;
        }

        if (args->mng_addr != NULL)
            fprintf(stdout, "Management service listening on %s:%u UDP\n", args->mng_addr, args->mng_port);
        else
            fprintf(stdout, "Management service listening on ::1:%u UDP\n", args->mng_port);

        setsockopt(udp_server_ipv6, SOL_SOCKET, SO_REUSEADDR, &(int){ 1 }, sizeof(int));
        setsockopt(udp_server_ipv6, SOL_IPV6, IPV6_V6ONLY, &(int) { 1 }, sizeof(int));

        if (bind(udp_server_ipv6, (struct sockaddr *) &udp_addr_ipv6, sizeof(udp_addr_ipv6)) < 0) {
            err_msg = "Unable to bind UDP socket";
            goto finally;
        }

        if (fcntl(udp_server_ipv6, F_SETFL, O_NONBLOCK) < 0) {
            err_msg = "Unable to put UDP socket into non-blocking mode";
            goto finally;
        }

        config_ret_ipv6 = malloc(sizeof(struct config));
        if (config_ret_ipv6 == NULL) {
            err_msg = "Malloc failed";
            goto finally;
        }
        memset(config_ret_ipv6, 0x00, sizeof(*config_ret_ipv6));
        buffer_init(&config_ret_ipv6->read_buffer, N(config_ret_ipv6->raw_buff_a), config_ret_ipv6->raw_buff_a);
        buffer_init(&config_ret_ipv6->write_buffer, N(config_ret_ipv6->raw_buff_b), config_ret_ipv6->raw_buff_b);
    }

    signal(SIGTERM, sigterm_handler);
    signal(SIGINT,  sigterm_handler);

    const struct selector_init conf = {
        .signal = SIGALRM,
        .select_timeout = {
            .tv_sec  = 10,
            .tv_nsec = 0,
        },
    };
    if (0 != selector_init(&conf)) {
        err_msg = "Initializing selector";
        goto finally;
    }

    selector = selector_new(1024);
    if (selector == NULL) {
        err_msg = "Unable to create selector";
        goto finally;
    }

    const struct fd_handler socksv5 = {
        .handle_read = socksv5_passive_accept,
        .handle_write = NULL,
        .handle_close = NULL,
    };
    if (domain == AF_INET || domain == -1) {
        if (selector_fd_set_nio(server) == -1) {
            err_msg = "Getting server socket flags";
            goto finally;
        }

        ss = selector_register(selector, server, &socksv5, OP_READ, NULL);
        if (ss != SELECTOR_SUCCESS) {
            err_msg = "Registering fd";
            goto finally;
        }
    }

    if (domain == AF_INET6 || domain == -1) {
        if (selector_fd_set_nio(server_ipv6) == -1) {
            err_msg = "Getting server socket flags";
            goto finally;
        }

        ss = selector_register(selector, server_ipv6, &socksv5, OP_READ, NULL);
        if (ss != SELECTOR_SUCCESS) {
            err_msg = "Registering fd";
            goto finally;
        }
    }

    const struct fd_handler config = {
            .handle_read = config_read,
            .handle_write = NULL,
            .handle_close = NULL,
    };
    if (domain_udp == AF_INET || domain_udp == -1) {
        if (selector_fd_set_nio(udp_server) == -1) {
            err_msg = "Getting config socket flags";
            goto finally;
        }

        ss = selector_register(selector, udp_server, &config, OP_READ, config_ret);
        if (ss != SELECTOR_SUCCESS) {
            err_msg = "Registering fd";
            goto finally;
        }
    }

    if (domain_udp == AF_INET6 || domain_udp == -1) {
        if (selector_fd_set_nio(udp_server_ipv6) == -1) {
            err_msg = "Getting config socket flags";
            goto finally;
        }

        ss = selector_register(selector, udp_server_ipv6, &config, OP_READ, config_ret_ipv6);
        if (ss != SELECTOR_SUCCESS) {
            err_msg = "Registering fd";
            goto finally;
        }
    }

    for(;!done;) {
        err_msg = NULL;
        ss = selector_select(selector);
        if (ss != SELECTOR_SUCCESS) {
            err_msg = "Serving";
            goto finally;
        }
    }
    if (err_msg == NULL) {
        err_msg = "Closing";
    }

    int ret = 0;
finally:
    for (unsigned int i = 0; i < args->nusers; i++) {
        free(args->users[i].name);
        free(args->users[i].pass);
    }
    free(args);
    if (config_ret != NULL)
        free(config_ret);
    if (config_ret_ipv6 != NULL)
        free(config_ret_ipv6);

    if (ss != SELECTOR_SUCCESS) {
        fprintf(stderr, "%s: %s\n", (err_msg == NULL) ? "": err_msg, ss == SELECTOR_IO ? strerror(errno) : selector_error(ss));
        ret = 2;
    } else if (err_msg) {
        if (errno == 0) {
            fprintf(stderr, "%s\n", err_msg);
        } else {
            perror(err_msg);
        }
        ret = 1;
    }
    if (selector != NULL) {
        selector_destroy(selector);
    }
    selector_close();

    socksv5_pool_destroy();

    if (server >= 0) {
        close(server);
    }
    if (server_ipv6 >= 0) {
        close(server_ipv6);
    }
    if (udp_server >= 0) {
        close(udp_server);
    }
    if (udp_server_ipv6 >= 0) {
        close(udp_server);
    }
    return ret;
}

int findUser(char * username) {
    for (unsigned int i = 0; i < args->nusers; i++) {
        if (!strcmp(args->users[i].name, username))
            return i;
    }
    return -1;
}

int check_password(int idx, char * password) {
    return !strcmp(args->users[idx].pass, password);
}

int add_user(char * username, char * password) {
    if (findUser(username) == -1) {
        if (args->nusers < MAX_USERS) {
            int len = strlen(username) + 1;
            args->users[args->nusers].name = malloc(len);
            if (args->users[args->nusers].name == NULL) {
                return -1;
            }
            strcpy(args->users[args->nusers].name, username);
            args->users[args->nusers].pass = malloc(strlen(password) + 1);
            if (args->users[args->nusers].pass == NULL) {
                return -1;
            }
            strcpy(args->users[args->nusers].pass, password);
            args->users[args->nusers].ulen = len;
            args->nusers++;
            return 0;
        }
        return -1;
    }
    return -2;
}

int modify_username(char * old_uname, char * new_uname) {
    int idx = findUser(old_uname);
    if (idx == -1)
        return 0;
    int len = strlen(new_uname) + 1;
    char * aux = realloc(args->users[idx].name, len);
    if (aux == NULL)
        return -1;
    args->users[idx].name = aux;
    strcpy(args->users[idx].name, new_uname);
    args->users[idx].ulen = len;
    return 1;
}

int modify_password(char * uname, char * new_pwd) {
    int idx = findUser(uname);
    if (idx == -1)
        return 0;
    char * aux = realloc(args->users[idx].pass, strlen(new_pwd) + 1);
    if (aux == NULL)
        return -1;
    args->users[idx].pass = aux;
    strcpy(args->users[idx].pass, new_pwd);
    return 1;
}

bool delete_user(char * uname) {
    int idx = findUser(uname);
    if (idx == -1)
        return false;
    free(args->users[idx].name);
    free(args->users[idx].pass);
    for (unsigned int i = idx; i < args->nusers; i++) {
        args->users[i] = args->users[i+1];
    }
    args->nusers--;
    return true;
}

size_t get_user_pages() {
    return (args->nusers / PAGE_SIZE) + ((args->nusers % PAGE_SIZE != 0) ? 1 : 0);
}

bool get_users(char * buff, int page) {
    size_t start = page * PAGE_SIZE, len = 0;
    unsigned i = 0;

    if (start >= args->nusers) {
        return false;
    }

    i = start;
    while (i < start + PAGE_SIZE && i < args->nusers) {
        strcpy(buff + len, args->users[i].name);
        len += args->users[i++].ulen;
        buff[len - 1] = '\n';
    }

    if (i % PAGE_SIZE) {
        buff[len] = '\n';
    }
    return true;
}

uint8_t verify_token(uint64_t token) {
    return token == config_token;
}
