#include <security/pam_modules.h>
#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include <sqlite3.h>
#include <string>
#include <sys/stat.h> // for mkdir
#include <unistd.h> // for chown
#include <errno.h>
#include <string.h>
#include <fstream>

#define MYNAME "pam_firstcomefirstserve"

struct database {
	sqlite3 *sql = nullptr;
	~database() {
		if(sql) {
			if(SQLITE_OK != sqlite3_close(sql)) {
				fprintf(stderr, MYNAME ": sqlite3_close failed: %s\n", sqlite3_errmsg(sql));
				// still continue, but with leaking memory
			}
			sql = nullptr;
		}
	}

	bool commit();

	database() {}
	database(database const&) = delete;
	database& operator=(database const&) = delete;
};

// error handling designed so that if an error occurs, all operations are ignored and then you check once at the end.
struct statement {
	sqlite3_stmt *stmt = nullptr;
	sqlite3 *db;
	bool ok = true;
	statement(database& db, const char *command) {
		this->db = db.sql;
		if(SQLITE_OK != sqlite3_prepare_v2(db.sql, command, -1, &stmt, nullptr)) {
			fprintf(stderr, MYNAME ": %s: %s\n", command, sqlite3_errmsg(db.sql));
			ok = false;
		}
	}

	// sqlite3_step returns an error, finished, or a row data.
	bool step() { // returns false if statement is finished (possibly with error), or true if we paused because another row is available to read from the results
		if(!stmt || !ok)
			return false;
		int result = sqlite3_step(stmt);
		if(result == SQLITE_ROW)
			return true;
		if(result != SQLITE_DONE) {
			fprintf(stderr, MYNAME ": %s\n", sqlite3_errmsg(db));
			ok = false;
		}
		return false;
	}

	void bind(int paramNum, std::string const& string) { // string must outlive statement object is destroyed (that's what SQLITE_STATIC means)
		if(!stmt || !ok) return;

		if(SQLITE_OK != sqlite3_bind_text(stmt, paramNum, string.data(), string.size(), SQLITE_STATIC)) {
			fprintf(stderr, MYNAME ": sqlite3_bind_text: %s\n", sqlite3_errmsg(db));
			ok = false;
		}
	}

	// undefined results (including crash) if row data is accessed any time except after step returns true, or if type doesn't match
	// note: parameters start at 1, but columns start at 0
	std::string getColumnText(int columnNum) {
		const char *data = (const char *)sqlite3_column_text(stmt, columnNum);
		int length = sqlite3_column_bytes(stmt, columnNum);
		return std::string(data, data+length);
	}

	~statement() {
		if(stmt) {
			if(sqlite3_finalize(stmt) != SQLITE_OK) {
				fprintf(stderr, MYNAME ": sqlite3_finalize: %s\n", sqlite3_errmsg(db));
				// leak memory, but continue anyway
			}
			stmt = nullptr;
		}
	}

	statement(statement const&) = delete;
	statement& operator=(statement const&) = delete;
};

bool database::commit() {
	statement st(*this, "commit");
	while(st.step()) {} // should complete in one step
	return st.ok;
}

