/*	$NetBSD: quotarestore.c,v 1.3 2012/09/13 21:44:50 joerg Exp $	*/
/*-
 * Copyright (c) 2012 The NetBSD Foundation, Inc.
 * All rights reserved.
 *
 * This code is derived from software contributed to The NetBSD Foundation
 * by David A. Holland.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/cdefs.h>
__RCSID("$NetBSD: quotarestore.c,v 1.3 2012/09/13 21:44:50 joerg Exp $");

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <getopt.h>
#include <limits.h>
#include <errno.h>
#include <err.h>

#include <quota.h>

static const char ws[] = " \t\r\n";

static char **idtypenames;
static unsigned numidtypes;

static char **objtypenames;
static unsigned numobjtypes;

////////////////////////////////////////////////////////////
// table of quota keys

struct qklist {
	struct quotakey *keys;
	unsigned num, max;
};

static struct qklist *
qklist_create(void)
{
	struct qklist *l;

	l = malloc(sizeof(*l));
	if (l == NULL) {
		err(EXIT_FAILURE, "malloc");
	}
	l->keys = 0;
	l->num = 0;
	l->max = 0;
	return l;
}

static void
qklist_destroy(struct qklist *l)
{
	free(l->keys);
	free(l);
}

static void
qklist_truncate(struct qklist *l)
{
	l->num = 0;
}

static void
qklist_add(struct qklist *l, const struct quotakey *qk)
{
	assert(l->num <= l->max);
	if (l->num == l->max) {
		l->max = l->max ? l->max * 2 : 4;
		l->keys = realloc(l->keys, l->max * sizeof(l->keys[0]));
		if (l->keys == NULL) {
			err(EXIT_FAILURE, "realloc");
		}
	}
	l->keys[l->num++] = *qk;
}

static int
qk_compare(const void *av, const void *bv)
{
	const struct quotakey *a = av;
	const struct quotakey *b = bv;

	if (a->qk_idtype < b->qk_idtype) {
		return -1;
	}
	if (a->qk_idtype > b->qk_idtype) {
		return 1;
	}

	if (a->qk_id < b->qk_id) {
		return -1;
	}
	if (a->qk_id > b->qk_id) {
		return 1;
	}

	if (a->qk_objtype < b->qk_objtype) {
		return -1;
	}
	if (a->qk_objtype > b->qk_objtype) {
		return 1;
	}

	return 0;
}

static void
qklist_sort(struct qklist *l)
{
	qsort(l->keys, l->num, sizeof(l->keys[0]), qk_compare);
}

static int
qklist_present(struct qklist *l, const struct quotakey *key)
{
	void *p;

	p = bsearch(key, l->keys, l->num, sizeof(l->keys[0]), qk_compare);
	return p != NULL;
}

////////////////////////////////////////////////////////////
// name tables and string conversion

static void
maketables(struct quotahandle *qh)
{
	unsigned i;

	numidtypes = quota_getnumidtypes(qh);
	idtypenames = malloc(numidtypes * sizeof(idtypenames[0]));
	if (idtypenames == NULL) {
		err(EXIT_FAILURE, "malloc");
	}

	for (i=0; i<numidtypes; i++) {
		idtypenames[i] = strdup(quota_idtype_getname(qh, i));
		if (idtypenames[i] == NULL) {
			err(EXIT_FAILURE, "strdup");
		}
	}

	numobjtypes = quota_getnumobjtypes(qh);
	objtypenames = malloc(numobjtypes * sizeof(objtypenames[0]));
	if (objtypenames == NULL) {
		err(EXIT_FAILURE, "malloc");
	}

	for (i=0; i<numobjtypes; i++) {
		objtypenames[i] = strdup(quota_objtype_getname(qh, i));
		if (objtypenames[i] == NULL) {
			err(EXIT_FAILURE, "strdup");
		}
	}
}

static int
getidtype(const char *name, int *ret)
{
	unsigned i;

	for (i=0; i<numidtypes; i++) {
		if (!strcmp(name, idtypenames[i])) {
			*ret = i;
			return 0;
		}
	}
	return -1;
}

static int
getid(const char *name, int idtype, id_t *ret)
{
	unsigned long val;
	char *s;

	if (!strcmp(name, "default")) {
		*ret = QUOTA_DEFAULTID;
		return 0;
	}
	errno = 0;
	val = strtoul(name, &s, 10);
	if (errno || *s != 0) {
		return -1;
	}
	if (idtype == QUOTA_IDTYPE_USER && val > UID_MAX) {
		return -1;
	}
	if (idtype == QUOTA_IDTYPE_GROUP && val > GID_MAX) {
		return -1;
	}
	*ret = val;
	return 0;
}

static int
getobjtype(const char *name, int *ret)
{
	unsigned i;
	size_t len;

	for (i=0; i<numobjtypes; i++) {
		if (!strcmp(name, objtypenames[i])) {
			*ret = i;
			return 0;
		}
	}

	/*
	 * Sigh. Some early committed versions of quotadump used
	 * "blocks" and "files" instead of "block" and "file".
	 */
	len = strlen(name);
	if (len == 0) {
		return -1;
	}
	for (i=0; i<numobjtypes; i++) {
		if (name[len-1] == 's' &&
		    !strncmp(name, objtypenames[i], len-1)) {
			*ret = i;
			return 0;
		}
	}
	return -1;
}

