/*
    us - User Switcher
    Copyright (C) 2021  Alessandro Mauri

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License version 3.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

    e-mail: alemauri001@tuta.io
*/

#define _POSIX_C_SOURCE 200809L
#ifdef __linux__
#define _DEFAULT_SOURCE
#else
#define _BSD_SOURCE
#endif

#include <sys/types.h>
#include <sys/wait.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <pwd.h>
#include <unistd.h>
#include <limits.h>
#include <grp.h>
#include <termios.h>
#include <stdarg.h>

#if !defined(_XOPEN_CRYPT) || _XOPEN_CRYPT == -1
#include <crypt.h>
#endif

#if defined(__linux__)
#include <shadow.h>
#endif

#define MAX_HASH 1024

static void *emalloc(size_t);
static char *estrdup(const char *);
static void usage(void);
static void die(const char *, ...);
static int perm_set(struct passwd *, struct group *);
static int authenticate(uid_t, char);
static struct passwd* user_to_passwd(const char *);
static struct group* group_to_grp(const char *);
//static int execvpe(const char *, char *const *, char *const *);

extern char **environ;

int main(int argc, char *argv[])
{
	// TODO: modify signal handlers to ignore sigchld to not make zombies
	char *t_usr = "root", *t_grp = NULL;
	struct passwd *t_pw;
	struct group *t_gr;
	int opt, err;
	int shellflag = 0, envflag = 0;
	while ((opt = getopt(argc, argv, "A:u:g:C:se")) != -1) {
		switch (opt) {
		case 'A':
			printf("-A is not yet implemented\n");
			exit(EXIT_FAILURE);
			break;
		case 'u':
			t_usr = optarg;
			break;
		case 'g':
			t_grp = optarg;
			break;
		case 'C':
			printf("-C is not yet implemented\n");
			exit(EXIT_FAILURE);
			break;
		case 's':
			shellflag = 1;
			break;
		case 'e':
			envflag = 1;
			break;
		case '?':
			usage();
			exit(EINVAL);
			break;
		}
	}

	/* Get user info */
	char *shell;
	uid_t ruid = getuid();
	struct passwd *my_pw = getpwuid(ruid);
	if (!my_pw) {
		fprintf(stderr, "getpwid: %s\n", strerror(errno));
		return errno;
	}
	char *my_name = estrdup(my_pw->pw_name);

	/* From now on most actions require root */
	if (setuid(0) == -1)
		die("setuid:");

	/* Authenticate, we will be root from now on */
	if (authenticate(my_pw->pw_uid, 0))
		exit(EXIT_FAILURE);

	/* Get target user and group info */
	t_pw = user_to_passwd(t_usr);
	if (!t_pw)
		die("usr_to_passwd:");
	t_gr = group_to_grp(t_grp);

	/* Get target user's shell */
	if (!shellflag)
		shell = t_pw->pw_shell;
	else
		shell = getenv("SHELL");
	if (!shell)
		shell = "/bin/sh";

	/* Set argc and argv */
	int c_argc = argc - optind;
	char **c_argv;
	if (c_argc) {
		c_argv = emalloc(sizeof(char *) * (c_argc + 1));
		for (int i = 0; optind < argc; optind++, i++)
			c_argv[i] = estrdup(argv[optind]);
	} else {
		c_argc = 1;
		c_argv = emalloc(sizeof(char *) * (c_argc + 1));
		c_argv[0] = estrdup(shell);
	}
	c_argv[c_argc] = NULL;

	struct env_elem {
		char *name;
		char *value;
	};

	struct env_elem env_keep[] = {
		{"PATH", NULL},
		{"TERM", NULL},
		{"EDITOR", NULL},
		{"VISUAL", NULL},
		{"DISPLAY", NULL},
		{"XAUTHORITY", NULL},
		{NULL, NULL}
	};

	struct env_elem env_mod[] = {
		{"USER", t_pw->pw_name},
		{"LOGNAME", t_pw->pw_name},
		{"SHELL", t_pw->pw_shell},
		{"HOME", t_pw->pw_dir},
		{NULL, NULL}
	};

	if (envflag) {	/* clear env */
		for (int i = 0; env_keep[i].name; i++)
			env_keep[i].value = estrdup(getenv(env_keep[i].name));
		environ = NULL; // in place of clearenv
	}

	for (int i = 0; env_mod[i].name; i++) {
		err = setenv(env_mod[i].name, env_mod[i].value, 1);
		if (err == -1) {
			fprintf(stderr, "setenv: %s\n", strerror(errno));
			goto fail_end;
		}
	}

	if (envflag) {
		for (int i = 0; env_keep[i].name; i++) {
			if (env_keep[i].value) {
				err = setenv(env_keep[i].name, env_keep[i].value, 1);
				if (err == -1) {
					fprintf(stderr, "setenv: %s\n", strerror(errno));
					goto fail_end;
				}
			}
		}
	}
	// do not override, we might be under more levels of 'us'
	err = setenv("US_USER", my_name, 0);
	free(my_name);

	errno = 0;
	/* Set permissions */
	if (perm_set(t_pw, t_gr) == -1) {
		fprintf(stderr, "perm_set: %s\n", strerror(errno));
		goto fail_end;
	}

	/* Execute the command */
	err = execvp(c_argv[0], c_argv);
	if (err == -1)
		fprintf(stderr, "execvp: %s\n", strerror(errno));

	/* Cleanup and return */
	fail_end:
	/* Free up the copied argv */
	for (int i=0; c_argv[i]; i++)
		free(c_argv[i]);
	free(c_argv);
	return errno;
}

