/*
    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 <sys/stat.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>
#include <ctype.h>
#include <signal.h>

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

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

#define MAX_HASH 1024
#define CONF_LINE_MAX 1024
#define GROUPS_MAX 256
#define FLAG_PERSIST 0x1
#define FLAG_NOPASS 0x2
#define FLAG_NOLOG 0x4

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

struct config {
	int type;
	int flags;
	char *who;
	char *as;
	struct env_elem *env;
	int env_n;
};

static void *emalloc(size_t);
static char *estrdup(const char *);
void *erealloc(void *, size_t);
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 get_config(struct config **, int *);

extern char **environ;
char *config_file = "/etc/us.conf";

int main(int argc, char *argv[])
{
	char *t_usr = "root", *t_grp = NULL;
	struct passwd *t_pw;
	struct group *t_gr;
	int opt, err;
	int shellflag = 0, envflag = 0, askpass = 0;
	while ((opt = getopt(argc, argv, "Au:g:C:se")) != -1) {
		switch (opt) {
		case 'A':
			askpass = 1;
			break;
		case 'u':
			t_usr = optarg;
			break;
		case 'g':
			t_grp = optarg;
			break;
		case 'C':
			config_file = optarg;
			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 = my_pw->pw_name;
	gid_t my_groups[GROUPS_MAX];
	int n_groups = 0;
	if ((n_groups = getgroups(GROUPS_MAX, my_groups)) == -1)
		die("getgroups:");

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

	/* Don't have to wait for children */
	struct sigaction sa = {0};
	sa.sa_handler = SIG_DFL;
	sa.sa_flags = SA_NOCLDWAIT;
	if (sigaction(SIGCHLD, &sa, NULL) == -1)
		die("sigaction:");

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

	/* get info from the config file and check if the action we want to
	 * do is permitted */
	struct env_elem *env_extra = NULL;
	struct config *conf = NULL;
	int conf_num, conf_flags = 0, env_extra_n = 0;
	if (get_config(&conf, &conf_num) == -1)
		die("get_config: invalid arguments");
	int here = 0;
	for (int i = 0; i < conf_num; i++) {
		struct passwd *who_pw, *as_pw;
		struct group *who_gr, *as_gr;
		int who_usr = conf[i].who[0] == ':' ? 0 : 1;
		int as_usr = conf[i].as[0] == ':' ? 0 : 1;
		if (who_usr) {
			who_pw = user_to_passwd(conf[i].who);
			if (my_pw->pw_uid != who_pw->pw_uid) {
				free(who_pw);
				continue;
			}
		} else {
			who_gr = group_to_grp(conf[i].who);
			gid_t w_gid = who_gr->gr_gid;
			int x = 0;
			for (; x < n_groups && w_gid != my_groups[x]; x++);
			if (w_gid != my_groups[x]) {
				free(who_gr);
				continue;
			}
		}

		if (as_usr) {
			as_pw = user_to_passwd(conf[i].as);
			if (!as_pw)
				die("%s not a valid user", conf[i].as);
			if (t_pw->pw_uid != as_pw->pw_uid) {
				free(as_pw);
				continue;
			}
		} else if (t_gr) {
			as_gr = group_to_grp(conf[i].as);
			if (t_gr->gr_gid != as_gr->gr_gid) {
				free(as_gr);
				continue;
			}
		} else {
			continue;
		}
		here = 1;
		if (conf[i].type == 0)
			die("Permission denied");

		conf_flags |= conf[i].flags;
		if (conf[i].env_n) {
			env_extra = erealloc(env_extra,
			    (env_extra_n+conf[i].env_n+1)*sizeof(struct env_elem));
			for (int j = env_extra_n, x = 0; j < env_extra_n + conf[i].env_n; j++, x++)
				env_extra[j] = conf[i].env[x];
			env_extra_n += conf[i].env_n;
			env_extra[env_extra_n] = (struct env_elem){NULL,NULL};
		}

	}
	if (!here)
		die("no rule found for user %s", my_name);

	/* Authenticate, we will be root from now on */
	if (!(conf_flags & FLAG_NOPASS))
		if (authenticate(my_pw->pw_uid, askpass))
			exit(EXIT_FAILURE);

	/* 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 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;
		}
	}

	for (int i = 0; envflag &&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;
			}
		}
	}

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

	// do not override, we might be under more levels of 'us'
	err = setenv("US_USER", my_name, 0);

	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 [-seA] [-u user] [-g group] [-C config] 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__)
		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, parent = getpid();
		int pipefd[2];
		if (pipe(pipefd) == -1)
			die("pipe:");
		switch (pid = fork()) {
		case -1:
			die("fork:");
			break;
		case 0:
			/* we are still root, drop off permissions before
			 * disasters happen */
			if (setuid(uid) == -1) {
				kill(parent, SIGTERM);
				die("askpass: setuid:");
			}
			if (dup2(pipefd[1], STDOUT_FILENO) == -1) {
				kill(parent, SIGTERM);
				die("askpass: dup2:");
			}
			close(pipefd[0]);
			execl("/bin/sh", "sh", "-c", askpass, (char *)NULL);
			kill(parent, SIGTERM);
			die("askpass: execl:");
			break;
		default:
			fd = pipefd[0];
			close(pipefd[1]);
			break;
		}
	} else {
		printf("Password (%s): ", pw->pw_name);
		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 && askpass) {
		close(fd);
	} else {
		if (tcsetattr(tty_fd, TCSANOW, &tio_before) == -1)
			die("tcsetattr:");
	}

	char *enc = crypt(pass, hash);
	/* clean pass from memory */
	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, *pr;
	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;
	}
	pr = emalloc(sizeof(struct passwd));
	memcpy(pr, pw, sizeof(struct passwd));
	return pr;
}

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

	struct group *gr, *rr;
	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;
	}
	rr = emalloc(sizeof(struct group));
	memcpy(rr, gr, sizeof(struct group));
	return rr;
}

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 = calloc(1, s);
	if (!p)
		die("bad malloc:");
	return p;
}

