/*
 * Copyright (C) 2016 Guido Berhoerster <guido+pwm@berhoerster.name>
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include "compat.h"

#ifdef	HAVE_ERR_H
#include <err.h>
#endif /* HAVE_ERR_H */
#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <sys/stat.h>
#ifdef	HAVE_SYS_TREE_H
#include <sys/tree.h>
#endif
#include <unistd.h>

#include "pwfile.h"
#include "util.h"

struct record_id_entry {
	RB_ENTRY(record_id_entry) record_id_entry;
	unsigned int	id;
	unsigned char	uuid[PWS3_UUID_SIZE];
};

RB_HEAD(record_id_tree, record_id_entry);

static int	record_id_entry_cmp(struct record_id_entry *,
    struct record_id_entry *);
RB_PROTOTYPE_STATIC(record_id_tree, record_id_entry, record_id_entry,
    record_id_entry_cmp)
RB_GENERATE_STATIC(record_id_tree, record_id_entry, record_id_entry,
    record_id_entry_cmp)

static int
record_id_entry_cmp(struct record_id_entry *entry1,
    struct record_id_entry *entry2)
{
	if (entry1->id > entry2->id) {
		return (-1);
	} else if (entry1->id < entry2->id) {
		return (1);
	}
	return (0);
}

static int
pws_record_cmp(const void *p1, const void *p2)
{
	int		retval;
	struct pws3_record *record1 = *(struct pws3_record **)p1;
	struct pws3_record *record2 = *(struct pws3_record **)p2;
	struct pws3_field *group_field1;
	const char	*group1;
	struct pws3_field *group_field2;
	const char	*group2;
	struct pws3_field *title_field1;
	const char	*title1;
	struct pws3_field *title_field2;
	const char	*title2;

	group_field1 = pws3_record_get_field(record1, PWS3_RECORD_FIELD_GROUP);
	group1 = (group_field1 != NULL) ?  pws3_field_get_text(group_field1) :
	    "";
	group_field2 = pws3_record_get_field(record2, PWS3_RECORD_FIELD_GROUP);
	group2 = (group_field2 != NULL) ?  pws3_field_get_text(group_field2) :
	    "";
	retval = strcmp(group1, group2);
	if (retval != 0) {
		return (retval);
	}

	title_field1 = pws3_record_get_field(record1, PWS3_RECORD_FIELD_TITLE);
	title1 = (title_field1 != NULL) ?  pws3_field_get_text(title_field1) :
	    "";
	title_field2 = pws3_record_get_field(record2, PWS3_RECORD_FIELD_TITLE);
	title2 = (title_field2 != NULL) ?  pws3_field_get_text(title_field2) :
	    "";
	return (strcmp(title1, title2));
}

static void
record_id_tree_clear(struct record_id_tree *tree)
{
	struct record_id_entry *entry;
	struct record_id_entry *entry_tmp;

	RB_FOREACH_SAFE(entry, record_id_tree, tree, entry_tmp) {
		RB_REMOVE(record_id_tree, tree, entry);
		free(entry);
	}
}

static void
record_id_tree_destroy(struct record_id_tree *tree)
{
	if (tree == NULL) {
		return;
	}

	record_id_tree_clear(tree);
	free(tree);
}

static const unsigned char *
record_id_tree_get_uuid(struct record_id_tree *tree, unsigned int id)
{
	struct record_id_entry *entry;

	entry = RB_FIND(record_id_tree, tree,
	    &(struct record_id_entry){ .id = id });
	if (entry == NULL) {
		return (NULL);
	}

	return (entry->uuid);
}

void
pwfile_init(struct pwm_ctx *ctx)
{
	ctx->file = pws3_file_create();
	if (ctx->file == NULL) {
		err(1, "pws3_file_create");
	}

	ctx->next_id = 1;

	ctx->record_id_tree = xmalloc(sizeof (struct record_id_tree));
	RB_INIT(ctx->record_id_tree);
}