static inline void usage(void)
{
	// TODO: planned options
	// -a [program]: like sudo's askpass
	// -u [user]: change the default user from root to user
	// -g [group]: change the primary group to [gorup]
	// 	both -a and -g will accept numbers with #[num] like sudo
	// -c [file]: manually select config file
	// something about environment
	// something about non interactiveness
	printf("usage: us [-se] [-u user] [-g group] command [args]\n");
}

static int perm_set(struct passwd *pw, struct group *gr)
{
	if (!pw) {
		errno = EINVAL;
		return -1;
	}

	uid_t uid;
	gid_t gid;

	uid = pw->pw_uid;
	gid = pw->pw_gid;

	if (gr)
		gid = gr->gr_gid;

	/* Set permissions, setting group perms first because in the case of
	 * dropping from higher permissions setting the uid first results in
	 * an error */
	int err;
	/* Non POSIX but implemented in most systems anyways */
	err = initgroups(pw->pw_name, pw->pw_gid);
	if (err == -1) {
		printf("initgroups failed\n");
		return -1;
	}

	if (setregid(gid, gid) == -1) {
		printf("setregid failed\n");
		return -1;
	}

	if (setreuid(uid, uid) == -1) {
		printf("setreuid failed\n");
		return -1;
	}

	return 0;
}

