#define _GNU_SOURCE
#define _BSD_SOURCE
#define _POSIX_C_SOURCE 200809L
#include <errno.h>
#include <netdb.h>
#include <poll.h>
#include <pwd.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

struct pollfd *listens = NULL;
size_t num_listens = 0;

void log_error(const char * restrict format, ...) {
	va_list args;

	va_start(args, format);

	vfprintf(stderr, format, args);

	va_end(args);
}

void log_perror(const char *string) {
	perror(string);
}

void setup_listen(void) {
	struct addrinfo hints;
	struct addrinfo *getaddrinfo_result;

	// AF_UNSPEC: either IPv4 or IPv6
	// SOCK_STREAM: TCP
	// AI_PASSIVE: fill out my IP for me
	// AI_ADDRCONFIG: only return addresses I have a configured interface for
	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_flags = AI_PASSIVE | AI_ADDRCONFIG;

	int status = getaddrinfo(NULL, "79", &hints, &getaddrinfo_result);

	if(status != 0) {
		log_error("getaddrinfo failed: %s\n", gai_strerror(status));
		exit(1);
	}

	for(struct addrinfo *res = getaddrinfo_result; res != NULL; res = res->ai_next) {
		int yes = 1;

		// Add corresponding interface to table of sockets
		// Create socket
		int sock = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
		if(sock == -1) {
			log_perror("socket");
			exit(1);
		}

		// Disable the IPv4 over IPv6, as that results in IPv4 and IPv6 sockets conflicting and one of them not being able to be set up
		if(res->ai_family == AF_INET6) {
			if(setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY, &yes, sizeof(yes)) == -1) {
				log_perror("setsockopt");
				exit(1);
			}
		}

		// Set reuseaddr
		if(setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)) == -1) {
			log_perror("setsockopt");
			exit(1);
		}

		// Bind onto given address
		if(bind(sock, res->ai_addr, res->ai_addrlen) == -1) {
			log_perror("bind");
			exit(1);
		}

		// Listen for incoming connections
		if(listen(sock, 1) == -1) {
			log_perror("listen");
			exit(1);
		}

		// Grow the table of sockets
		size_t index = num_listens++;
		if(SIZE_MAX / sizeof(struct pollfd) < num_listens){
			log_error("Too many interfaces to listen on\n");
			exit(1);
		}

		listens = realloc(listens, num_listens * sizeof(struct pollfd));

		if(listens == NULL) {
			log_perror("realloc");
			exit(1);
		}

		// Add socket to the table
		struct pollfd new_pollobj = {.fd = sock, .events = POLLIN};
		listens[index] = new_pollobj;
	}

	freeaddrinfo(getaddrinfo_result);
}

void drop_privileges(void) {
	uid_t uid = getuid();
	gid_t gid = getgid();

	if(setresgid(gid, gid, gid) != 0) {
		log_perror("setresgid");
		exit(1);
	}
	if(setresuid(uid, uid, uid) != 0) {
		log_perror("setresuid");
		exit(1);
	}
}

ssize_t writeall(int fd, const char *buf, size_t amount) {
	size_t written = 0;
	while(written < amount) {
		ssize_t r = write(fd, buf + written, amount - written);
		if(r < 0) {
			return r;
		}

		written += r;
	}

	return written;
}

ssize_t writeall_str(int fd, const char *s) {
	return writeall(fd, s, strlen(s));
}