void
pwfile_destroy(struct pwm_ctx *ctx)
{
	record_id_tree_destroy(ctx->record_id_tree);
	ctx->record_id_tree = NULL;
	pws3_file_destroy(ctx->file);
	ctx->file = NULL;
}

int
pwfile_read_file(struct pwm_ctx *ctx, FILE *fp)
{
	struct pws3_record *pws3_record;
	size_t		record_list_len = 0;
	struct pws3_record **pws3_record_list;
	size_t		i;
	struct pws3_field *uuid_field;
	const unsigned char *uuid;
	struct record_id_entry *entry;

	if (pws3_file_read_stream(ctx->file, ctx->password, fp) != 0) {
		warnx("failed to read password database: %s",
		    pws3_file_get_error_message(ctx->file));
		return (-1);
	}

	record_id_tree_clear(ctx->record_id_tree);

	/* sort records by group and title */
	for (pws3_record = pws3_file_first_record(ctx->file);
	    pws3_record != NULL;
	    pws3_record = pws3_file_next_record(ctx->file, pws3_record)) {
		record_list_len++;
	}
	pws3_record_list = xmalloc(sizeof (struct pws3_record *) *
	    record_list_len);
	for (pws3_record = pws3_file_first_record(ctx->file), i = 0;
	    pws3_record != NULL;
	    pws3_record = pws3_file_next_record(ctx->file, pws3_record)) {
		pws3_record_list[i++] = pws3_record;
	}
	qsort(pws3_record_list, record_list_len,
	    sizeof (struct pws3_record *), pws_record_cmp);

	/* build the tree of record IDs */
	for (i = 0; i < record_list_len; i++) {
		uuid_field = pws3_record_get_field(pws3_record_list[i],
		    PWS3_RECORD_FIELD_UUID);
		uuid = pws3_field_get_uuid(uuid_field);

		entry = xmalloc(sizeof (struct record_id_entry));
		entry->id = ctx->next_id++;
		memcpy(entry->uuid, uuid, sizeof (entry->uuid));

		RB_INSERT(record_id_tree, ctx->record_id_tree, entry);
	}

	free(pws3_record_list);

	return (0);
}

static int
make_backup_copy(const char *filename)
{
	int		retval = -1;
	FILE		*fp_orig = NULL;
	char		*backup_filename = NULL;
	char		*tmpfilename = NULL;
	mode_t		old_mode;
	int		fd_backup = -1;
	unsigned char	buf[BUFSIZ];
	size_t		read_len;
	FILE		*fp_backup = NULL;

	fp_orig = fopen(filename, "r");
	if (fp_orig == NULL) {
		if (errno != ENOENT) {
			warn("fopen");
			return (-1);
		}
		return (0);
	}

	xasprintf(&backup_filename, "%s~", filename);
	xasprintf(&tmpfilename, "%s.XXXXXX", filename);

	/* create temporary file */
	old_mode = umask(S_IRWXG | S_IRWXO);
	fd_backup = mkstemp(tmpfilename);
	umask(old_mode);
	if (fd_backup == -1) {
		warn("mkstemp");
		goto out;
	}
	fp_backup = fdopen(fd_backup, "w");
	if (fp_backup == NULL) {
		warn("fdopen");
		goto out;
	}

	/* copy file contents */
	while (!feof(fp_orig)) {
		read_len = fread(buf, 1, sizeof (buf), fp_orig);
		if ((read_len < sizeof (buf)) && ferror(fp_orig)) {
			warn("fread");
			goto out;
		}
		if (fwrite(buf, 1, read_len, fp_backup) != read_len) {
			warn("fwrite");
			goto out;
		}
	}
	if (fflush(fp_backup) != 0) {
		warn("fflush");
		goto out;
	}
	if (fsync(fileno(fp_backup)) != 0) {
		warn("fsync");
		goto out;
	}

	retval = 0;

out:
	if ((fd_backup != -1) && (fp_backup == NULL)) {
		close(fd_backup);
	}
	if (fp_backup != NULL) {
		fclose(fp_backup);
	}
	if (fp_orig != NULL) {
		fclose(fp_orig);
	}
	if (retval == 0) {
		/* rename temporary file and overwrite existing file */
		if (rename(tmpfilename, backup_filename) != 0) {
			warn("rename");
			retval = -1;
		}
	}
	if ((retval != 0) && ((fd_backup != -1) || (fp_backup != NULL))) {
		unlink(tmpfilename);
	}
	free(tmpfilename);
	free(backup_filename);

	return (retval);
}

