#define _XOPEN_SOURCE 500
#include <ftw.h>
#include <librc.h>
#include <pwd.h>
#include <security/pam_appl.h>
#include <security/pam_modules.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <sys/file.h>
#include <syslog.h>
#include <unistd.h>

#include "einfo.h"
#include "queue.h"

static int
inc_dec_lockfile(const char *runtimedir, int val)
{
	char *lockfile_path = NULL;
	FILE *lockfile = NULL;

	int locknum = 0;

	elog(LOG_INFO, "locking lockfile");

	xasprintf(&lockfile_path, "%s/%s", runtimedir, "lock");
	lockfile = fopen(lockfile_path, "r+");
	if (!lockfile) {
		lockfile = fopen(lockfile_path, "w+");
		if (!lockfile) {
			eerror("fopen: failed to open file %s, %s", lockfile_path, strerror(errno));
			return -1;
		}
		if (flock(fileno(lockfile), LOCK_EX) != 0) {
			eerror("flock: %s", strerror(errno));
			fclose(lockfile);
			return -1;
		}
		locknum = 1;
	} else {
		if (flock(fileno(lockfile), LOCK_EX) != 0) {
			eerror("flock: %s", strerror(errno));
			fclose(lockfile);
			return -1;
		}
		fscanf(lockfile, "%d", &locknum);
		locknum += val;
		rewind(lockfile);
	}

	free(lockfile_path);

	fprintf(lockfile, "%d", locknum);

	if (flock(fileno(lockfile), LOCK_UN)) {
		eerror("flock: %s", strerror(errno));
	}
	fclose(lockfile);

	elog(LOG_INFO, "unlocking lockfile");

	return locknum;
}

static void load_envs_from_file(const char *path, RC_STRINGLIST *out) {
	FILE *fp = NULL;
	char *line = NULL;
	size_t n = 0;
	char *p = NULL;

	fp = fopen(path, "r");
	if (!fp) {
		return;
	}
	while (getline(&line, &n, fp) != -1) {
		if ((p = strchr(line, '\n'))) {
			*p = '\0';
		};
		rc_stringlist_addu(out, line);
	}
	fclose(fp);
}

static RC_STRINGLIST *load_dir(const char *dir) {
	RC_STRINGLIST *list = rc_stringlist_new();
	DIR *dp = NULL;
	struct dirent *d = NULL;
	char *path;

	if ((dp = opendir(dir)) != NULL) {
		while ((d = readdir(dp)) != NULL) {
			xasprintf(&path, "%s/%s", dir, d->d_name);
			load_envs_from_file(path, list);
		}
	}
	closedir(dp);

	return list;
}

static void set_user_env(pam_handle_t *pamh) {
	RC_STRINGLIST *allowed_env = NULL;
	RC_STRINGLIST *user_env = NULL;
	RC_STRING *env;
	RC_STRING *uenv;
	char *p;
	char *user_env_path;

	elog(LOG_INFO, "Loading allowed envs in %s", RC_USER_ENV_WHITELIST_D);
	allowed_env = load_dir(RC_USER_ENV_WHITELIST_D);

	elog(LOG_INFO, "Loading allowed envs in %s", RC_USER_ENV_WHITELIST);
	load_envs_from_file(RC_USER_ENV_WHITELIST, allowed_env);

	xasprintf(&user_env_path, "%s/openrc/env", pam_getenv(pamh, "XDG_RUNTIME_DIR"));

	elog(LOG_INFO, "Loading user envs in %s", user_env_path);
	user_env = load_dir(user_env_path);

	TAILQ_FOREACH(env, allowed_env, entries) {
		elog(LOG_INFO, "allowed env %s", env->value);
		TAILQ_FOREACH(uenv, user_env, entries) {
			p = strchr(uenv->value, '=');
			if (p) {
				*p = '\0';
				if (strcmp(env->value, uenv->value) == 0) {
					*p = '=';
					elog(LOG_INFO, "Exporting: %s", uenv->value);
					pam_putenv(pamh, uenv->value);
				} else {
					*p = '=';
				}
			}
		}
	}

	elog(LOG_INFO, "Finished loading user environment");

	rc_stringlist_free(allowed_env);
	rc_stringlist_free(user_env);
	free(user_env_path);
}

static int
exec_user_cmd(struct passwd *pw, char *cmd, char **envlist)
{
	int retval;
	const char *shellname = basename_c(pw->pw_shell);

	switch (fork()) {
		case 0:
			setgid(pw->pw_gid);
			setuid(pw->pw_uid);

			execle(pw->pw_shell, shellname, "-c", cmd, NULL, envlist);

			return -1;
			break;
		case -1:
			return -1;
			break;
	}
	wait(&retval);
	return retval;
}