bool execute_finger_response(int sock, const char *finger_response_file) {
	int pipes[2];
	if(pipe(pipes) != 0) {
		log_perror("pipe");
		return false;
	}

	pid_t child = fork();
	if(child < 0) {
		log_perror("fork");
		return false;
	}

	if(child == 0) {
		// Child

		// Close unneeded fds
		for(size_t i = 0; i < num_listens; i++) {
			close(listens[num_listens].fd);
		}
		close(sock);
		close(pipes[0]);

		// Set up stdout
		if(dup2(pipes[1], STDOUT_FILENO) < 0) {
			log_perror("dup2");
			_exit(1);
		}

		// execv's types require this cast
		char *argv[] = {(char*) finger_response_file, NULL};
		execv(finger_response_file, argv);

		log_perror("execv");
		_exit(1);
	} else {
		// Parent

		// Close unneeded fds
		close(pipes[1]);

		struct timespec start_time;
		// Doesn't seem this could fail in our configuration
		clock_gettime(CLOCK_MONOTONIC, &start_time);

		char iobuf[1024];

		struct pollfd pipe_poll = {.fd = pipes[0], .events = POLLIN};
		for(;;) {
			// Check the child is still running
			int status;
			waitpid(child, &status, WNOHANG);
			if(WIFEXITED(status) || WIFSIGNALED(status)) {
				// Child quit
				if(WIFSIGNALED(status)) {
					log_error("Child died of signal %i\n", WTERMSIG(status));
				} else if(WIFEXITED(status) && WEXITSTATUS(status) != 0) {
					log_error("Child exited with nonzero return code %i\n", WEXITSTATUS(status));
				}
				break;
			}

			struct timespec current_time;
			// See previous usage of this
			clock_gettime(CLOCK_MONOTONIC, &current_time);

			time_t secs_passed = current_time.tv_sec - start_time.tv_sec;
			long ns_passed = current_time.tv_nsec - start_time.tv_nsec;
			int ms_passed = secs_passed * 1000 + ns_passed / 1000;
			// Timeout of 5s
			int ms_remaining = 5000 - ms_passed;

			if(ms_remaining <= 0) {
				log_error("Child ran too long\n");
				if(kill(child, SIGKILL)) {
					log_perror("kill");
				}

				// Wait to get rid of the zombie
				waitpid(child, NULL, 0);

				if(writeall_str(sock, "Timeout\r\n") < 0) {
					log_perror("writeall_str");
					close(pipes[0]);
					return false;
				}

				break;
			}

			int amount_ready = poll(&pipe_poll, 1, ms_remaining);

			for(int i = 0; i < amount_ready; i++) {
				if(pipe_poll.revents & POLLIN) {
					ssize_t amount_read = read(pipes[0], iobuf, sizeof(iobuf));
					if(amount_read < 0) {
						log_perror("read");
						break;
					}

					char *newline_location = memchr(iobuf, '\n', amount_read);
					if(newline_location == NULL) {
						// No newline, write entire buffer
						if(writeall(sock, iobuf, amount_read) < 0) {
							log_perror("writeall");
							break;
						}
					} else {
						// Newline found. Write stuff before it and \r\n
						if(writeall(sock, iobuf, newline_location - iobuf) < 0 || writeall_str(sock, "\r\n") < 0) {
							log_perror("writeall / writeall_str");
							break;
						}

						// Write the part after it
						char *part_after_nl = newline_location + 1;
						size_t part_after_nl_length = amount_read - (part_after_nl - iobuf);
						if(writeall(sock, part_after_nl, part_after_nl_length) < 0) {
							log_perror("writeall");
							break;
						}
					}
				}
			}
		}

		close(pipes[0]);
	}

	return true;
}

bool read_finger_response(int sock, const char *finger_response_file) {
	FILE *f = fopen(finger_response_file, "r");
	if(f == NULL) {
		log_perror("fopen");
		return false;
	}

	size_t line_size = 0;
	char *line = NULL;
	ssize_t line_len;
	while((line_len = getline(&line, &line_size, f)) > 0) {
		// Strip the newline
		if(line[line_len - 1] == '\n') {
			line[line_len - 1] = '\0';
			line_len--;
		}

		// Send the line along with the correct newline
		if(writeall(sock, line, line_len) < 0 || writeall_str(sock, "\r\n") < 0) {
			log_perror("writeall / writeall_str");
			free(line);
			return false;
		}
	}

	if(!feof(f)) {
		log_perror("getline");
		free(line);
		return false;
	}

	free(line);
	fclose(f);
	return true;
}

bool handle_finger_response(int sock, const char *finger_response_file) {
	if(access(finger_response_file, X_OK) == 0) {
		return execute_finger_response(sock, finger_response_file);
	} else {
		return read_finger_response(sock, finger_response_file);
	}
}

