/*
 * ips.c: Internet Proxy Server
 *
 */

#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 LISTEN_BACKLOG 32
#define REQUEST_BUFSZ  48

static char req_buffer[REQUEST_BUFSZ] = {0};
static bool ips_running = true;

static void print_usage(const char *app_name)
{
	printf("Usage:\n");
	printf("    %s <forwarding port> <tunnel port>\n", app_name);
	printf("\n");
	printf("  forwarding port: Port in which clients from external networks\n");
	printf("                   will be able to reach a service inside the\n");
	printf("                   private subnet.\n");
	printf("      tunnel port: Port in which the IPS listen to. This is used\n");
	printf("                   by the SRC to maintain a control connection.\n");
}

static void parse_args(int argc, const char **argv, int *ports)
{
	int i;
	char *end;
	uid_t uid = getuid();

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

	for (i = 0; i < 2; i++) {
		ports[i] = strtol(argv[i + 1], &end, 10);

		if (*end != '\0' || errno == ERANGE) {
			print_error("Invalid port value.\n");
			goto _parse_abort;
		}

		if (ports[i] < 1 || ports[i] > 65535) {
			print_error("Port is out of range (1-65535).\n");
			goto _parse_abort;
		}

		if (ports[i] < 1024 && uid != 0) {
			print_error("Binding ports below 1024 require root privileges.\n");
			exit(-3);
		}
	}

	/* No error found */
	return;

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

int new_control_connection(int relay_sock)
{
	int ctl_sock;

	/* Control socket address info. */
	struct sockaddr_in ctl_addr;
	socklen_t ctl_len = sizeof(struct sockaddr_in);

	/* Buffer for storing addresses as a C string. */
	char addr_buf[INET_ADDRSTRLEN];

	/* Accept the new connection */
	ctl_sock = accept(relay_sock, (struct sockaddr *)&ctl_addr, &ctl_len);

	if (ctl_sock > 0) {
		/* Print relay client IP */
		inet_ntop(AF_INET, &ctl_addr.sin_addr, addr_buf, ctl_len);
		print_info("Connection from: %s <relay client>\n", addr_buf);
	}

	return ctl_sock;
}

void accept_new_client(int client_sock, int ctl_sock, struct node **head_ref)
{
	int new_sock;
	struct sockaddr_in addr;
	socklen_t len = sizeof addr;
	struct conn_context new_cc = empty_conn_context();

	/* Buffer for storing addresses as a C string */
	char addr_buf[INET_ADDRSTRLEN];

	if (!head_ref)
		return;

	/* Accept the new connection */
	new_sock = accept(client_sock, (struct sockaddr *)&addr, &len);

	/* Print client IP */
	inet_ntop(AF_INET, &addr.sin_addr, addr_buf, len);
	print_info("New connection from: %s <client: %d>\n", addr_buf, new_sock);

	/* Allocate a new node */
	new_cc = empty_conn_context();
	new_cc.status = S_UNPAIRED;
	new_cc.len[0] = len;
	new_cc.addr[0] = addr;
	new_cc.sockets[0] = new_sock;
	new_cc.timestamp = get_timestamp_ms();

	/* Copy and insert into the list */
	node_insert(head_ref, &new_cc, sizeof new_cc);

	/* Transmit a request to establish a relay tunnel */
	snprintf(req_buffer, REQUEST_BUFSZ, "REQ_RELAY: %d\r\n", new_sock);
	send(ctl_sock, req_buffer, strnlen(req_buffer, REQUEST_BUFSZ), MSG_NOSIGNAL);
}

void accept_new_relay(int relay_sock, struct node *head)
{
	int new_sock;
	struct sockaddr_in addr;
	socklen_t len = sizeof addr;
	struct conn_context *cc;

	/* Buffer for storing addresses as a C string */
	char addr_buf[INET_ADDRSTRLEN];

	/* Accept the new connection */
	new_sock = accept(relay_sock, (struct sockaddr *)&addr, &len);

	/* Print relay client IP */
	inet_ntop(AF_INET, &addr.sin_addr, addr_buf, len);
	print_info("New connection from: %s <relay: %d>\n", addr_buf, new_sock);

	/* Check for available slots (clients establishing a new connection) */
	if (!(cc = find_unpaired_node(head))) {
		print_warn("Unsolicited connection, closing...\n");
		close(new_sock);
		return;
	}

	/* Pair the relay tunnel with the available client */
	cc->status = S_WAITING;
	cc->len[1] = len;
	cc->addr[1] = addr;
	cc->sockets[1] = new_sock;
}

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

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

	/* Sockets */
	int ctl_sock;
	int relay_sock;
	int client_sock;

	/* Listening ports */
	int listen_ports[2];

	/* Parse arguments. */
	parse_args(argc, argv, listen_ports);

	/* Listen for connections */
	relay_sock = tcp_server_socket(listen_ports[1], LISTEN_BACKLOG);
	client_sock = tcp_server_socket(listen_ports[0], LISTEN_BACKLOG);

	if (relay_sock == -1 || client_sock == -1) {
		fprintf(stderr, "Either one or both ports are unavailable\n");
		exit(-5);
	}

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

	/* Wait for relay client */
	printf("\tInternet Proxy Server\n\n");
	printf("Press Ctrl+C to exit...\n\n");

	print_info("Waiting for relay client...\n");
	ctl_sock = new_control_connection(relay_sock);

	while (ips_running) {
		int retval, 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("Relay client disconnected.\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 disconnected, attempt to reconnect */
		if (ctl_sock < 0) {
			/* Drop any pending connection */
			close(relay_sock);
			close(client_sock);

			/* Rebind listening sockets */
			relay_sock = tcp_server_socket(listen_ports[1], LISTEN_BACKLOG);
			client_sock = tcp_server_socket(listen_ports[0], LISTEN_BACKLOG);
			if (relay_sock == -1 || client_sock == -1) {
				fprintf(stderr, "Unable to rebind listening sockets\n");
				exit(-5);
			}

			/* Wait for a new connection */
			print_info("Waiting for relay client...\n");
			ctl_sock = new_control_connection(relay_sock);
			continue;
		}

		/* Add sockets to the master (reference) descriptor set. */
		FD_ZERO(&working_fds);
		FD_SET(relay_sock, &working_fds);
		FD_SET(client_sock, &working_fds);
		maxfd = 1 + max(relay_sock, client_sock);

		/* Check if there is a new connection. */
		retval = select(maxfd, &working_fds, NULL, NULL, &sel_timeout);
		if (retval <= 0) {
			/* Select has got an error or timed out. */
			continue;
		}

		/* Self explanatory, isn't it? */
		if (FD_ISSET(client_sock, &working_fds))
			accept_new_client(client_sock, ctl_sock, &head);

		if (FD_ISSET(relay_sock, &working_fds))
			accept_new_relay(relay_sock, head);
	}

	close(ctl_sock);
	close(client_sock);
	close(relay_sock);

	return EXIT_SUCCESS;
}