int
pwfile_write_file(struct pwm_ctx *ctx)
{
	int	retval = -1;
	char	*tmpfilename = NULL;
	mode_t	old_mode;
	int	fd = -1;
	FILE	*fp = NULL;

	if (make_backup_copy(ctx->filename) != 0) {
		goto out;
	}

	xasprintf(&tmpfilename, "%s.XXXXXX", ctx->filename);

	/* create temporary file */
	old_mode = umask(S_IRWXG | S_IRWXO);
	fd = mkstemp(tmpfilename);
	if (fd == -1) {
		warn("mkstemp");
		goto out;
	}
	umask(old_mode);
	fp = fdopen(fd, "w");
	if (fp == NULL) {
		warn("fdopen");
		goto out;
	}

	/* write contents */
	if (pws3_file_write_stream(ctx->file, ctx->password, 10000, fp) != 0) {
		warnx("pws3_file_write_stream: %s",
		    pws3_file_get_error_message(ctx->file));
		goto out;
	}
	if (fflush(fp) != 0) {
		warn("fflush");
		goto out;
	}
	if (fsync(fileno(fp)) != 0) {
		warn("fsync");
		goto out;
	}

	retval = 0;

out:
	if ((fd != -1) && (fp == NULL)) {
		close(fd);
	}
	if (fp != NULL) {
		fclose(fp);
	}
	if (retval == 0) {
		/* rename temporary file and overwrite existing file */
		if (rename(tmpfilename, ctx->filename) != 0) {
			warn("rename");
			retval = -1;
		}
	}
	if ((retval != 0) && ((fd != -1) || (fp != NULL))) {
		unlink(tmpfilename);
	}
	free(tmpfilename);

	return (retval);
}

static int
list_item_cmp(const void *p1, const void *p2)
{
	int		retval;
	const union list_item *item1 = *(const union list_item **)p1;
	const union list_item *item2 = *(const union list_item **)p2;
	const char	*group1;
	const char	*group2;
	const char	*title1;
	const char	*title2;

	/* sort both groups and records first by group name */
	group1 = (item1->any.group != NULL) ? item1->any.group : "";
	group2 = (item2->any.group != NULL) ? item2->any.group : "";
	retval = strcmp(group1, group2);
	if ((retval != 0) || ((item1->any.type == ITEM_TYPE_GROUP) &&
	    (item2->any.type == ITEM_TYPE_GROUP))) {
		return (retval);
	} else if ((item1->any.type == ITEM_TYPE_GROUP) &&
	    (item2->any.type == ITEM_TYPE_RECORD)) {
		/* groups come before records belonging to it */
		return (-1);
	} else if ((item1->any.type == ITEM_TYPE_RECORD) &&
	    (item2->any.type == ITEM_TYPE_GROUP)) {
		return (1);
	}

	/* sort records also by title */
	title1 = (item1->record.title != NULL) ?  item1->record.title : "";
	title2 = (item2->record.title != NULL) ?  item2->record.title : "";
	return (strcmp(title1, title2));
}

