/*
 * login_totp - TOTP-based login method
 *
 * Written in 2023 by Lucas
 *
 * To the extent possible under law, the author(s) have dedicated all
 * copyright and related and neighboring rights to this software to the
 * public domain worldwide. This software is distributed without any
 * warranty.
 *
 * You should have received a copy of the CC0 Public Domain Dedication
 * along with this software. If not, see
 * <http://creativecommons.org/publicdomain/zero/1.0/>.
 */

#include <inttypes.h>
#include <limits.h>
#include <paths.h>
#include <readpassphrase.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <syslog.h>
#include <time.h>
#include <unistd.h>

#include <login_cap.h>

#include <openssl/evp.h>
#include <openssl/hmac.h>


#define MODE_LOGIN	0
#define MODE_CHALLENGE	1
#define MODE_RESPONSE	2

#define TOTP_KEYLEN	20

#define _PATH_TOTP	_PATH_VARDB "totp/"


#define PUT64BE(x, o) do {						\
	(o)[0] = ((x) >> 56);						\
	(o)[1] = ((x) >> 48) & 0xff;					\
	(o)[2] = ((x) >> 40) & 0xff;					\
	(o)[3] = ((x) >> 32) & 0xff;					\
	(o)[4] = ((x) >> 24) & 0xff;					\
	(o)[5] = ((x) >> 16) & 0xff;					\
	(o)[6] = ((x) >> 8) & 0xff;					\
	(o)[7] = (x) & 0xff;						\
} while (0)

#define GET32BE(o, n) do {						\
	(n) = ((o)[0] << 24) | ((o)[1] << 16) | ((o)[2] << 8) |		\
	    ((o)[3]);							\
} while (0)


static int32_t	totp(const void *, int, int, const EVP_MD *, uint64_t,
    uint64_t);
static int32_t	totp_login(char *, char *);
static void	user_key(char *, unsigned char *);

static int32_t
totp(const void *key, int key_len, int digits, const EVP_MD *evp_md,
    uint64_t ts, uint64_t step)
{
	unsigned char md[EVP_MAX_MD_SIZE];
	uint8_t ctr[8];
	uint64_t counter;
	unsigned int md_len, offset;
	int32_t value, modulo;
	int i;

	/*
	 * RFC 4226 reads "The HOTP value must be at least a 6-digit value."
	 * Given that values are AND'd with 0x7fffffff, they're bound by
	 * 2^31 - 1, 2147483647 in base 10; more than 10 digits is impossible.
	 */
	if (digits < 6 || digits > 10)
		return -1;

	if (digits < 10) {
		modulo = 1000000;
		for (i = 0; i < digits - 6; i++)
			modulo *= 10;
	}

	counter = ts / step;
	PUT64BE(counter, ctr);
	if (HMAC(evp_md, key, key_len, ctr, sizeof(ctr), md, &md_len) == NULL)
		return -1;
	offset = md[md_len - 1] & 0xf;
	GET32BE(&md[offset], value);
	value &= 0x7fffffff;

	return digits == 10 ? value : value % modulo;
}

static int
totp_login(char *user, char *pass)
{
	unsigned char key[TOTP_KEYLEN];
	char buf[16];
	uint64_t ts, step;
	time_t now;
	int32_t r;
	int digits = 6, tries;

	now = time(NULL);
	if (now < 0) {
		syslog(LOG_WARNING, "time underflow");
		return 0;
	}
	if (now > INT64_MAX) {
		syslog(LOG_WARNING, "time overflow");
		return 0;
	}

	ts = (uint64_t)now;
	step = 30;
	tries = 3;

	if (ts >= step)
		ts -= step;
	else
		tries--;

	user_key(user, &key[0]);

	do {
		r = totp(key, sizeof(key), digits, EVP_sha1(), ts, step);
		if (r == -1)
			syslog(LOG_ERR, "TOTP error");
		else {
			snprintf(buf, sizeof(buf), "%0*" PRIi32, digits, r);
			if (strcmp(buf, pass) == 0)
				return 1;
		}
		if (ts > UINT64_MAX - step)
			break;
		ts += step;
	} while (--tries > 0);
	explicit_bzero(key, sizeof(key));

	return 0;
}