static int authenticate(uid_t uid, char ask)
{
	// TODO: implement u2f compat
	/* get the encrypted password */
	struct passwd *pw = getpwuid(uid);
	char *hash_p, hash[MAX_HASH];
	char *p = pw->pw_passwd;
	int tty_fd = STDOUT_FILENO;

	if (!strcmp(p, "x") || *p == '*' || *p == '!') {
#if defined(__linux__)
		// get exclusive access to shadow
		if (lckpwdf() == -1)
			die("lckpwdf");
		setspent();
		struct spwd *sp = getspnam(pw->pw_name);
		if (sp == NULL)
			die("getspnam");
		hash_p = sp->sp_pwdp;
		endspent();
		ulckpwdf();
#elif defined(__OpenBSD__)
		// TODO: openbsd has getpwuid_passwd
		struct passwd *op = getpwuid_shadow(uid);
		if (!op)
			die("getpwuid_shadow:");
		hash_p = op->pw_passwd;
#endif
	} else {
		hash_p = pw->pw_passwd;
	}
	strncpy(hash, hash_p, MAX_HASH - 2);
	hash[MAX_HASH - 1] = '\0';
	if (strlen(hash) >= MAX_HASH - 3)
		die("password hash too long :^)");

	int fd = STDIN_FILENO;
	char *askpass = getenv("US_ASKPASS");
	char pass[1024] = {0};
	struct termios tio_before, tio_pass;
	if (ask && askpass) {
		pid_t pid;
		int pipefd[2];
		if (pipe(pipefd) == -1)
			die("pipe:");
		switch (pid = fork()) {
		case -1:
			die("fork:");
			break;
		case 0:
			if (dup2(pipefd[1], STDOUT_FILENO) == -1) {
				// TODO: send signal to parent dup failed
				die("dup2:");
			}
			close(pipefd[0]);
			execl(askpass, "", (char *)NULL);
			die("execl:");
			break;
		default:
			fd = pipefd[0];
			close(pipefd[1]);
			break;
		}
	} else if (ask) {
		die("askpass: no valid askpass specified");
	} else {
		printf("Password: ");
		fflush(stdout);
		if (tcgetattr(tty_fd, &tio_before) == -1)
			die("tcgetattr:");
		tio_pass = tio_before;
		/* Do not echo and accept when enter is pressed */
		tio_pass.c_lflag &= ~ECHO;
		tio_pass.c_lflag |= ICANON;
		if (tcsetattr(tty_fd, TCSANOW, &tio_pass) == -1)
			die("tcsetattr:");
	}
	int r = read(fd, pass, 1023);
	if (!r || r == -1) {
		if (errno)
			fprintf(stderr, "read: %s\n", strerror(errno));
		else
			printf("Password can't be zero length\n");
		// we may have been interrupted, wait askpass before bailing
		waitpid(-1, NULL, 0);
		exit(EXIT_FAILURE);
	}
	pass[1023] = '\0';
	// remove \n
	int l = strlen(pass);
	if (pass[l-1] == '\n')
		pass[--l] = '\0';

	if (ask) {
		close(fd);
	} else {
		if (tcsetattr(tty_fd, TCSANOW, &tio_before) == -1)
			die("tcsetattr:");
	}

	char *enc = crypt(pass, hash);
	memset(pass, 0, 1024);
	if (strncmp(hash, enc, 1024)) {
		printf("Authentication failure\n");
		setuid(uid);
		return -1;
	}
	printf("\n");
	return 0;
}

static struct passwd* user_to_passwd(const char *user)
{
	if (!user) {
		errno = EINVAL;
		return NULL;
	}

	struct passwd* pw;
	long uid_l;
	errno = 0;
	if (user[0] != '#') {
		pw = getpwnam(user);
	} else {
		uid_l = strtol(&user[1], NULL, 10);
		if (uid_l < 0 || errno) {
			errno = errno ? errno : EINVAL;
			return NULL;
		}
		pw = getpwuid((uid_t)uid_l);
	}

	if (!pw) {
		if (!errno)
			errno = EINVAL;
		return NULL;
	}

	return pw;
}

static struct group* group_to_grp(const char *group)
{
	if (!group) {
		errno = EINVAL;
		return NULL;
	}

	struct group* gr;
	long gid_l;
	errno = 0;
	if (group[0] != '#') {
		gr = getgrnam(group);
	} else {
		gid_l = strtol(&group[1], NULL, 10);
		if (gid_l < 0 || errno) {
			errno = errno ? errno : EINVAL;
			return NULL;
		}
		gr = getgrgid((gid_t)gid_l);
	}

	if (!gr) {
		if (!errno)
			errno = EINVAL;
		return NULL;
	}
	return gr;
}

void die(const char *fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);

	vfprintf(stderr, fmt, ap);

	if (fmt[0] && fmt[strlen(fmt) - 1] == ':') {
		fputc(' ', stderr);
		perror(NULL);
	} else {
		fputc('\n', stderr);
	}

	va_end(ap);
	exit(errno ? errno : EXIT_FAILURE);
}

void *emalloc(size_t s)
{
	if (!s || s == (size_t)-1)
		die("bad malloc: invalid size");
	void *p = malloc(s);
	if (!p)
		die("bad malloc:");
	return p;
}

char *estrdup(const char *s)
{
	if (!s)
		die("bad strdup: cannot duplicate NULL pointer");
	char *r = strdup(s);
	if (!r)
		die("bad strdup:");
	return r;
}