/*-
 * Copyright (c) 2015 Taylor R. Campbell
 * All rights reserved.
 *
 * 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 AUTHOR 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 AUTHOR 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.
 */

#define	_POSIX_C_SOURCE	200809L

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "keccak.h"

/*
 * XXX Yes, sha3.c, not sha3.h.  We want the internal sha3_* routines,
 * not the OpenSSL-style wrappers, in order to more easily handle
 * multiple possible sizes.
 */
#include "sha3.c"

static const char *progname;

void
setprogname(const char *argv0)
{
	const char *p;

	p = strrchr(argv0, '/');
	if (p != NULL)
		progname = p + 1;
	else
		progname = argv0;
}

const char *
getprogname(void)
{

	return progname;
}

static void
vwarnx(const char *fmt, va_list va)
{

	(void)fprintf(stderr, "%s: ", getprogname());
	(void)vfprintf(stderr, fmt, va);
	(void)fprintf(stderr, "\n");
}

static void
warnx(const char *fmt, ...)
{
	va_list va;

	va_start(va, fmt);
	vwarnx(fmt, va);
	va_end(va);
}

static void
vwarn(const char *fmt, va_list va)
{
	int error = errno;

	(void)fprintf(stderr, "%s: ", getprogname());
	(void)vfprintf(stderr, fmt, va);
	(void)fprintf(stderr, ": %s\n", strerror(error));
}

static void
warn(const char *fmt, ...)
{
	va_list va;

	va_start(va, fmt);
	vwarn(fmt, va);
	va_end(va);
}

static void
errx(int code, const char *fmt, ...)
{
	va_list va;

	va_start(va, fmt);
	vwarnx(fmt, va);
	va_end(va);

	exit(code);
}

static void
err(int code, const char *fmt, ...)
{
	va_list va;

	va_start(va, fmt);
	vwarn(fmt, va);
	va_end(va);

	exit(code);
}

static void
usage(void)
{

	(void)fprintf(stderr, "Usage: %s -s <size> [-x <cap>] [<file>...]\n",
	    getprogname());
	exit(EXIT_FAILURE);
}

static ssize_t
read_block(int fd, void *buf, size_t len)
{
	ssize_t nread;

	while ((nread = read(fd, buf, len)) == -1) {
		switch (errno) {
		case EINTR:
			continue;
		case EAGAIN: {
			struct pollfd pfd = { .fd = fd, .events = POLLIN };

			(void)poll(&pfd, 1, -1);
			continue;
		}
		default:
			return -1;
		}
	}

	return nread;
}

static void
doit(size_t cap, size_t size, int fd, const char *fname)
{
	const size_t rate = sha3_rate(cap == 0? size/8 : cap/8);
	static uint8_t buf[0x10000];
	static uint8_t h[0x1000];
	static struct sha3 sha3;
	ssize_t nread;
	const uint8_t *p = h;
	size_t i;

	sha3_init(&sha3, rate);
	while ((nread = read_block(fd, buf, sizeof buf)) != 0) {
		if (nread == -1)
			err(EXIT_FAILURE, "read");
		sha3_update(&sha3, buf, (size_t)nread, rate);
	}
	(*(cap == 0? sha3_final : shake_final))(h, size/8, &sha3, rate);

	if (fname != NULL) {
		if (cap == 0) {
			if (printf("SHA3-%zu", size) < 0)
				err(EXIT_FAILURE, "printf");
		} else {
			if (printf("SHAKE%zu-%zu", cap, size) < 0)
				err(EXIT_FAILURE, "printf");
		}
		if (printf(" (%s) = ", fname) < 0)
			err(EXIT_FAILURE, "printf");
	}
	for (i = 0; i < size/8; i++) {
		if (printf("%02hhx", *p++) < 0)
			err(EXIT_FAILURE, "printf");
	}
	if (printf("\n") < 0)
		err(EXIT_FAILURE, "printf");
}

static void
ignore_signal(int signo)
{

	(void)signo;
}

int
main(int argc, char **argv)
{
	extern char *optarg;
	extern int optind;
	size_t cap = 0;
	size_t size = 0;
	int ch;

	setprogname(argv[0]);
	if (signal(SIGPIPE, &ignore_signal) == SIG_ERR)
		err(EXIT_FAILURE, "ignore SIGPIPE");

	if (SHA3_Selftest() == -1)
		errx(EXIT_FAILURE, "SHA-3 self-test failed");

	while ((ch = getopt(argc, argv, "+s:x:")) != -1) {
		switch (ch) {
		case 's': {
			char *end;
			unsigned long s;

			if (size != 0) {
				warnx("specify only one `-s <size>'");
				usage();
			}

			errno = 0;
			s = strtoul(optarg, &end, 0);
			if ((errno == ERANGE && s == ULONG_MAX) ||
			    (s > SIZE_MAX)) {
				warnx("invalid output size: %s", optarg);
				usage();
			}
			size = s;
			break;
		}
		case 'x': {
			char *end;
			unsigned long c;

			if (cap != 0) {
				warnx("specify only one `-x <cap>'");
				usage();
			}

			errno = 0;
			c = strtoul(optarg, &end, 0);
			if ((errno == ERANGE && c == ULONG_MAX) ||
			    (c > SIZE_MAX) ||
			    (c == 0)) {
				warnx("invalid XOF capacity: %s", optarg);
				usage();
			}
			cap = c;
			break;
		}
		case '?':
		case '+':	/* WTF GNU */
			usage();
		}
	}

	argc -= optind;
	argv += optind;

	if (size == 0) {
		warnx("missing `-s <size>'");
		usage();
	}
	if (cap == 0) {
		if (size != 224 && size != 256 && size != 384 && size != 512) {
			warnx("invalid hash size %zu"
			    " -- must be 224, 256, 384, or 512",
			    size);
			usage();
		}
	} else {
		if (cap != 128 && cap != 256) {
			warnx("invalid XOF capacity %zu -- must be 128 or 256",
			    cap);
			usage();
		}
		if ((size % 8) != 0) {
			warnx("invalid XOF size %zu -- must be divisible by 8",
			    size);
			usage();
		}
		if (size > 65536) {
			warnx("invalid XOF size %zu -- must be <=65536", size);
			usage();
		}
	}

	if (argc == 0) {
		doit(cap, size, STDIN_FILENO, NULL);
		return EXIT_SUCCESS;
	}

	while (argc--) {
		const char *fname = *argv++;
		int fd;

		if (strcmp(fname, "-") == 0) {
			fd = STDIN_FILENO;
		} else {
			fd = open(fname, O_RDONLY);
			if (fd == -1) {
				warn("%s", fname);
				continue;
			}
		}
		doit(cap, size, fd, fname);
		(void)close(fd);
	}

	return EXIT_SUCCESS;
}