#define _POSIX_C_SOURCE 200809L  // mkdtemp
#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <poll.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/wait.h>
#include "gen/libintercept.so.h"
#include "lib/libintercept.h"

#ifndef __linux__
#error This program currently only works on Linux
#endif


static void* mallocerr(size_t num) {
	void *ptr = malloc(num);
	if (ptr == NULL) {
		fprintf(stderr, "Cannot allocate memory\n");
		exit(1);
	}
	return ptr;
}

// Returns whether successful
static bool write_library(const char *path) {
	FILE *f = fopen(path, "w");
	if (!f) {
		fprintf(stderr, "Could not place library in temporary directory\n");
		return false;
	}

	fwrite(libintercept_so, 1, libintercept_so_len, f);
	if (fclose(f) != 0) {
		fprintf(stderr, "Could not write to temporary directory (%s)\n", strerror(errno));
		return false;
	}

	return true;
}

// Returns socket, or -1 on error
static int create_socket(const char *path) {
	int sock = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
	if (sock < 0) {
		perror("Cannot create unix socket: socket");
		return -1;
	}

	struct sockaddr_un addr;
	memset(&addr, 0, sizeof addr);
	addr.sun_family = AF_UNIX;
	strncpy(addr.sun_path, path, sizeof addr.sun_path - 1);

	if (bind(sock, (const struct sockaddr*)&addr, sizeof addr) < 0) {
		perror("Cannot create unix socket: bind");
		close(sock);
		return -1;
	}

	if (listen(sock, 10) < 0) {
		perror("listen");
		close(sock);
		unlink(path);
		return -1;
	}

	return sock;
}

struct paths {
	char *tempdir;
	char *libpath;
	char *socketpath;
	int socket;
};

// returns {NULL} on failure
static struct paths create_tempfiles(void) {
	const char *template = "/tmp/exec-intercept.XXXXXX";
	const char *libsuffix = "/libintercept.so";
	const char *socketsuffix = "/exec-intercept.socket";

	char *tempdir = mallocerr(strlen(template) + 1);
	memcpy(tempdir, template, strlen(template) + 1);

	if (mkdtemp(tempdir) == NULL) {
		perror("Cannot create temporary directory");
		exit(1);
	}

	char *libpath = mallocerr(strlen(template) + strlen(libsuffix) + 1);
	memcpy(libpath, tempdir, strlen(template));
	memcpy(libpath + strlen(template), libsuffix, strlen(libsuffix) + 1);

	if (!write_library(libpath)) {
		free(libpath);
		rmdir(tempdir);
		free(tempdir);
		return (struct paths){NULL, NULL, NULL, -1};
	}

	char *socketpath = mallocerr(strlen(template) + strlen(socketsuffix) + 1);
	memcpy(socketpath, tempdir, strlen(template));
	memcpy(socketpath + strlen(template), socketsuffix, strlen(socketsuffix) + 1);

	int socket = create_socket(socketpath);
	if (socket < 0) {
		free(socketpath);
		unlink(libpath);
		free(libpath);
		rmdir(tempdir);
		free(tempdir);
		return (struct paths){NULL, NULL, NULL, -1};
	}

	return (struct paths){tempdir, libpath, socketpath, socket};
}

static void cleanup_tempfiles(struct paths info) {
	close(info.socket);
	unlink(info.socketpath);
	free(info.socketpath);
	unlink(info.libpath);
	free(info.libpath);
	rmdir(info.tempdir);
	free(info.tempdir);
}

// Returns PID of child, or -1 on error
static pid_t start_child(char **argv) {
	pid_t pid = fork();
	if (pid < 0) {
		perror("fork");
		return -1;
	}

	if (pid == 0) {
		// Use execvp, which finds the executable in PATH
		execvp(argv[0], argv);
		perror("execvp");
		exit(1);
	}

	return pid;
}

struct message_buffer {
	size_t cap, len;
	char *buffer;
};

// If .buffer is NULL, memory allocation failed
static struct message_buffer message_buffer_alloc(void) {
	struct message_buffer mb;
	mb.cap = 1024;
	mb.len = 0;
	mb.buffer = malloc(mb.cap);
	if (mb.buffer == NULL) {
		fprintf(stderr, "Cannot allocate memory\n");
		return mb;
	}
	return mb;
}