void handle_userlist(int sock, bool force_multiline) {
	// Go over users' home dirs and see if .finger-response is present
	bool first_line = true;
	for(;;) {
		errno = 0;
		struct passwd *passwd_entry = getpwent();
		if(passwd_entry == NULL) {
			if(errno == 0) {
				// End of file
				endpwent();
				break;
			} else {
				log_perror("getpwent");
				return;
			}
		}

		char *finger_response_file;
		if(asprintf(&finger_response_file, "%s/.finger-response", passwd_entry->pw_dir) < 0) {
			log_perror("asprintf");
			endpwent();
			return;
		}

		if(access(finger_response_file, F_OK) == 0) {
			if(force_multiline && !first_line) {
				if(writeall_str(sock, "\r\n") < 0) {
					log_perror("writeall_str (separating newline)");
					endpwent();
					return;
				}
			}

			if(writeall_str(sock, passwd_entry->pw_name) < 0) {
				log_perror("writeall_str (username)");
				endpwent();
				return;
			}

			const char *suffix = "\r\n";
			if(force_multiline) {
				suffix = ":\r\n";
			}

			if(writeall_str(sock, suffix) < 0) {
				log_perror("writeall_str (suffix)");
				endpwent();
				return;
			}

			if(force_multiline) {
				if(handle_finger_response(sock, finger_response_file) == false) {
					endpwent();
					return;
				}
			}

			first_line = false;
		}
	}
}

void handle_query(int sock, const char *username) {
	errno = 0;
	struct passwd *passwd_entry = getpwnam(username);
	if(passwd_entry == NULL) {
		if(errno == 0) {
			if(writeall_str(sock, "Can't finger\r\n") < 0) {
				log_perror("writeall_str");
				return;
			}
		} else {
			log_perror("getpwent");
		}
		return;
	}

	char *finger_response_file;
	if(asprintf(&finger_response_file, "%s/.finger-response", passwd_entry->pw_dir) < 0) {
		log_perror("asprintf");
		return;
	}

	if(access(finger_response_file, F_OK) != 0) {
		if(writeall_str(sock, "Can't finger\r\n") < 0) {
			log_perror("writeall_str (not found error)");
			return;
		}
	} else {
		handle_finger_response(sock, finger_response_file);
	}
}

void handle_connection(int sock) {
	size_t request_size = 0;
	char *request = NULL;

	for(;;) {
		ssize_t amount_read;
		char iobuf[1024];

		amount_read = read(sock, &iobuf, sizeof(iobuf));

		if(amount_read == 0) {
			log_error("Client hung up\n");
			shutdown(sock, SHUT_RDWR);
			close(sock);
			return;
		} else if(amount_read < 0) {
			log_perror("read");
			shutdown(sock, SHUT_RDWR);
			close(sock);
			return;
		}

		size_t index = request_size;
		request_size += amount_read;
		if((request = realloc(request, request_size)) == NULL) {
			log_perror("realloc");
			exit(1);
		}
		memmove(request + index, iobuf, amount_read);

		if(request_size >= 2 && request[request_size-2] == '\r' && request[request_size-1] == '\n') {
			// Request read
			break;
		}
	}

	char *request_buffer = request;
	request[request_size-2] = '\0';

	// See if the multi-line format was requested
	bool force_multiline = false;
	if(request[0] == '/' && request[1] == 'W') {
		if(request[2] == '\0') {
			// Just "/W"
			force_multiline = true;
			request = &request[2];
		} else if(request[2] == ' ') {
			// "/W " as prefix
			force_multiline = true;
			request = &request[3];
		}
	}

	if(!strcmp(request, "")) {
		handle_userlist(sock, force_multiline);
	} else {
		handle_query(sock, request);
	}

	free(request_buffer);

	shutdown(sock, SHUT_RDWR);
	close(sock);
}

int main(void) {
	setup_listen();
	drop_privileges();

	for(;;) {
		int amount_ready = poll(listens, num_listens, -1);
		if(amount_ready < 0) {
			log_perror("poll");
			exit(1);
		}

		for(size_t i = 0; i < num_listens && amount_ready > 0; i++) {
			if(listens[i].revents & POLLIN) {
				struct sockaddr_storage client_addr;
				socklen_t addr_size = sizeof(client_addr);

				int sock = accept(listens[i].fd, (struct sockaddr *)&client_addr, &addr_size);

				handle_connection(sock);

				amount_ready--;
			}
		}
	}
}