static int
getlimit(const char *name, uint64_t *ret)
{
	unsigned long long val;
	char *s;

	if (!strcmp(name, "-")) {
		*ret = QUOTA_NOLIMIT;
		return 0;
	}
	errno = 0;
	val = strtoull(name, &s, 10);
	if (errno || *s != 0) {
		return -1;
	}
	*ret = val;
	return 0;
}

static int
gettime(const char *name, int64_t *ret)
{
	long long val;
	char *s;

	if (!strcmp(name, "-")) {
		*ret = QUOTA_NOTIME;
		return 0;
	}
	errno = 0;
	val = strtoll(name, &s, 10);
	if (errno || *s != 0 || val < 0) {
		return -1;
	}
	*ret = val;
	return 0;
}

////////////////////////////////////////////////////////////
// parsing tools

static int
isws(int ch)
{
	return ch != '\0' && strchr(ws, ch) != NULL;
}

static char *
skipws(char *s)
{
	while (isws(*s)) {
		s++;
	}
	return s;
}

////////////////////////////////////////////////////////////
// deletion of extra records

static void
scankeys(struct quotahandle *qh, struct qklist *seenkeys,
	struct qklist *dropkeys)
{
	struct quotacursor *qc;
#define MAX 8
	struct quotakey keys[MAX];
	struct quotaval vals[MAX];
	int num, i;

	qc = quota_opencursor(qh);
	if (qc == NULL) {
		err(EXIT_FAILURE, "quota_opencursor");
	}

	while (quotacursor_atend(qc) == 0) {
		num = quotacursor_getn(qc, keys, vals, MAX);
		if (num < 0) {
			if (errno == EDEADLK) {
				quotacursor_rewind(qc);
				qklist_truncate(dropkeys);
				continue;
			}
			err(EXIT_FAILURE, "quotacursor_getn");
		}
		for (i=0; i<num; i++) {
			if (qklist_present(seenkeys, &keys[i]) == 0) {
				qklist_add(dropkeys, &keys[i]);
			}
		}
	}

	quotacursor_close(qc);
}

static void
purge(struct quotahandle *qh, struct qklist *dropkeys)
{
	unsigned i;

	for (i=0; i<dropkeys->num; i++) {
		if (quota_delete(qh, &dropkeys->keys[i])) {
			err(EXIT_FAILURE, "quota_delete");
		}
	}
}

////////////////////////////////////////////////////////////
// dumpfile reader

static void
readdumpfile(struct quotahandle *qh, FILE *f, const char *path,
	     struct qklist *seenkeys)
{
	char buf[128];
	unsigned lineno;
	unsigned long version;
	char *s;
	char *fields[8];
	unsigned num;
	char *x;
	struct quotakey key;
	struct quotaval val;
	int ch;

	lineno = 0;
	if (fgets(buf, sizeof(buf), f) == NULL) {
		errx(EXIT_FAILURE, "%s: EOF before quotadump header", path);
	}
	lineno++;
	if (strncmp(buf, "@format netbsd-quota-dump v", 27) != 0) {
		errx(EXIT_FAILURE, "%s: Missing quotadump header", path);
	}
	s = buf+27;
	errno = 0;
	version = strtoul(s, &s, 10);
	if (errno) {
		errx(EXIT_FAILURE, "%s: Corrupted quotadump header", path);
	}
	s = skipws(s);
	if (*s != '\0') {
		errx(EXIT_FAILURE, "%s: Trash after quotadump header", path);
	}

	switch (version) {
	    case 1: break;
	    default:
		errx(EXIT_FAILURE, "%s: Unsupported quotadump version %lu",
		     path, version);
	}