// Returns length of (string + '\0'), or 0 if '\0' is not found
static size_t consume_string(const char *buffer, size_t available) {
	for (size_t i = 0; i < available; i++) {
		if (buffer[i] == '\0') return i + 1;
	}
	return 0;
}

static bool shell_safe(char c) {
	return ('+' <= c && c <= ':') ||
			c == '=' ||
			('@' <= c && c <= '[') ||
			(']' <= c && c <= '_') ||
			('a' <= c && c <= 'z');
}

static void write_shell_escaped(FILE *logfile, const char *str, size_t length) {
	if (length == 0) {
		fprintf(logfile, "''");
		return;
	}

	bool safe = true;
	for (size_t i = 0; i < length; i++) {
		if (!shell_safe(str[i])) {
			safe = false;
			break;
		}
	}

	if (safe) {
		fwrite(str, 1, length, logfile);
		return;
	}

	int quotemode = 0;  // 0=none, 1='', 2="", 3=$''
	for (size_t i = 0; i < length; i++) {
		if (str[i] < ' ' || str[i] > '~') {
			if (quotemode == 0) fprintf(logfile, "$'");
			else if (quotemode == 1) fprintf(logfile, "'$'");
			else if (quotemode == 2) fprintf(logfile, "\"$'");
			quotemode = 3;
			fprintf(logfile, "\\x%c%c",
					"0123456789abcdef"[str[i] / 16], "0123456789abcdef"[str[i] % 16]);
		} else if (str[i] == '\'') {
			if (quotemode == 0) fputc('"', logfile);
			else if (quotemode == 1 || quotemode == 3) fprintf(logfile, "'\"");
			quotemode = 2;
			fputc('\'', logfile);
		} else {
			if (quotemode == 0) fputc('\'', logfile);
			else if (quotemode == 2) fprintf(logfile, "\"'");
			else if (quotemode == 3) fprintf(logfile, "''");
			quotemode = 1;
			fputc(str[i], logfile);
		}
	}
	if (quotemode == 1 || quotemode == 3) fputc('\'', logfile);
	else if (quotemode == 2) fputc('"', logfile);
}

// Returns amount consumed, which may be 0
static size_t parse_write_entry(const char *buffer, size_t available, FILE *logfile) {
	if (available < 8) return 0;
	const size_t nargs = *(const size_t*)buffer;

	size_t cursor = 8;
	for (size_t i = 0; i < nargs + 1; i++) {
		const size_t len = consume_string(buffer + cursor, available - cursor);
		if (len == 0) return 0;
		cursor += len;
	}

	// All null bytes were found, so let's log the full message
	cursor = 8;
	for (size_t i = 0; i < nargs + 1; i++) {
		const size_t len = consume_string(buffer + cursor, available - cursor);
		if (i != 1) {  // skip argv[0]
			if (i > 0) fputc(' ', logfile);
			write_shell_escaped(logfile, buffer + cursor, len - 1);
		}
		cursor += len;
	}

	fputc('\n', logfile);

	return cursor;
}

// Returns whether successful
static bool message_buffer_add_data(
		struct message_buffer *mb, const char *data, size_t length,
		FILE *logfile) {
	if (mb->len + length > mb->cap) {
		do mb->cap *= 2;
		while (mb->len + length > mb->cap);
		mb->buffer = realloc(mb->buffer, mb->cap);
		if (mb->buffer == NULL) return false;
	}

	memcpy(mb->buffer + mb->len, data, length);
	mb->len += length;

	size_t nconsumed = parse_write_entry(mb->buffer, mb->len, logfile);
	if (nconsumed > 0) {
		memmove(mb->buffer, mb->buffer + nconsumed, mb->len - nconsumed);
		mb->len -= nconsumed;
	}
	return true;
}

static void message_buffer_free(struct message_buffer mb) {
	free(mb.buffer);
}

