/*
 * src.c: Subnet relay client
 *
 */

#include "common.h"
#include "linked_list.h"
#include "net_utils.h"

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <errno.h>
#include <pthread.h>
#include <signal.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/time.h>

#define REQUEST_BUFSZ  4096
#define MAX_ATTEMPTS   5

static char req_buffer[REQUEST_BUFSZ] = {0};
static bool src_running = true;
static int reconn_attempts = 0;

struct client_config {
	const char *proxy_server;
	int proxy_port;

	const char *target_server;
	int target_port;
};

static void print_usage(const char *app_name)
{
	printf("Usage:\n");
	printf("    %s <proxy server> <proxy port> <target server> <target port>\n", app_name);
	printf("\n");
	printf("   proxy server: Host or IP address of the machine running the IPS.\n");
	printf("     proxy port: Port in which the IPS listens to.\n");
	printf("  target server: Host or IP address of the service we want to forward.\n");
	printf("    target port: Port in which the service listens to.\n");
}

int connect_to_proxy(const char *hostname, uint16_t port)
{
	int new_sock;

	print_info("Connecting to proxy server... ");
	new_sock = tcp_client_socket(hostname, port);
	if (new_sock < 0) {
		printf("[" COLOR_RED "failed" COLOR_OFF "]\n");
		sleep(2); /* Sleep for two seconds before trying again */
	} else {
		printf("[" COLOR_GREEN "success" COLOR_OFF "]\n");
	}

	return new_sock;
}

void parse_new_request(int ctl_sock, struct client_config *cfg,
	struct node **head_ref)
{
	int service_sock, client_sock;
	struct conn_context new_cc = empty_conn_context();
	ssize_t nbr = recv(ctl_sock, req_buffer, REQUEST_BUFSZ, 0);

	if (!head_ref || !cfg)
		return;

	if (nbr <= 0 || errno == EAGAIN)
		return; /* Resource is unavailable, try again. */

	service_sock = tcp_client_socket(cfg->target_server, cfg->target_port);
	client_sock = tcp_client_socket(cfg->proxy_server, cfg->proxy_port);

	if (service_sock && client_sock) {
		print_info("New relay request <client: %d, relay: %d>\n",
			client_sock, service_sock);
		new_cc.status = S_WAITING;
		new_cc.sockets[0] = service_sock;
		new_cc.sockets[1] = client_sock;
		new_cc.timestamp = get_timestamp_ms();
		node_insert(head_ref, &new_cc, sizeof new_cc);
	}
}

static inline int sanitize_port_str(const char *port_s)
{
	char *end;

	int tmp = strtol(port_s, &end, 10);
	if (*end != '\0' || errno == ERANGE) {
		print_error("Invalid port value <%s>.\n", port_s);
		return -EINVAL;
	} else if (tmp < 1 || tmp > 65535) {
		print_error("Port is out of range (1-65535) <%s>.\n", port_s);
		return -ERANGE;
	}

	return tmp;
}

static void parse_args(int argc, const char **argv, struct client_config *cfg)
{
	int dummy_sock;
	struct sockaddr_in addr;

	/* Check the number of arguments and perform some input sanitization */
	if (argc != 5) {
		print_error("Insufficient arguments.\n");
		goto _parse_abort;
	}

	/* Parse arguments */
	cfg->proxy_server = argv[1];
	cfg->proxy_port = sanitize_port_str(argv[2]);
	cfg->target_server = argv[3];
	cfg->target_port = sanitize_port_str(argv[4]);

	if (cfg->proxy_port < 0 || cfg->target_port < 0)
		goto _parse_abort;

	/* Test proxy server name resolution */
	if (resolve_host(cfg->proxy_server, &addr) == NULL) {
		print_error("Cannot resolve proxy server: %s\n", cfg->target_server);
		exit(-4);
	}

	/* Check if the target service is accepting connections */
	dummy_sock = tcp_client_socket(cfg->target_server, cfg->target_port);
	if (dummy_sock < 0) {
		print_error("Service is unavaiable <%s:%d>\n",
			cfg->target_server, cfg->target_port);
		exit(-4);
	}

	/* No error found */
	close(dummy_sock);
	return;

_parse_abort:
	print_usage(argv[0]);
	exit(-1);
}

static void handle_signal(int signum)
{
	switch (signum) {
	case SIGINT: /* Falls through */
	case SIGTERM:
		src_running = false;
		printf("\n");
		print_warn("Closing program...\n");
		break;
	}
}

int main(int argc, const char *argv[])
{
	struct node *head = NULL;
	struct client_config cfg;
	struct sigaction sa;

	int ctl_sock = -1; /* Arbitrarily chosen (must be negative) */

	/* Register signal handler */
	sa.sa_handler = &handle_signal;
	sigemptyset(&sa.sa_mask);
	sa.sa_flags = 0;
	sigaction(SIGINT, &sa, NULL);
	sigaction(SIGTERM, &sa, NULL);

	parse_args(argc, argv, &cfg);
	printf("\tSource Relay Client\n\n");
	printf("Press Ctrl+C to exit...\n\n");

	while (src_running) {
		int retval;
		int maxfd;
		struct timeval sel_timeout = {.tv_sec = 0, .tv_usec = 1000};
		fd_set working_fds;

		/* Check if the control socket is still alive. */
		if (ctl_sock > 0 && test_connection(ctl_sock)) {
			print_warn("Disconnected from proxy server.\n");

			/* Schedule current connections to be dropped */
			mark_all_connections(head, S_TERMINATED);

			/* Close the socket descriptor */
			close(ctl_sock);
			ctl_sock = -1;
		}

		/* Connection scheduler */
		run_connection_scheduler(&head);

		if (ctl_sock < 0) {
			if (++reconn_attempts > MAX_ATTEMPTS) {
				print_error("Maximum attempts reached.\n");
				exit(-1);
			}

			ctl_sock = connect_to_proxy(cfg.proxy_server, cfg.proxy_port);
			continue;
		}

		/* Clear reconnection attempts counter */
		reconn_attempts = 0;

		/* Check if there is a new connection request */
		FD_ZERO(&working_fds);
		FD_SET(ctl_sock, &working_fds);
		maxfd = 1 + ctl_sock;
		retval = select(maxfd, &working_fds, NULL, NULL, &sel_timeout);
		if (retval <= 0)
			continue; /* Either error or timed out. */

		/* Parse a new request and create a relay between the proxy
		 * and the server */
		if (FD_ISSET(ctl_sock, &working_fds))
			parse_new_request(ctl_sock, &cfg, &head);
	}

	close(ctl_sock);

	return EXIT_SUCCESS;
}