	while (fgets(buf, sizeof(buf), f)) {
		lineno++;
		if (buf[0] == '#') {
			continue;
		}
		if (!strncmp(buf, "@end", 4)) {
			s = skipws(buf+4);
			if (*s != '\0') {
				errx(EXIT_FAILURE, "%s:%u: Invalid @end tag",
				     path, lineno);
			}
			break;
		}

		num = 0;
		for (s = strtok_r(buf, ws, &x);
		     s != NULL;
		     s = strtok_r(NULL, ws, &x)) {
			if (num < 8) {
				fields[num++] = s;
			} else {
				errx(EXIT_FAILURE, "%s:%u: Too many fields",
				     path, lineno);
			}
		}
		if (num < 8) {
			errx(EXIT_FAILURE, "%s:%u: Not enough fields",
			     path, lineno);
		}

		if (getidtype(fields[0], &key.qk_idtype)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid/unknown ID type %s",
			     path, lineno, fields[0]);
		}
		if (getid(fields[1], key.qk_idtype, &key.qk_id)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid ID number %s",
			     path, lineno, fields[1]);
		}
		if (getobjtype(fields[2], &key.qk_objtype)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid/unknown object "
			     "type %s",
			     path, lineno, fields[2]);
		}

		if (getlimit(fields[3], &val.qv_hardlimit)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid hard limit %s",
			     path, lineno, fields[3]);
		}
		if (getlimit(fields[4], &val.qv_softlimit)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid soft limit %s",
			     path, lineno, fields[4]);
		}
		if (getlimit(fields[5], &val.qv_usage)) {
			/*
			 * Make this nonfatal as it'll be ignored by
			 * quota_put() anyway.
			 */
			warnx("%s:%u: Invalid current usage %s",
			     path, lineno, fields[5]);
			val.qv_usage = 0;
		}
		if (gettime(fields[6], &val.qv_expiretime)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid expire time %s",
			     path, lineno, fields[6]);
		}
		if (gettime(fields[7], &val.qv_grace)) {
			errx(EXIT_FAILURE, "%s:%u: Invalid grace period %s",
			     path, lineno, fields[7]);
		}

		if (quota_put(qh, &key, &val)) {
			err(EXIT_FAILURE, "%s:%u: quota_put", path, lineno);
		}

		if (seenkeys != NULL) {
			qklist_add(seenkeys, &key);
		}
	}
	if (feof(f)) {
		return;
	}
	if (ferror(f)) {
		errx(EXIT_FAILURE, "%s: Read error", path);
	}
	/* not at EOF, not an error... what's left? */
	while (1) {
		ch = fgetc(f);
		if (ch == EOF)
			break;
		if (isws(ch)) {
			continue;
		}
		warnx("%s:%u: Trash after @end tag", path, lineno);
	}
}

////////////////////////////////////////////////////////////
// top level control logic

__dead static void
usage(void)
{
	fprintf(stderr, "usage: %s [-d] volume [dump-file]\n",
		getprogname());
	exit(EXIT_FAILURE);
}

int
main(int argc, char *argv[])
{
	int ch;
	FILE *f;
	struct quotahandle *qh;

	int dflag = 0;
	const char *volume = NULL;
	const char *dumpfile = NULL;

	while ((ch = getopt(argc, argv, "d")) != -1) {
		switch (ch) {
		    case 'd': dflag = 1; break;
		    default: usage(); break;
		}
	}

	if (optind >= argc) {
		usage();
	}
	volume = argv[optind++];
	if (optind < argc) {
		dumpfile = argv[optind++];
	}
	if (optind < argc) {
		usage();
	}

	qh = quota_open(volume);
	if (qh == NULL) {
		err(EXIT_FAILURE, "quota_open: %s", volume);
	}
	if (dumpfile != NULL) {
		f = fopen(dumpfile, "r");
		if (f == NULL) {
			err(EXIT_FAILURE, "%s", dumpfile);
		}
	} else {
		f = stdin;
		dumpfile = "<stdin>";
	}

	maketables(qh);

	if (dflag) {
		struct qklist *seenkeys, *dropkeys;

		seenkeys = qklist_create();
		dropkeys = qklist_create();

		readdumpfile(qh, f, dumpfile, seenkeys);
		qklist_sort(seenkeys);
		scankeys(qh, seenkeys, dropkeys);
		purge(qh, dropkeys);

		qklist_destroy(dropkeys);
		qklist_destroy(seenkeys);
	} else {
		readdumpfile(qh, f, dumpfile, NULL);
	}

	if (f != stdin) {
		fclose(f);
	}
	quota_close(qh);
	return EXIT_SUCCESS;
}