static void message_loop(int listensock, pid_t child_pid, FILE *logfile) {
	// pfds[0] is always listensock
	size_t conn_cap = 4, conn_len = 1;

	struct pollfd *pfds = mallocerr(conn_cap * sizeof(struct pollfd));
	pfds[0].fd = listensock;
	pfds[0].events = POLLIN;

	// mbufs[0] is always unused
	struct message_buffer *mbufs = mallocerr(conn_cap * sizeof(struct message_buffer));

	while (true) {
		int ret = poll(pfds, conn_len, -1);
		if (ret < 0 && errno != EINTR) {
			perror("poll");
			goto cleanup;
		}

		if (pfds[0].revents & POLLIN) {
			int sock = accept(listensock, NULL, NULL);
			if (sock >= 0) {
				if (conn_len == conn_cap) {
					conn_cap *= 2;
					pfds = realloc(pfds, conn_cap * sizeof(struct pollfd));
					mbufs = realloc(mbufs, conn_cap * sizeof(struct message_buffer));
					if (!pfds || !mbufs) {
						fprintf(stderr, "Cannot allocate memory!\n");
						exit(1);
					}
				}
				pfds[conn_len].fd = sock;
				pfds[conn_len].events = POLLIN;
				mbufs[conn_len] = message_buffer_alloc();
				conn_len++;
			}
		}

		for (size_t i = 1; i < conn_len; i++) {
			if (!(pfds[i].revents & POLLIN)) continue;

			char buf[1024];
			ssize_t nr = read(pfds[i].fd, buf, sizeof buf);
			if (nr < 0 && errno != EINTR && errno != ECONNRESET) {
				perror("read");
				close(pfds[i].fd);
				message_buffer_free(mbufs[i]);
				if (i < conn_len - 1) {
					pfds[i] = pfds[conn_len - 1];
					mbufs[i] = mbufs[conn_len - 1];
				}
				conn_len--;
			}

			if (nr > 0) {
				message_buffer_add_data(&mbufs[i], buf, (size_t)nr, logfile);
				continue;  // see if there's more data before wait()'ing on the child
			}
		}

		int status;
		pid_t waitret = waitpid(child_pid, &status, WNOHANG);
		if (waitret < 0) {
			perror("waitpid");
			goto cleanup;
		}

		if (waitret > 0 && WIFEXITED(status)) {
			goto cleanup;
		}
	}

cleanup:
	for (size_t i = 1; i < conn_len; i++) {
		close(pfds[i].fd);
		message_buffer_free(mbufs[i]);
	}
}

static void signal_handler_nop(int sig) {
	(void)sig;
}

int main(int argc, char **argv) {
	if (argc < 4) {
		fprintf(stderr,
				"Usage: %s -o <log.txt> <command...>\n"
				"Captures (and passes through) all execve calls in the command,\n"
				"and logs them to the specified file. The commands in the log\n"
				"file are shell-escaped in bash syntax.\n",
				argv[0]);
		return 1;
	}

	const char *logfname = argv[2];

	if (getenv("LD_LIBRARY_PATH") != NULL || getenv("LD_PRELOAD") != NULL) {
		fprintf(stderr,
				"This program manipulates the dynamic loader environment\n"
				"variables LD_LIBRARY_PATH and LD_PRELOAD; having them set\n"
				"while running this program is insecure. Please unset them\n"
				"first.\n");
		return 1;
	}

	if (signal(SIGCHLD, signal_handler_nop) == SIG_ERR) {
		perror("signal");
		return 1;
	}

	struct paths paths = create_tempfiles();
	if (paths.tempdir == NULL) return 1;
	FILE *logfile = NULL;

	// This new environment only applies to new child processes
	if (setenv("LD_PRELOAD", paths.libpath, 1) < 0 ||
			setenv(COMM_SOCKET_ENVVAR, paths.socketpath, 1) < 0) {
		perror("setenv");
		goto cleanup;
	}

	logfile = fopen(logfname, "w");
	if (logfile == NULL) {
		fprintf(stderr, "Cannot open log file <%s>\n", logfname);
		goto cleanup;
	}

	pid_t child_pid = start_child(argv + 3);
	if (child_pid == -1) {
		goto cleanup;
	}

	message_loop(paths.socket, child_pid, logfile);

cleanup:
	if (logfile != NULL) fclose(logfile);
	cleanup_tempfiles(paths);
}