void *erealloc(void *p, size_t s)
{
	if (!s || s == (size_t)-1)
		die("bad realloc: invalid size");
	void *r = realloc(p, s);
	if (!r)
		die("bad realloc:");
	return r;
}

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;
}

static int get_config(struct config **conf, int *num)
{
	if (!conf || !num)
		return -1;
	FILE *fp = fopen(config_file, "r");
	if (!fp)
		die("fopen:");
	struct stat st;
	if (fstat(fileno(fp), &st) == -1)
		die("fstat:");
	if (st.st_uid != 0 || st.st_gid != 0)
		die("config file must be owned by root:root");
	if (!S_ISREG(st.st_mode))
		die("config file must be a regular file");
	if (st.st_mode & S_IRWXO || st.st_mode & S_IROTH ||
	    st.st_mode & S_IWOTH || st.st_mode & S_IXOTH)
	    	die("others may not modify, read or execute config file");
	char line[CONF_LINE_MAX];
	*num = 0;
	*conf =  NULL;
	for (int i = 0; fgets(line, CONF_LINE_MAX, fp); i++) {
		char *s, *t, *sv;
		int n = 0;
		for (int ll = strlen(line)-1; isspace(line[ll]); line[ll--] = '\0');
		struct config c = {0};
		for (s = line;; s = NULL, n++) {
			int getflags = 0;
			if (!(t = strtok_r(s, " ", &sv)))
				break;
			if (*t == '#')
				break;
			switch (n) {
			case 0:
				if (!strcmp(t, "+"))
					c.type = 1;
				else if (!strcmp(t, "-"))
					c.type = 0;
				else
					die("non valid config line %d", i);
				break;
			case 1:
				c.who = estrdup(t);
				break;
			case 2:
				if (strcmp(t, "as"))
					die("non valid config line %d", i);
				break;
			case 3:
				c.as = estrdup(t);
				break;
			default:
				getflags = 1;
				break;
			}
			if (!getflags)
				continue;

			if (isupper(*t)) {
				char *e, *se, *et;
				for (e = t;; e = NULL) {
					if (!(et = strtok_r(e, ",", &se)))
						break;
					char *sep;
					int x;
					if (!(sep = strchr(et, ',')))
						die("invalid env at %d", i);
					c.env = erealloc(c.env, (c.env_n+1)*sizeof(struct env_elem));
					x = c.env_n;
					*sep = '\0';
					for (char *p = et; *p; p++)
						if (!isupper(*p) && *p != '-'
						    && *p != '_')
						    	die("non valid"
						    	"env at %d", i);
					c.env[x].name = estrdup(et);
					c.env[x].value = estrdup(sep+1);
					c.env_n++;
				}
			} else if (!strcmp(t, "persist")) {
				c.flags |= FLAG_PERSIST;
			} else if (!strcmp(t, "nopass")) {
				c.flags |= FLAG_NOPASS;
			} else if (!strcmp(t, "nolog")) {
				c.flags |= FLAG_NOLOG;
			}
		}
		if (n < 3)
			die("non valid config line %d", i);
		*conf = erealloc(*conf, ((*num)+1)*sizeof(struct config));
		(*conf)[(*num)] = c;
		*num += 1;
	}
	fclose(fp);
	return 0;
}