#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

#define _XOPEN_SOURCE 600
#define _BSD_SOURCE

#include "session.h"

#include "xwrap.h"
#include "socktool.h"
#include "connection.h"
#include "server.h"
#include "shell.h"
#include "basic-array.h"

#include <fcntl.h>
#include <errno.h>
#include <grp.h>
#include <pwd.h>
#include <signal.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

int ras_session_init (RasSession* session, RasServer* server, int fd, int id) {
	/* chain up to parent constructor */
	int rval = ras_conn_init (
			RAS_CONN (session),
			RAS_CONN (server)->perm,
			RAS_CONN (server)->domain,
			"ras-session",
			RAS_CONN (server)->log_file,
			id);

	if (rval < 0) {
		return rval;
	}

	RAS_CONN (session)->fd = fd;
	RAS_CONN (session)->fd_is_set = true;

	ras_socktool_buffer_clear (&session->buffer, true);
	session->shell_pid = 0;
	session->shell_stdio = -1;
	session->shell_use_pty = true;

	ras_session_log (session, "session started");

	return 0;
}

void ras_session_destroy (RasSession* session) {
	ras_session_log (session, "session closed");
	ras_socktool_buffer_clear (&session->buffer, false);

	/* chain up to parent destructor */
	ras_conn_destroy (RAS_CONN (session));
}

char* ras_session_getline (RasSession* session, int* len) {
	char* newline;
	session->buffer.buf_error = false;
	session->buffer.buf_eof = false;
	while ((newline = ras_socktool_getline (
			RAS_CONN (session)->fd, &session->buffer, '\n', len)) == NULL &&
			!(session->buffer.buf_error) && !(session->buffer.buf_eof));
	return newline;
}

