/*
    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
#define _DEFAULT_SOURCE

#include <sys/types.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 <security/pam_appl.h>
#include <security/pam_misc.h>

static void usage (void);
static int perm_set (struct passwd *, struct group *);
static int authenticate (const 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 *);

// FIXME: misc_conv is a separate library, should stick to plain PAM or make
// our own pam module
static struct pam_conv conv = {misc_conv, NULL};

int main (int argc, char *argv[])
{
	// TODO: Add arguments
	// FIXME: change the default program to execute SHELL
	char *t_usr = "root", *t_grp = NULL;
	struct passwd *t_pw;
	struct group *t_gr;
	int opt, err;
	int shellflag = 0;
	while ((opt = getopt(argc, argv, "A:u:g:C:s")) != -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 '?':
			usage();
			exit(EINVAL);
			break;
		}
	}

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

	t_pw = user_to_passwd(t_usr);
	if (!t_pw) {
		fprintf(stderr, "user_to_passwd: %s\n", strerror(errno));
		return errno;
	}
	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 = malloc(sizeof(char *) * (c_argc + 1));
		if (!c_argv) {
			fprintf(stderr, "malloc: %s\n", strerror(errno));
			exit(errno);
		}
		for (int i = 0; optind < argc; optind++, i++) {
			c_argv[i] = strdup(argv[optind]);
			if (!c_argv[i]) {
				fprintf(stderr, "strdup: %s\n", strerror(errno));
				exit(errno);
			}
		}
	} else {
		c_argc = 1;
		c_argv = malloc(sizeof(char *) * (c_argc + 1));
		if (!c_argv) {
			fprintf(stderr, "malloc: %s\n", strerror(errno));
			exit(errno);
		}
		c_argv[0] = strdup(shell);
		if (!c_argv[0]) {
			fprintf(stderr, "strdup: %s\n", strerror(errno));
			exit(errno);
		}
	}
	c_argv[c_argc] = NULL;

	/* Authenticate */
	// FIXME: move this up
	if (authenticate(uname) != PAM_SUCCESS)
		exit(EXIT_FAILURE);

	// TODO: clean up env
	/* copy and filter env */
/*
	char **c_env;
	extern char **environ;
	int size = 0;
	for (int i = 0; environ[i]; i++, size++);
	c_env = malloc(sizeof(char *) * (size + 1));
	if (!c_env) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(errno);
	}
	for (int i = 0; environ[i]; i++) {
		c_env[i] = strdup(environ[i]);
		if (!c_env[i]) {
			fprintf(stderr, "strdup: %s\n", strerror(errno));
			exit(errno);
		}
	}
	c_env[size] = NULL;
*/
	// TODO: check err value
	// TODO: add all this to list and loop over it
	err = setenv("USER", t_pw->pw_name, 1);
	err = setenv("LOGNAME", t_pw->pw_name, 1);
	err = setenv("SHELL", t_pw->pw_shell, 1);
	err = setenv("HOME", t_pw->pw_dir, 1);

	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, "execl: %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 [-s] [-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;
	}

	// FIXME: ideally when failing reset the permissions
	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 (const char *uname)
{
	pam_handle_t *pamh;
	int pam_err, count = 0;
	pam_err = pam_start("User Switcher", uname, &conv, &pamh);
	if (pam_err != PAM_SUCCESS) {
		fprintf(stderr, "pam_start: %s\n", pam_strerror(pamh, pam_err));
		return pam_err;
	}

	do {
		pam_err = pam_authenticate(pamh, 0);
		if (pam_err != PAM_SUCCESS)
			printf("Auth failed: %s\n", pam_strerror(pamh, pam_err));
		// FIXME: count gets ignored because authentication service has
		// a set amount of retries giving an error:
		// Have exhausted maximum number of retries for service
		count++;
	} while (pam_err != PAM_SUCCESS && count < 4);
	if (pam_err != PAM_SUCCESS) {
		fprintf(stderr, "better luck next time\n");
		pam_end(pamh, pam_err);
		return pam_err;
	}
	// FIXME: check again for the validity of the login for more security
	// as in: https://docs.oracle.com/cd/E19120-01/open.solaris/819-2145/pam-20/index.html
	// FIXME: ^C [SIGINT] will interrupt this call possibly causing a
	// vulnerability

	return pam_end(pamh, pam_err);
}

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

/*
static int execvpe(const char *file, char *const argv[], char *const envp[])
{
	const char *p, *z, *path = getenv("PATH");
	size_t l, k;

	errno = ENOENT;
	if (!*file) return -1;

	if (strchr(file, '/'))
		return execve(file, argv, envp);

	if (!path) path = "/usr/local/bin:/bin:/usr/bin";
	k = strnlen(file, NAME_MAX+1);
	if (k > NAME_MAX) {
		errno = ENAMETOOLONG;
		return -1;
	}
	l = strnlen(path, PATH_MAX-1)+1;

	for(p=path; ; p=z) {
		char b[l+k+1];
		z = strchr(p, ':');
		if (!z) z = p+strlen(p);
		if ((size_t)(z-p) >= l) {
			if (!*z++) break;
			continue;
		}
		memcpy(b, p, z-p);
		b[z-p] = '/';
		memcpy(b+(z-p)+(z>p), file, k+1);
		execve(b, argv, envp);
		if (errno != ENOENT) return -1;
		if (!*z++) break;
	}
	return -1;
}
*/