generate.c

#include <sys/stat.h>
#include <sys/wait.h>

#include <err.h>
#include <fcntl.h>
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <strings.h>
#include <unistd.h>

#include "lootsack.h"

void wait_and_report(pid_t pid, const char *n);
__dead void usage(const char *progname);

struct treasure *generate_treasure(uint32_t n);
__dead void write_treasure(struct treasure *t, size_t n, char csv_fmt,
    const char *out);

int verbose = 0;
static struct option longopts[] = {
	{ "number", required_argument, NULL, 'n' },
	{ "verbose",      no_argument, NULL, 'v' },
	{ "out",    required_argument, NULL, 'o' },
	{ "fmt",    required_argument, NULL, 'f' },
};

__dead void
usage(const char *progname)
{
	fprintf(stderr, "Usage:\n\t%s [-v] [-n items] [-f fmt] -o outputfile\n",
	    progname);
	exit(1);
}

int
main(int argc, char **argv)
{
	pid_t pid;
	uint32_t num_items;
	const char *out;
	struct treasure *items;
	char ch;
	char csv_fmt = 0;

	out = NULL;
	num_items = 1400;

	while ((ch = getopt_long(argc, argv, "o:n:vf:", longopts, NULL)) != -1)
	{
		switch (ch) {
		case 'v':
			verbose = 1;
			break;
		case 'n':
			num_items = (uint32_t) atoll(optarg);
			break;
		case 'o':
			out = optarg;
			break;
		case 'f':
			if (strncasecmp("csv", optarg, 3) == 0) {
				csv_fmt = 1;
			} else if (strncasecmp("bin", optarg, 3) == 0) {
				/* use default */
			} else {
				fprintf(stderr, "Unknown format, \"%s\".\n",
				    optarg);
				usage(*argv); /* does not return */
			}
			break;
		default:
			usage(*argv); /* does not return */
		}
	}

	if (!out)
		usage(*argv); /* does not return */

	if (num_items > TREASURE_ITEM_MAX) {
		if (verbose)
			fprintf(stderr, "Truncating provided num_items from %u "
			    "to %u.\n", num_items, TREASURE_ITEM_MAX);
		num_items = TREASURE_ITEM_MAX;
	}

	if (pledge("stdio proc unveil wpath cpath", NULL) == -1)
		err(1, "pledge");

	items = generate_treasure(num_items);

	if (verbose)
		print_treasure(items, num_items > 10 ? 10 : num_items);

	if (verbose)
		fprintf(stderr, "Writing to \"%s\".\n", out);

	if ((pid = fork()) == 0)
		write_treasure(items, num_items, csv_fmt, out); /* no return */
	else
		wait_and_report(pid, "write_treasure");

	return 0;
}

struct treasure *
generate_treasure(uint32_t n)
{
	uint32_t i,j;
	uint32_t r;
	struct treasure *result;
	struct treasure_type_limits *lim;
	if ((result = calloc(n, sizeof(struct treasure))) == NULL)
		err(1, "calloc");
	for (i = 0; i < n; i++) {
		r = arc4random_uniform(1000) + 1;
		for (j = 0; j < NUM_TREASURE_TYPES; j++) {
			result[i].t = j;
			if (type_frequency[j] >= r)
				break;
		}
		lim = limits + result[i].t;
		result[i].w = lim->w_min + arc4random_uniform(
		    lim->w_max - lim->w_min);
		result[i].v = lim->v_min + arc4random_uniform(
		    lim->v_max - lim->v_min);
	}
	return result;
}


void
wait_and_report(pid_t pid, const char *n)
{
	int status;
	waitpid(pid, &status, 0);
	if (WIFEXITED(status) && WEXITSTATUS(status) != 0) {
		fprintf(stderr, "Child (%s) non-zero exit: %d\n",
		    n, WEXITSTATUS(status));
	} else if (WIFSIGNALED(status)) {
		fprintf(stderr, "Child (%s) exited by signal: %d\n",
		    n, WTERMSIG(status));
	} else if (WIFSTOPPED(status)) {
		fprintf(stderr, "Child (%s) stopped by signal: %d\n",
		    n, WSTOPSIG(status));
	}
}

__dead void
write_treasure(struct treasure *items, size_t n, char csv_fmt, const char *out)
{
	int fd;

	if (pledge("stdio unveil wpath cpath", NULL) == -1)
		err(1, "pledge");
	if (pledge(NULL, NULL) == -1)
		err(1, "finalize pledge");

	if (unveil(out, "wc") == -1)
		err(1, "unveil");
	if (unveil(NULL, NULL) == -1)
		err(1, "finalize unveil");

	if (verbose)
		fprintf(stderr, "Opening output file \"%s\".\n", out);

	if ((fd = open(out, O_WRONLY | O_CREAT, S_IRUSR | S_IWUSR | S_IRGRP))
	    == -1)
		err(1, "%s", out);
	if (csv_fmt) {
		/* write the csv format */
		write_treasure_csv(fd, items, n);
	} else {
		/* write the binary format */
		write_treasure_bin(fd, items, n);
	}
	exit(0);
}