int pam_sm_authenticate(pam_handle_t *pamh, int flags, int argc, const char **argv) {
	struct pam_conv *conv;
	if(PAM_SUCCESS != pam_get_item(pamh, PAM_CONV, (const void**)&conv)) {
		fprintf(stderr, MYNAME ": auth: pam_get_item(PAM_CONV) failed (programming error in application)\n");
		return PAM_AUTHINFO_UNAVAIL;
	}

	const char *username_;
	if(PAM_SUCCESS != pam_get_item(pamh, PAM_USER, (const void**)&username_)) {
		fprintf(stderr, MYNAME ": auth: pam_get_item(PAM_USER) failed (programming error in application)\n");
		return PAM_AUTHINFO_UNAVAIL;
	}
	std::string username = username_;

	// Since we create directories based on usernames, we have to enforce some basic restrictions
	// so we don't create /srv/upload/../../etc/passwd or something.
	// Usernames are not empty, don't start with "." and contain alphanumeric characters, '_', '-' and '.'.
	for(char c : username) {
		if(!isalnum(c) && c != '_' && c != '.' && c != '-') {
			fprintf(stderr, MYNAME ": invalid username\n");
			return PAM_AUTH_ERR;
		}
	}
	if(username.size() == 0 || username[0] == '.') {
		fprintf(stderr, MYNAME ": invalid username\n");
		return PAM_AUTH_ERR;
	}

	// argv does not include the module name: argv[0] is the first argument, and argc is 0 if there are no arguments.
	if(argc != 4) {
		fprintf(stderr, MYNAME ": auth: wrong number of arguments. arguments must be database path, parent of user home directories, home dir UID and home dir GID\n");
		return PAM_AUTHINFO_UNAVAIL;
	}
	int home_dir_uid = atoi(argv[2]), home_dir_gid = atoi(argv[3]);
	if(!home_dir_uid || !home_dir_gid) { // conveniently catches non-numbers and root
		fprintf(stderr, MYNAME ": auth: please check home uid/gid (args 3/4)\n");
		return PAM_AUTHINFO_UNAVAIL;
	}


	std::string password;
	{
		struct pam_message msg = {PAM_PROMPT_ECHO_OFF, "Password: "};
		const struct pam_message *msgs = {&msg};
		struct pam_response *responses = NULL;
		if(PAM_SUCCESS != conv->conv(1, &msgs, &responses, conv->appdata_ptr)) {
			fprintf(stderr, MYNAME ": auth: password prompt failed\n");
			// do not free responses in this case? "The application should not set *resp"
			return PAM_AUTHINFO_UNAVAIL;
		}
		password = responses[0].resp;
		free(responses[0].resp);
		free(responses);
	};

	database db;
	if(SQLITE_OK != sqlite3_open(argv[0], &db.sql)) {
		fprintf(stderr, MYNAME ": auth: sqlite3_open failed\n");
		return PAM_AUTHINFO_UNAVAIL;
	}

	if(SQLITE_OK != sqlite3_busy_timeout(db.sql, 10000)) { // milliseconds. Call should always succeed?
		fprintf(stderr, MYNAME ": auth: sqlite3_busy_timeout: %s\n", sqlite3_errmsg(db.sql));
	}

	// Don't make any database changes until after reading the password, because SQLite may take a lock

	{
		statement st(db, "begin transaction");
		while(st.step()) {}
		if(!st.ok) {
			fprintf(stderr, MYNAME ": auth: failed to begin database transaction\n");
			return PAM_AUTHINFO_UNAVAIL;
		}
	}

	{
		statement st(db, "create table if not exists users (username varchar not null primary key, password varchar not null)");
		while(st.step()) {} // should finish in one step, but anyway
		// no st.ok check - if an error occurred, it was already printed, and we continue anyway -
		// if the error is still present, we'll get the same error when accessing the data, anyway.
	}

	auto make_home_dir_and_subdirs = [&]() {
		auto make_home_dir = [&](std::string const& subdir) {
			std::string createpath = std::string(argv[1])+"/"+username;
			if(subdir != "") createpath += "/" + subdir;
			if(mkdir(createpath.c_str(), 0700) && errno != EEXIST) {
				fprintf(stderr, "mkdir %s: %s\n", createpath.c_str(), strerror(errno));
			}
			if(chown(createpath.c_str(), home_dir_uid, home_dir_gid)) {
				fprintf(stderr, "chown %s: %s\n", createpath.c_str(), strerror(errno));
			}
		};
		make_home_dir("");

		std::ifstream in("/srv/ftp-subdirs");
		std::string subdir;
		while(std::getline(in, subdir)) {
			if(subdir != "")
				make_home_dir(subdir);
		}
	};

	{
		statement st(db, "select password from users where username = ?1");
		st.bind(1, username);
		if(st.step()) {
			// Found the username. Does their password match?
			std::string right_pw = st.getColumnText(0);
			if(right_pw == password) {
				make_home_dir_and_subdirs();
				return PAM_SUCCESS;
			} else {
				return PAM_AUTH_ERR;
			}

		} else if(st.ok) {
			// New user! (Query succeeded, but didn't find any row)
			statement st2(db, "insert into users (username, password) values (?1, ?2)");
			st2.bind(1, username);
			st2.bind(2, password);
			while(st2.step()) {}

			if(st2.ok && db.commit()) {
				make_home_dir_and_subdirs();
				return PAM_SUCCESS;
			}
		}
	}

	return PAM_AUTHINFO_UNAVAIL;
}

int pam_sm_setcred(pam_handle_t *pamh, int flags, int argc, const char **argv) {
	return PAM_SUCCESS;
}

#if 0

int pam_sm_acct_mgmt(pam_handle_t *pamh, int flags, int argc, const char **argv) {

}

int pam_sm_open_session(pam_handle_t *pamh, int flags, int argc, const char **argv) {

}

int pam_sm_close_session(pam_handle_t *pamh, int flags, int argc, const char **argv) {

}
#endif