int ras_session_read_header (RasSession* session) {
	int onelinelen;
	char* oneline = ras_session_getline (session, &onelinelen);
	char* space;
	char* next;

	ras_socktool_buffer_free_line (&session->buffer);
	if (session->buffer.buf_error) {
		free (oneline);
		ras_session_log (session, "read error: %s", strerror (errno));
		return -1;
	}
	if (session->buffer.buf_eof) {
		free (oneline);
		ras_session_log (session, "connection closed by peer");
		return -1;
	}

	const char usage_msg[] =
		"Usage: LOGIN    [USERNAME] (login without allocating a pseudo-tty)\n"
		"  or:  LOGINTTY [USERNAME] (login and allocate a pseudo-tty)\n"
		"If USERNAME is not specified, system default will be used.\n";
	const int usage_len = STATIC_STRLEN (usage_msg);

	space = strchr (oneline, ' ');
	if (space == NULL) {
		next = oneline + onelinelen;
	} else {
		*space = '\0';
		for (space++; *space == '\0' && *space == ' '; space++);
		next = space;
	}

	if (strcmp (oneline, "LOGIN") == 0) {
		session->shell_use_pty = false;
	} else if (strcmp (oneline, "LOGINTTY") == 0) {
		session->shell_use_pty = true;
	} else {
		const char badcmd_msg[] = "Error: Unknown login command.\n";
		const int badcmd_len = STATIC_STRLEN (badcmd_msg);
		ras_socktool_write_string (RAS_CONN (session)->fd, badcmd_msg, badcmd_len);
		ras_socktool_write_string (RAS_CONN (session)->fd, usage_msg, usage_len);
		ras_session_log (session, "command %s is not understood", oneline);
		free (oneline);
		return -1;
	}

	struct passwd* login_user = NULL;
	if (*next != '\0') {
		login_user = getpwnam (next);
		if (login_user == NULL) {
			const char baduser_msg[] = "Error: User not found.\n";
			const int baduser_len = STATIC_STRLEN (baduser_msg);
			ras_socktool_write_string (RAS_CONN (session)->fd, baduser_msg, baduser_len);
			ras_session_log (session, "user %s cannot be found", next);
			free (oneline);
			return -1;
		}
	}

	free (oneline);

	if (RAS_CONN (session)->perm == RAS_CONN_PERM_RESTRICTED) {
		uid_t real_uid = getuid ();
		if (login_user != NULL && login_user->pw_uid != real_uid) {
			const char rest_msg[] = "Error: You are not permitted to login from here.\n";
			const int rest_len = STATIC_STRLEN (rest_msg);
			ras_socktool_write_string (RAS_CONN (session)->fd, rest_msg, rest_len);
			ras_session_log (session, "user %s is not permitted", login_user->pw_name);
			return -1;
		}
		setuid (real_uid);
	} else {
		const char sufail_msg[] = "Error: Unable to change the user credential.\n";
		const int sufail_len = STATIC_STRLEN (sufail_msg);
		int rval;

#if defined(HAVE_INITGROUPS)
		rval = initgroups (login_user->pw_name, login_user->pw_gid);
#elif defined(HAVE_SETGROUPS)
		struct group* groupentry;
		Array* grouplist = array_create (sizeof (gid_t), 1);
		int errno_saved;

		array_pushback (grouplist, &login_user->pw_gid);
		setgrent ();
		while ((groupentry = getgrent ()) != NULL) {
			for (int i = 0; groupentry->gr_mem[i] != NULL; i++) {
				if (strcmp (groupentry->gr_mem[i], login_user->pw_name) == 0) {
					array_pushback (grouplist, &groupentry->gr_gid);
				}
			}
		}
		endgrent ();

		rval = setgroups (array_getlen (grouplist), array_data (grouplist));

		errno_saved = errno;
		array_free (grouplist);
		errno = errno_saved;
#endif

		if (rval < 0) {
			write (RAS_CONN (session)->fd, sufail_msg, sufail_len);
			ras_session_log (session,
					"cannot set supplementary groups: %s", strerror (errno));
			return -1;
		}

		rval = setgid (login_user->pw_gid);
		if (rval < 0) {
			write (RAS_CONN (session)->fd, sufail_msg, sufail_len);
			ras_session_log (session,
					"cannot set primary group to %d: %s",
					login_user->pw_gid, strerror (errno));
			return -1;
		}

		rval = setuid (login_user->pw_uid);
		if (rval < 0) {
			write (RAS_CONN (session)->fd, sufail_msg, sufail_len);
			ras_session_log (session,
					"cannot set user to %s (uid = %d): %s",
					login_user->pw_name, login_user->pw_uid, strerror (errno));
			return -1;
		}
	}

	return 0;
}

static volatile sig_atomic_t child_terminated;
static pid_t child_pid;
static void session_child_terminated_setter (int signo) {
	int status;
	if (waitpid (child_pid, &status, WUNTRACED | WNOHANG) < 0) {
		return;
	}
	if (WIFEXITED (status) || WIFSIGNALED (status)) {
		child_terminated = true;
	} else if (WIFSTOPPED (status)) {
		kill (child_pid, SIGCONT);
	}
}