static char *create_rc_runtime_dir(struct passwd *pw) {
	char *path = NULL;

	if (mkdir("/run/openrc/user", 0755) != 0 && errno != EEXIST) {
		elog(LOG_ERR, "Failed to mkdir %s: %s", path, strerror(errno));
		return NULL;
	}

	xasprintf(&path, "/run/openrc/user/%d/", pw->pw_uid);

	if (mkdir(path, 0700) != 0 && errno != EEXIST) {
		elog(LOG_ERR, "Failed to mkdir %s: %s", path, strerror(errno));
		free(path);
		return NULL;
	}

	return path;
}

static int clear_runtime_dir(const char *path, const struct stat *sb, int typeflag, struct FTW *ftwbuf) {
	int ret = remove(path) != 0;
	(void)sb;
	(void)typeflag;
	(void)ftwbuf;

	if (ret != 0) {
		elog(LOG_ERR, "Failed to remove %s: %s", path, strerror(errno));
	}

	return ret;
}

static bool exec_openrc(pam_handle_t *pamh, const char *runlevel, bool lock) {
	int lockval;
	char *cmd = NULL;
	const char *username;
	struct passwd *pw = NULL;
	char *openrc_runtime_dir;
	const char *xdg_runtime_dir;
	char *openrc_runtime_dir_env;
	char **envlist;
	char **env;

	if (pam_get_user(pamh, &username, "username:") != PAM_SUCCESS)
		return false;
	pw = getpwnam(username);
	if (!pw)
		return false;

	/* avoid setting RC_RUNTIME_DIR if XDG_RUNTIME_DIR is already set by the user's stack */
	if ((xdg_runtime_dir = pam_getenv(pamh, "XDG_RUNTIME_DIR"))) {
		elog(LOG_INFO, "Setting up RC_RUNTIME_DIR");
		openrc_runtime_dir = create_rc_runtime_dir(pw);
		if (!openrc_runtime_dir) {
			return false;
		}

		xasprintf(&openrc_runtime_dir_env, "RC_RUNTIME_DIR=%s", openrc_runtime_dir);
		pam_putenv(pamh, openrc_runtime_dir_env);
		elog(LOG_INFO, "Exporting: %s", openrc_runtime_dir_env);
		free(openrc_runtime_dir_env);
	} else {
		elog(LOG_INFO, "Using XDG_RUNTIME_DIR");
		xasprintf(&openrc_runtime_dir, "%s/%s", xdg_runtime_dir, "openrc");
		if (mkdir(openrc_runtime_dir, 0700) != 0 && errno != EEXIST) {
			free(openrc_runtime_dir);
			return false;
		}
	}

	if (chown(openrc_runtime_dir, pw->pw_uid, pw->pw_gid) != 0) {
		elog(LOG_ERR, "failed to chown %s", openrc_runtime_dir);
		rmdir(openrc_runtime_dir);
		free(openrc_runtime_dir);
		return false;
	}

	envlist = pam_getenvlist(pamh);

	xasprintf(&cmd, "openrc --user %s", runlevel);

	/* if we are locking, reduce the count by 1,
	 * because we don't want to count ourselves */
	lockval = inc_dec_lockfile(openrc_runtime_dir, lock ? 1 : -1) - lock == true ? 1 : 0;

	if (lockval == 0) {
		elog(LOG_INFO, "Executing %s for user %s", cmd, username);
		exec_user_cmd(pw, cmd, envlist);

		if (!lock) {
			nftw(openrc_runtime_dir, clear_runtime_dir, 64, FTW_DEPTH | FTW_PHYS);
		}
	}

	if (lock) {
		elog(LOG_INFO, "Setting the user's environment");
		set_user_env(pamh);
	}

	for (env = envlist; *env; env++)
		free(*env);
	free(openrc_runtime_dir);
	free(envlist);
	free(cmd);
	return true;
}

PAM_EXTERN int pam_sm_open_session(pam_handle_t *pamh, int flags, int argc, const char **argv) {
	const char *runlevel = argc > 0 ? runlevel = argv[0] : "default";
	int ret = PAM_SUCCESS;
	(void)flags;

	setenv("EINFO_LOG", "openrc-pam", 1);
	elog(LOG_INFO, "Opening openrc session");

	if (exec_openrc(pamh, runlevel, true)) {
		elog(LOG_INFO, "Openrc session opened");
	} else {
		elog(LOG_ERR, "Failed to open session");
		ret = PAM_SESSION_ERR;
	}

	unsetenv("EINFO_LOG");
	return ret;
}

PAM_EXTERN int pam_sm_close_session(pam_handle_t *pamh, int flags, int argc, const char **argv) {
	const char *runlevel = argc > 1 ? argv[1] : "none";
	int ret = PAM_SUCCESS;
	(void)flags;

	setenv("EINFO_LOG", "openrc-pam", 1);
	elog(LOG_INFO, "Closing openrc session");

	if (exec_openrc(pamh, runlevel, false)) {
		elog(LOG_INFO, "Openrc session closed");
	} else {
		elog(LOG_ERR, "Failed to close session");
		ret = PAM_SESSION_ERR;
	}

	unsetenv("EINFO_LOG");
	return ret;
}