union list_item **
pwfile_create_list(struct pwm_ctx *ctx)
{
	union list_item	**list;
	struct record_id_entry *entry;
	size_t		list_capacity = 0;
	struct pws3_field *empty_group_field;
	struct pws3_record *pws3_record;
	struct pws3_field *group_field;
	const char	*group;
	struct pws3_field *title_field;
	const char	*title;
	union list_item	*item;
	size_t		list_len = 0;
	size_t		i;
	size_t		j;
	const char	*prev_group = NULL;

	RB_FOREACH(entry, record_id_tree, ctx->record_id_tree) {
		list_capacity++;
	}
	list_capacity *= 2; /* maximum number of group items */
	for (empty_group_field = pws3_file_first_empty_group(ctx->file);
	    empty_group_field != NULL;
	    empty_group_field = pws3_file_next_empty_group(ctx->file,
	    empty_group_field)) {
		list_capacity++;
	}
	list_capacity++; /* terminating NULL */
	list = xmalloc(sizeof (union list_item *) * list_capacity);

	RB_FOREACH(entry, record_id_tree, ctx->record_id_tree) {
		pws3_record = pws3_file_get_record(ctx->file, entry->uuid);
		group_field = pws3_record_get_field(pws3_record,
		    PWS3_RECORD_FIELD_GROUP);
		group = (group_field != NULL) ?
		    pws3_field_get_text(group_field) : NULL;
		title_field = pws3_record_get_field(pws3_record,
		    PWS3_RECORD_FIELD_TITLE);
		title = (title_field != NULL) ?
		    pws3_field_get_text(title_field) : NULL;

		item = xmalloc(sizeof (union list_item));
		item->record.type = ITEM_TYPE_RECORD;
		item->record.group = (group != NULL) ? xstrdup(group) : NULL;
		item->record.title = (title != NULL) ? xstrdup(title) : NULL;
		item->record.id = entry->id;
		memcpy(item->record.uuid, entry->uuid,
		    sizeof (item->record.uuid));

		list[list_len++] = item;
	}
	/* sort records by group and title in order to find unqiue groups */
	qsort(list, list_len, sizeof (union list_item *), list_item_cmp);

	/* add groups based on the sorted records */
	for (i = 0, j = list_len; i < list_len; i++) {
		group = list[i]->record.group;
		if ((group != NULL) && ((prev_group == NULL) ||
		    (strcmp(prev_group, group) != 0))) {
			item = xmalloc(sizeof (union list_item));
			item->group.type = ITEM_TYPE_GROUP;
			item->group.group = (group != NULL) ? xstrdup(group) :
			    NULL;
			list[j++] = item;
			prev_group = group;
		}
	}
	list_len = j;

	/* add empty groups to the list */
	for (empty_group_field = pws3_file_first_empty_group(ctx->file);
	    empty_group_field != NULL;
	    empty_group_field = pws3_file_next_empty_group(ctx->file,
	    empty_group_field)) {
		group = pws3_field_get_text(empty_group_field);

		item = xmalloc(sizeof (union list_item));
		item->group.type = ITEM_TYPE_GROUP;
		item->group.group = xstrdup(group);

		list[list_len++] = item;
	}

	/* terminate the list */
	list[list_len] = NULL;
	list = xrealloc(list, sizeof (union list_item *) * (list_len + 1));

	/* sort the final list by group and title */
	qsort(list, list_len, sizeof (union list_item *), list_item_cmp);

	return (list);
}

void
pwfile_destroy_list(union list_item **list)
{
	size_t	i;

	if (list == NULL) {
		return;
	}

	for (i = 0; list[i] != NULL; i++) {
		if (list[i]->any.type == ITEM_TYPE_RECORD) {
			free(list[i]->record.title);
		}
		free(list[i]->any.group);
		free(list[i]);
	}

	free(list);
}