int ras_session_start_shell (RasSession* session) {
	const char noshell_msg[] = "Error: Unable to start shell.\n";
	const int noshell_len = STATIC_STRLEN (noshell_msg);
	const char nopty_msg[] = "Error: Cannot allocate a pseudo-tty.\n";
	const int nopty_len = STATIC_STRLEN (nopty_msg);
	int sockpair[2]; /* parent = 0, child = 1 */
	char* slave_dev = NULL;

	child_terminated = false;
	child_pid = -1;
	struct sigaction session_action = {
		.sa_handler = session_child_terminated_setter,
		.sa_flags = 0
	};
	sigaction (SIGCHLD, &session_action, NULL);

	if (session->shell_use_pty) {
		session->shell_stdio = posix_openpt (O_RDWR | O_NOCTTY);
		if (session->shell_stdio < 0) {
			write (RAS_CONN (session)->fd, nopty_msg, nopty_len);
			ras_session_log (session, "cannot posix_openpt: %s", strerror (errno));
			return -1;
		}
		if (grantpt (session->shell_stdio) < 0) {
			write (RAS_CONN (session)->fd, nopty_msg, nopty_len);
			ras_session_log (session, "cannot grantpt: %s", strerror (errno));
			close (session->shell_stdio);
			return -1;
		}
		if (unlockpt (session->shell_stdio) < 0) {
			write (RAS_CONN (session)->fd, nopty_msg, nopty_len);
			ras_session_log (session, "cannot unlockpt: %s", strerror (errno));
			close (session->shell_stdio);
			return -1;
		}
		if ((slave_dev = ptsname (session->shell_stdio)) == NULL) {
			write (RAS_CONN (session)->fd, nopty_msg, nopty_len);
			ras_session_log (session, "cannot ptsname: %s", strerror (errno));
			close (session->shell_stdio);
			return -1;
		}
	} else {
		if (socketpair (AF_LOCAL, SOCK_STREAM, 0, sockpair) < 0) {
			write (RAS_CONN (session)->fd, noshell_msg, noshell_len);
			ras_session_log (session, "cannot socketpair: %s", strerror (errno));
			return -1;
		}
		session->shell_stdio = sockpair[0];
	}

	fflush (stdout);
	fflush (stderr);

	session->shell_pid = fork ();
	if (session->shell_pid < 0) {
		write (RAS_CONN (session)->fd, noshell_msg, noshell_len);
		ras_session_log (session, "cannot fork to run shell: %s", strerror (errno));
		return -1;
	} else if (session->shell_pid > 0) {
		/* parent process */
		fd_set rset, rresult;
		int child, sock, nfds;

		child_pid = session->shell_pid;
		child = session->shell_stdio;
		sock = RAS_CONN (session)->fd;
		nfds = xmax (child, sock) + 1;

		FD_ZERO (&rset);
		FD_SET (child, &rset);
		FD_SET (sock, &rset);

		while (!child_terminated) {
			rresult = rset;
			if (select (nfds, &rresult, NULL, NULL, NULL) >= 0) {
				if (session->buffer.buf_len == 0) {
					if (FD_ISSET (child, &rresult)) {
						int rval = read (child, session->buffer.buf, RAS_BUFFER_SIZE);
						if (rval <= 0) {
							continue;
						}
						ras_socktool_write_string (sock, session->buffer.buf, rval);
					}
					if (FD_ISSET (sock, &rresult)) {
						int rval = read (sock, session->buffer.buf, RAS_BUFFER_SIZE);
						if (rval <= 0) {
							continue;
						}
						ras_socktool_write_string (child, session->buffer.buf, rval);
					}
				} else {
					ras_socktool_write_string (
							child,
							session->buffer.buf + session->buffer.buf_start,
							session->buffer.buf_len);
					session->buffer.buf_len = 0;
				}
			}
		}
	} else {
		/* child process */
		int child_stdio;

		if (slave_dev != NULL) {
			child_stdio = open (slave_dev, O_RDWR);
			if (child_stdio < 0) {
				write (RAS_CONN (session)->fd, nopty_msg, nopty_len);
				ras_session_log (session, "cannot open slave pty: %s", strerror (errno));
				exit (1);
			}
#ifdef BSD
			if (ioctl (child_stdio, TIOCSCTTY, NULL) < 0) {
				write (RAS_CONN (session)->fd, nopty_msg, nopty_len);
				ras_session_log (session, "cannot set controlling terminal: %s", strerror (errno));
				exit (1);
			}
#endif
		} else {
			child_stdio = sockpair[1];
		}

		ras_socktool_reset_signals ();

		ras_socktool_close_except (child_stdio);
		dup2 (child_stdio, STDIN_FILENO);
		dup2 (child_stdio, STDOUT_FILENO);
		dup2 (child_stdio, STDERR_FILENO);
		close (child_stdio);

		char* argv[] = { "-ras-shell", NULL };
		exit (ras_shell_main (1, argv));
	}

	return 0;
}