static void
user_key(char *user, unsigned char *key)
{
	FILE *f;
	char *path;
	size_t r;
	unsigned char dummy;

	if (asprintf(&path, _PATH_TOTP "%s.key", user) == -1) {
		syslog(LOG_ERR, "asprintf: %m");
		exit(1);
	}
	f = fopen(path, "r");
	if (f == NULL) {
		syslog(LOG_ERR, "Can't open keyfile for %s: %m", user);
		exit(1);
	}

	r = fread(key, sizeof(*key), TOTP_KEYLEN, f);
	if (r != TOTP_KEYLEN) {
		syslog(LOG_ERR, "short key");
		exit(1);
	}
	r = fread(&dummy, sizeof(dummy), 1, f);
	if (r != 0 || !feof(f)) {
		syslog(LOG_ERR, "long key");
		exit(1);
	}

	fclose(f);
	free(path);
}

int
main(int argc, char *argv[])
{
	FILE *back;
	char *curtime, *user, *pass;
	char pbuf[1024], rbuf[1024];
	time_t clock;
	int ch, mode;

	(void)signal(SIGQUIT, SIG_IGN);
	(void)signal(SIGINT, SIG_IGN);

	openlog("login_-totp", LOG_ODELAY, LOG_AUTH);

	back = NULL;
	mode = MODE_LOGIN;
	while ((ch = getopt(argc, argv, "ds:v:")) != -1) {
		switch (ch) {
		case 'd':
			back = stdout;
			break;
		case 's':
			if (strcmp(optarg, "login") == 0)
				mode = MODE_LOGIN;
			else if (strcmp(optarg, "challenge") == 0)
				mode = MODE_CHALLENGE;
			else if (strcmp(optarg, "response") == 0)
				mode = MODE_RESPONSE;
			else {
				syslog(LOG_ERR, "%s: invalid service", optarg);
				exit(1);
			}
			break;
		case 'v':
			break;
		default:
			syslog(LOG_ERR, "usage error");
			exit(1);
		}
	}
	argc -= optind;
	argv += optind;

	switch (argc) {
	case 2:
	case 1:
		user = argv[0];
		break;
	default:
		syslog(LOG_ERR, "usage error");
		exit(1);
	}

	if (pledge("stdio rpath tty", NULL) == -1) {
		syslog(LOG_ERR, "pledge: %m");
		exit(1);
	}

	if (back == NULL && (back = fdopen(3, "r+")) == NULL) {
		syslog(LOG_ERR, "reopening back channel: %m");
		exit(1);
	}

	if (mode == MODE_CHALLENGE) {
		fprintf(back, BI_VALUE " challenge totp: \n");
		fprintf(back, BI_CHALLENGE "\n");
		exit(0);
	}

	pass = NULL;
	if (mode == MODE_RESPONSE) {
		int i, npass;

		i = -1;
		npass = 0;
		while (++i < sizeof(rbuf) && read(3, &rbuf[i], 1) == 1) {
			if (rbuf[i] == '\0') {
				if (strcmp(rbuf, "totp: ") != 0) {
					syslog(LOG_ERR,
					    "%s: unknown challenge\n", rbuf);
					exit(1);
				}
				npass++;
			}
			if (npass == 2)
				break;
			if (rbuf[i] == '\0' && npass == 1)
				pass = rbuf + i + 1;
		}
	} else {
		pass = readpassphrase("TOTP code: ", pbuf, sizeof(pbuf),
		    RPP_ECHO_ON);
		if (pass == NULL)
			goto reject;
	}

	if (!totp_login(user, pass)) {
		clock = time(NULL);
		curtime = ctime(&clock);
		if (curtime != NULL)
			/* XXX syslog too? */
			fprintf(stderr, "Current time is %s", curtime);

		goto reject;
	}

	fprintf(back, BI_AUTH "\n");
	exit(0);

 reject:
	fprintf(back, BI_REJECT "\n");
	exit(1);
}