static void
update_record(struct pws3_record *pws3_record, struct record *record)
{
	struct pws3_field *title_field;
	struct pws3_field *group_field;
	struct pws3_field *username_field;
	struct pws3_field *password_field;
	struct pws3_field *notes_field;
	struct pws3_field *url_field;

	if (record->title != NULL) {
		title_field = pws3_field_create(0, PWS3_RECORD_FIELD_TITLE);
		if (title_field == NULL) {
			err(1, "pws3_record_field_create");
		}
		if (pws3_field_set_text(title_field,
		    record->title) != 0) {
			err(1, "pws3_field_set_text");
		}
		pws3_record_set_field(pws3_record, title_field);
	}
	if (record->group != NULL) {
		group_field = pws3_field_create(0, PWS3_RECORD_FIELD_GROUP);
		if (group_field == NULL) {
			err(1, "pws3_record_field_create");
		}
		if (pws3_field_set_text(group_field,
		    record->group) != 0) {
			err(1, "pws3_field_set_text");
		}
		pws3_record_set_field(pws3_record, group_field);
	}
	if (record->username != NULL) {
		username_field = pws3_field_create(0,
		    PWS3_RECORD_FIELD_USERNAME);
		if (username_field == NULL) {
			err(1, "pws3_record_field_create");
		}
		if (pws3_field_set_text(username_field,
		    record->username) != 0) {
			err(1, "pws3_field_set_text");
		}
		pws3_record_set_field(pws3_record, username_field);
	}
	if (record->password != NULL) {
		password_field = pws3_field_create(0,
		    PWS3_RECORD_FIELD_PASSWORD);
		if (password_field == NULL) {
			err(1, "pws3_record_field_create");
		}
		if (pws3_field_set_text(password_field,
		    record->password) != 0) {
			err(1, "pws3_field_set_text");
		}
		pws3_record_set_field(pws3_record, password_field);
	}
	if (record->notes != NULL) {
		notes_field = pws3_field_create(0, PWS3_RECORD_FIELD_NOTES);
		if (notes_field == NULL) {
			err(1, "pws3_record_field_create");
		}
		if (pws3_field_set_text(notes_field, record->notes) != 0) {
			err(1, "pws3_field_set_text");
		}
		pws3_record_set_field(pws3_record, notes_field);
	}
	if (record->url != NULL) {
		url_field = pws3_field_create(0, PWS3_RECORD_FIELD_URL);
		if (url_field == NULL) {
			err(1, "pws3_record_field_create");
		}
		if (pws3_field_set_text(url_field, record->url) != 0) {
			err(1, "pws3_field_set_text");
		}
		pws3_record_set_field(pws3_record, url_field);
	}
}

int
pwfile_create_record(struct pwm_ctx *ctx, struct record *record)
{
	struct pws3_record *pws3_record;
	const unsigned char *uuid;
	struct record_id_entry *entry;

	pws3_record = pws3_record_create();
	if (pws3_record == NULL) {
		err(1, "pws3_record_create");
	}
	update_record(pws3_record, record);
	pws3_file_insert_record(ctx->file, pws3_record);

	uuid = pws3_field_get_uuid(pws3_record_get_field(pws3_record,
	    PWS3_RECORD_FIELD_UUID));
	entry = xmalloc(sizeof (struct record_id_entry));
	entry->id = ctx->next_id++;
	memcpy(entry->uuid, uuid, sizeof (entry->uuid));
	RB_INSERT(record_id_tree, ctx->record_id_tree, entry);

	return (0);
}

int
pwfile_modify_record(struct pwm_ctx *ctx, unsigned int id,
    struct record *record)
{
	const unsigned char *uuid;

	uuid = record_id_tree_get_uuid(ctx->record_id_tree, id);
	if (uuid == NULL) {
		return (-1);
	}

	update_record(pws3_file_get_record(ctx->file, uuid), record);

	return (0);
}

int
pwfile_remove_record(struct pwm_ctx *ctx, unsigned int id)
{
	const unsigned char *uuid;
	struct record_id_entry *entry;

	uuid = record_id_tree_get_uuid(ctx->record_id_tree, id);
	if (uuid == NULL) {
		return (-1);
	}

	pws3_record_destroy(pws3_file_remove_record(ctx->file, uuid));

	entry = RB_FIND(record_id_tree, ctx->record_id_tree,
	    &(struct record_id_entry){ .id = id });
	free(RB_REMOVE(record_id_tree, ctx->record_id_tree, entry));

	return (0);
}

int
pwfile_create_group(struct pwm_ctx *ctx, const char *group)
{
	struct pws3_record *pws3_record;
	struct pws3_field *group_field;
	struct pws3_field *empty_group_field;

	/* check for a record in the given group */
	for (pws3_record = pws3_file_first_record(ctx->file);
	    pws3_record != NULL; pws3_record = pws3_file_next_record(ctx->file,
	    pws3_record)) {
		group_field = pws3_record_get_field(pws3_record,
		    PWS3_RECORD_FIELD_GROUP);
		if ((group_field != NULL) &&
		    (strcmp(group, pws3_field_get_text(group_field)) == 0)) {
			return (-1);
		}
	}

	empty_group_field = pws3_field_create(1,
	    PWS3_HEADER_FIELD_EMPTY_GROUPS);
	if (empty_group_field == NULL) {
		err(1, "pws3_field_create");
	}
	if (pws3_field_set_text(empty_group_field, group) != 0) {
		err(1, "pws3_field_set_text");
	}
	pws3_file_insert_empty_group(ctx->file, empty_group_field);

	return (0);
}

int
pwfile_remove_group(struct pwm_ctx *ctx, const char *group)
{
	struct pws3_field *empty_group_field;

	empty_group_field = pws3_file_remove_empty_group(ctx->file, group);
	if (empty_group_field == NULL) {
		return (-1);
	}
	pws3_field_destroy(empty_group_field);

	return (0);
}

struct record *
pwfile_get_record(struct pwm_ctx *ctx, unsigned int id)
{
	struct record	*record;
	const unsigned char *uuid;
	struct pws3_record *pws3_record;
	struct pws3_field *title_field;
	struct pws3_field *group_field;
	struct pws3_field *username_field;
	struct pws3_field *password_field;
	struct pws3_field *notes_field;
	struct pws3_field *url_field;

	uuid = record_id_tree_get_uuid(ctx->record_id_tree, id);
	if (uuid == NULL) {
		return (NULL);
	}
	pws3_record = pws3_file_get_record(ctx->file, uuid);

	record = xmalloc(sizeof (struct record));

	title_field = pws3_record_get_field(pws3_record,
	    PWS3_RECORD_FIELD_TITLE);
	record->title = (title_field != NULL) ?
	    xstrdup(pws3_field_get_text(title_field)) : NULL;

	group_field = pws3_record_get_field(pws3_record,
	    PWS3_RECORD_FIELD_GROUP);
	record->group = (group_field != NULL) ?
	    xstrdup(pws3_field_get_text(group_field)) : NULL;

	username_field = pws3_record_get_field(pws3_record,
	    PWS3_RECORD_FIELD_USERNAME);
	record->username = (username_field != NULL) ?
	    xstrdup(pws3_field_get_text(username_field)) : NULL;

	password_field = pws3_record_get_field(pws3_record,
	    PWS3_RECORD_FIELD_PASSWORD);
	record->password = (password_field != NULL) ?
	    xstrdup(pws3_field_get_text(password_field)) : NULL;

	notes_field = pws3_record_get_field(pws3_record,
	    PWS3_RECORD_FIELD_NOTES);
	record->notes = (notes_field != NULL) ?
	    xstrdup(pws3_field_get_text(notes_field)) : NULL;

	url_field = pws3_record_get_field(pws3_record, PWS3_RECORD_FIELD_URL);
	record->url = (url_field != NULL) ?
	    xstrdup(pws3_field_get_text(url_field)) : NULL;

	return (record);
}

void
pwfile_destroy_record(struct record *record)
{
	if (record == NULL) {
		return;
	}

	free(record->title);
	free(record->group);
	free(record->username);
	free(record->password);
	free(record->notes);
	free(record->url);
	free(record);
}