/* Copyright (C) 2022  June McEnroe <june@causal.agency>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 * Additional permission under GNU GPL version 3 section 7:
 *
 * If you modify this Program, or any covered work, by linking or
 * combining it with OpenSSL (or a modified version of that library),
 * containing parts covered by the terms of the OpenSSL License and the
 * original SSLeay license, the licensors of this Program grant you
 * additional permission to convey the resulting work. Corresponding
 * Source for a non-source form of such a combination shall include the
 * source code for the parts of OpenSSL used as well as that of the
 * covered work.
 */

#include <assert.h>
#include <err.h>
#include <errno.h>
#include <limits.h>
#include <netdb.h>
#include <netinet/in.h>
#include <resolv.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sysexits.h>
#include <tls.h>
#include <unistd.h>

char *readpassphrase(const char *prompt, char *buf, size_t bufsiz, int flags);

static bool verbose;
static void log(const char *format, ...) __attribute__((format(printf, 1, 2)));
static void log(const char *format, ...) {
	va_list ap;
	if (!verbose) return;
	va_start(ap, format);
	fprintf(stderr, "# ");
	vfprintf(stderr, format, ap);
	fprintf(stderr, "\n");
	va_end(ap);
}

static void wrap(size_t margin, const char *str, size_t len) {
	static size_t cols;
	if (!cols) {
		struct winsize ws;
		int error = ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws);
		if (error) {
			cols = 80;
		} else {
			cols = ws.ws_col;
		}
	}
	while (len) {
		if (str[0] == ' ') {
			str++;
			len--;
		}
		size_t line = cols - margin - 1;
		if (line > len) {
			line = len;
		} else {
			while (line && str[line] != ' ') line--;
			if (!line) line = len;
		}
		fwrite(str, line, 1, stdout);
		if (line < len) putchar('\n');
		str += line;
		len -= line;
	}
}

static void wrapf(const char *format, ...)
__attribute__((format(printf, 1, 2)));
static void wrapf(const char *format, ...) {
	va_list ap;
	char buf[1024];
	va_start(ap, format);
	int len = vsnprintf(buf, sizeof(buf), format, ap);
	assert((size_t)len < sizeof(buf));
	va_end(ap);
	wrap(0, buf, len);
	putchar('\n');
}

static int prompt1(char **resp, const char *def, const char *desc) {
	static char *buf;
	static size_t cap;

	size_t head = strcspn(desc, "\n");
	for (*resp = NULL; !*resp;) {
		wrap(strlen(def ?: "") * 2 + 4, desc, head);
		printf(" [%s] ", (def ?: ""));
		ssize_t len = getline(&buf, &cap, stdin);
		if (len < 0 && feof(stdin)) {
			printf("\n");
			clearerr(stdin);
			return -1;
		} else if (len < 0) {
			err(EX_IOERR, "getline");
		}
		if (len && buf[len-1] == '\n') buf[len-1] = '\0';

		if (buf[0] == '?') {
			printf("\n");
			wrap(0, &desc[head+1], strlen(&desc[head+1]));
			printf("\n\n");
		} else if (buf[0]) {
			*resp = strdup(buf);
			if (!*resp) err(EX_OSERR, "strdup");
		} else if (def) {
			*resp = strdup(def);
			if (!*resp) err(EX_OSERR, "strdup");
		}
	}
	return 0;
}

static int prompt(char **resp, const char *def, const char *format, ...)
__attribute__((format(printf, 3, 4)));
static int prompt(char **resp, const char *def, const char *format, ...) {
	va_list ap;
	char desc[1024];
	va_start(ap, format);
	vsnprintf(desc, sizeof(desc), format, ap);
	va_end(ap);
	return prompt1(resp, def, desc);
}

static int yesno(bool *resp, bool def, const char *format, ...)
__attribute__((format(printf, 3, 4)));
static int yesno(bool *resp, bool def, const char *format, ...) {
	va_list ap;
	char desc[1024];
	va_start(ap, format);
	vsnprintf(desc, sizeof(desc), format, ap);
	va_end(ap);
	char *str;
	for (;;) {
		int pop = prompt1(&str, (def ? "yes" : "no"), desc);
		if (pop) return pop;
		char yn = str[0];
		free(str);
		if (yn == 'Y' || yn == 'y') {
			*resp = true;
			return 0;
		} else if (yn == 'N' || yn == 'n') {
			*resp = false;
			return 0;
		}
	}
}

static int password(char **resp, bool confirm, const char *format, ...)
__attribute__((format(printf, 3, 4)));
static int password(char **resp, bool confirm, const char *format, ...) {
	va_list ap;
	char desc[1024];
	va_start(ap, format);
	vsnprintf(desc, sizeof(desc), format, ap);
	va_end(ap);

	enum { Cap = 1024 };
	char *buf1 = malloc(Cap);
	if (!buf1) err(EX_OSERR, "malloc");

	char *buf2 = NULL;
	if (confirm) {
		buf2 = malloc(Cap);
		if (!buf2) errx(EX_OSERR, "malloc");
	}

	for (;;) {
		*resp = readpassphrase(desc, buf1, Cap, 0);
		if (!*resp) errx(EX_IOERR, "could not read passphrase");
		if (!buf1[0]) {
			*resp = NULL;
			free(buf1);
			free(buf2);
			return -1;
		}
		if (!confirm) return 0;

		*resp = readpassphrase("Confirm password: ", buf2, Cap, 0);
		if (!*resp) errx(EX_IOERR, "could not read passphrase");
		if (!strcmp(buf2, buf1)) {
			explicit_bzero(buf1, Cap);
			free(buf1);
			return 0;
		}
		printf("Passwords do not match.\n");
	}
}

static struct {
	char *network;
	char *config;
	char *host;
	char *port;
	struct addrinfo *addr;
	struct tls_config *tls;
	char *trust;
	struct {
		char *name;
		char *pass;
	} account;
	char *nick;
	char *pronouns;
} info;

static void unset(char **resp) {
	free(*resp);
	*resp = NULL;
}

static int getNetwork(const char *def) {
	return prompt(
		&info.network, def,
		"Which network would you like to connect to?\n"
		"Enter the domain name of the IRC network you wish to connect to. "
		"Examples: tilde.chat, libera.chat "
	);
}

static int getConfig(const char *configHome) {
	char *name = strdup(info.network);
	if (!name) err(EX_OSERR, "strdup");
	char *prev = NULL, *next = NULL;
	while (name) {
		prev = next;
		next = strsep(&name, ".");
	}

	const char *home = getenv("HOME");
	const char *suffix = configHome;
	size_t homeLen = strlen(home);
	if (!strncmp(configHome, home, homeLen)) {
		suffix += homeLen;
	}
	int pop = prompt(
		&info.config, prev,
		"What would you like to name this configuration?\n"
		"This is the name you will pass to catgirl to connect to %s. "
		"The configuration will be written to: %s%s/catgirl ",
		prev, (suffix != configHome ? "~" : ""), suffix
	);
	free(name);
	if (pop) return pop;

	char path[PATH_MAX];
	snprintf(path, sizeof(path), "%s/catgirl/%s", configHome, info.config);
	if (!access(path, F_OK)) {
		wrapf("There is already a configuration named %s.", info.config);
		wrapf(
			"Please remove or rename the file: %s%s/catgirl/%s",
			(suffix != configHome ? "~" : ""), suffix, info.config
		);
		wrapf("Otherwise, enter a different name.");
		free(info.config);
		info.config = NULL;
	}
	return 0;
}

static int srvLookup(char **host, char **port, const char *name) {
	char dname[256];
	snprintf(dname, sizeof(dname), "_ircs._tcp.%s", name);
	log("Looking up SRV record for %s", dname);

	uint8_t msg[512];
	int len = res_query(dname, 1 /* IN */, 33 /* SRV */, msg, sizeof(msg));
	if (len < 12) return -1;

	uint8_t *ptr = &msg[12];
#define NAME_SKIP \
	for (uint8_t n; ptr < &msg[len] && (n = *ptr++); ptr += n) { \
		if (n & 0xC0) { ptr++; break; } \
	}

	uint16_t qdcount = msg[4] << 8 | msg[5];
	for (uint16_t q = 0; ptr < &msg[len] && q < qdcount; ++q) {
		NAME_SKIP; // QNAME
		ptr += 4; // QTYPE, QCLASS
	}

	NAME_SKIP; // NAME
	ptr += 14; // TYPE, CLASS, TTL, RDLENGTH, Priority, Weight
	if (&msg[len] < ptr + 3) return -1;
#undef NAME_SKIP

	int n = asprintf(port, "%d", ptr[0] << 8 | ptr[1]);
	if (n < 0) err(EX_OSERR, "asprintf");
	ptr += 2;

	// Name compression is not used for Target.
	if (!ptr[0]) return -1;
	char *host0 = (char *)&ptr[1];
	for (uint8_t n; ptr < &msg[len] && (n = *ptr); ptr += n) {
		*ptr++ = '.';
	}
	*host = strdup(host0);
	if (!*host) err(EX_OSERR, "strdup");
	return 0;
}

static int getHostPort(void) {
	int error = srvLookup(&info.host, &info.port, info.network);
	if (!error) {
		log("Found SRV record %s:%s", info.host, info.port);
		return 0;
	}

	char *ircdot = info.network;
	if (strncmp(ircdot, "irc.", 4)) {
		int n = asprintf(&ircdot, "irc.%s", ircdot);
		if (n < 0) err(EX_OSERR, "asprintf");
	}

	int pop = prompt(
		&info.host, ircdot,
		"Which host should catgirl connect to?\n"
		"Enter the hostname of the IRC server to connect to. "
		"Usually this starts with \"irc.\". "
	);
	free(ircdot);
	if (pop) return pop;

	return prompt(
		&info.port, "6697",
		"Which port should catgirl connect to?\n"
		"Enter the port number of the IRC server. "
		"This must be the port for TLS, also called SSL. "
	);
}

static int getAddr(void) {
	log("Looking up %s:%s", info.host, info.port);
	struct addrinfo hints = {
		.ai_family = AF_UNSPEC,
		.ai_socktype = SOCK_STREAM,
		.ai_protocol = IPPROTO_TCP,
	};
	int error = getaddrinfo(info.host, info.port, &hints, &info.addr);
	if (error) {
		log("%s", gai_strerror(error));
		wrapf("Could not find %s.", info.host);
		return -1;
	}
	return 0;
}

static int connectSock(void) {
	log("Connecting to %s:%s", info.host, info.port);
	for (struct addrinfo *ai = info.addr; ai; ai = ai->ai_next) {
		int sock = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
		if (sock < 0) err(EX_OSERR, "socket");

		int error = connect(sock, ai->ai_addr, ai->ai_addrlen);
		if (!error) return sock;

		close(sock);
	}
	wrapf(
		"Could not connect to %s:%s: %s.",
		info.host, info.port, strerror(errno)
	);
	return -1;
}

static struct tls *connectTLS(int sock) {
	struct tls *client = tls_client();
	if (!client) errx(EX_SOFTWARE, "tls_client");

	int error = tls_configure(client, info.tls);
	if (error) errx(EX_SOFTWARE, "tls_configure: %s", tls_error(client));

	error = tls_connect_socket(client, sock, info.host)
		|| tls_handshake(client);
	if (!error) return client;

	wrapf(
		"Could not establish TLS with %s:%s: %s",
		info.host, info.port, tls_error(client)
	);
	tls_close(client);
	tls_free(client);
	close(sock);
	return NULL;
}

static int getTLS(const char *configHome) {
	info.tls = tls_config_new();
	if (!info.tls) errx(EX_SOFTWARE, "tls_config_new");

	struct tls *client = tls_client();
	if (!client) errx(EX_SOFTWARE, "tls_client");

	int error = tls_configure(client, info.tls);
	if (error) errx(EX_SOFTWARE, "tls_configure: %s", tls_error(client));

	int sock = connectSock();
	if (sock < 0) goto fail;

	log("Performing TLS handshake");
	error = tls_connect_socket(client, sock, info.host)
		|| tls_handshake(client);
	if (!error) {
		log("TLS established");
		tls_close(client);
		tls_free(client);
		close(sock);
		return 0;
	}
	wrapf(
		"Could not establish TLS with %s:%s: %s",
		info.host, info.port, tls_error(client)
	);

	// XXX: Comparing error strings, don't have a better way.
	bool selfSigned = !strcmp(
		tls_error(client),
		"certificate verification failed: self-signed certificate"
	);
	tls_close(client);
	tls_free(client);
	close(sock);
	if (!selfSigned) goto fail;

	int n = asprintf(&info.trust, "%s/catgirl/%s.pem", configHome, info.host);
	if (n < 0) err(EX_OSERR, "asprintf");

	char *suffix = info.trust;
	const char *home = getenv("HOME");
	size_t homeLen = strlen(home);
	if (!strncmp(suffix, home, homeLen)) {
		suffix += homeLen;
	}

	bool trustIt;
	int pop = yesno(
		&trustIt, true,
		"Trust the server's self-signed certificate?\n"
		"If you are sure you are connecting to the right place, answer yes. "
		"The current certificate will be saved and explicitly trusted for "
		"future connections to this network. "
		"The certificate will be written to: %s%s ",
		(suffix != info.trust ? "~" : ""), suffix
	);
	if (pop || !trustIt) goto fail;

	tls_config_insecure_noverifycert(info.tls);
	tls_config_insecure_noverifyname(info.tls);
	sock = connectSock();
	if (sock < 0) goto fail;
	client = connectTLS(sock);
	if (!client) goto fail;

	size_t pemLen;
	const uint8_t *pem = tls_peer_cert_chain_pem(client, &pemLen);
	if (!pem) {
		wrapf(
			"Could not obtain TLS certificate from %s:%s.",
			info.host, info.port
		);
		goto fail;
	}

	log("Writing certificate to %s", info.trust);
	FILE *file = fopen(info.trust, "w");
	if (!file) err(EX_CANTCREAT, "%s", info.trust);
	fprintf(file, "subject= %s\n", tls_peer_cert_subject(client));
	fwrite(pem, pemLen, 1, file);
	error = fclose(file);
	if (error) err(EX_IOERR, "%s", info.trust);

	tls_close(client);
	tls_free(client);
	close(sock);

	tls_config_verify(info.tls);
	tls_config_insecure_noverifyname(info.tls);
	error = tls_config_set_ca_file(info.tls, info.trust);
	if (error) {
		errx(EX_SOFTWARE, "%s: %s", info.trust, tls_config_error(info.tls));
	}

	return 0;

fail:
	unset(&info.trust);
	tls_config_free(info.tls);
	info.tls = NULL;
	return -1;
}

static int getAccountName(void) {
	int pop = prompt(
		&info.account.name, "",
		"If you already have an account on %s, what is its name?\n"
		"If you have a NickServ account already, "
		"enter the name of the account. "
		"If you aren't sure, leave this blank. ",
		info.network
	);
	if (pop) return pop;
	if (!info.account.name[0]) return 0;

	info.nick = strdup(info.account.name);
	if (!info.nick) err(EX_OSERR, "strdup");
	return 0;
}

static int getAccountPass(void) {
	return password(
		&info.account.pass, false,
		"Password for %s on %s: ", info.account.name, info.network
	);
}

static int getNick(void) {
	return prompt(
		&info.nick, getenv("USER"),
		"What name would you like to use on %s?\n"
		"This is the name others will see on IRC, called a nick. "
		"It must be unique, and can usually contain letters, digits, "
		"underscores and hyphens, but cannot start with a digit. ",
		info.network
	);
}

static int getPronouns(void) {
	return prompt(
		&info.pronouns, NULL,
		"What are your pronouns?\n"
		"This will be added to your \"real name\" so that other IRC users can "
		"find out how to refer to you with /whois. "
		"Examples: they/them, she/her, he/him "
	);
}

int main(int argc, char *argv[]) {
	for (int opt; 0 < (opt = getopt(argc, argv, "v"));) {
		switch (opt) {
			break; case 'v': verbose = true;
			break; default: return EX_USAGE;
		}
	}

	char configHome[PATH_MAX];
	if (getenv("XDG_CONFIG_HOME")) {
		snprintf(
			configHome, sizeof(configHome), "%s", getenv("XDG_CONFIG_HOME")
		);
	} else if (getenv("HOME")) {
		snprintf(configHome, sizeof(configHome), "%s/.config", getenv("HOME"));
	} else {
		errx(EX_USAGE, "HOME unset");
	}

	char path[PATH_MAX];
	snprintf(path, sizeof(path), "%s/catgirl", configHome);
	int error = mkdir(path, S_IRWXU);
	if (error && errno != EEXIST) err(EX_CANTCREAT, "%s", path);
	umask(S_IRWXG | S_IRWXO);

	wrapf(
		"Welcome to catgirl enrollment! "
		"This tool should get you connected to IRC using catgirl :3 "
	);
	printf("\n");
	wrapf(
		"Type ? at a prompt for a more detailed explanation. "
		"The default response is shown in [brackets]. "
		"Press enter at a prompt to select the default response. "
		"Press control-D at a prompt to return to the previous one. "
	);
	printf("\n");

	int pop;
	for (;;) {
		if (!info.network) {
			pop = getNetwork((optind < argc ? argv[optind] : NULL));
			if (pop) return EX_USAGE;
		} else if (!info.config) {
			pop = getConfig(configHome);
			if (pop) unset(&info.network);
		} else if (!info.host || !info.port) {
			pop = getHostPort();
			if (pop) unset(&info.config);
		} else if (!info.addr) {
			pop = getAddr();
			if (pop) {
				unset(&info.host);
				unset(&info.port);
			}
		} else if (!info.tls) {
			pop = getTLS(configHome);
			if (pop) {
				freeaddrinfo(info.addr);
				info.addr = NULL;
				unset(&info.host);
				unset(&info.port);
			}
		} else if (!info.account.name) {
			pop = getAccountName();
			if (pop) {
				unset(&info.trust);
				freeaddrinfo(info.addr);
				info.addr = NULL;
				unset(&info.host);
				unset(&info.port);
			}
		} else if (info.account.name[0] && !info.account.pass) {
			pop = getAccountPass();
			if (pop) unset(&info.account.name);
		} else if (!info.nick) {
			pop = getNick();
			if (pop) {
				unset(&info.account.name);
				unset(&info.account.pass);
			}
		} else if (!info.pronouns) {
			pop = getPronouns();
			if (pop) unset(&info.nick);
		} else {
			break;
		}
	}

	snprintf(path, sizeof(path), "%s/catgirl/%s", configHome, info.config);
	log("Writing configuration to %s", path);
	FILE *file = fopen(path, "w");
	if (!file) err(EX_CANTCREAT, "%s", path);

	fprintf(file, "host = %s\n", info.host);
	if (strcmp(info.port, "6697")) {
		fprintf(file, "port = %s\n", info.port);
	}
	if (info.trust) {
		fprintf(file, "trust = %s.pem\n", info.host);
	}

	// TODO: Check for CertFP instead.
	bool saslPass;
	pop = yesno(
		&saslPass, true,
		"Would you like to save your account password?\n"
		"If yes, your account password will be written to the configuration "
		"file and sent automatically using SASL when catgirl connects. "
		"Otherwise, only your account name will be saved, and catgirl will "
		"prompt you for the password every time. "
	);
	if (!pop && saslPass) {
		fprintf(
			file, "sasl-plain = %s:%s\n", info.account.name, info.account.pass
		);
	} else {
		fprintf(file, "sasl-plain = %s:\n", info.account.name);
	}

	fprintf(file, "nick = %s\n", info.nick);
	if (strlen(info.pronouns) > 1) {
		fprintf(file, "real = %s [%s]\n", info.nick, info.pronouns);
	}

	error = fclose(file);
	if (error) err(EX_IOERR, "%s", path);

	printf("\n");
	wrapf("You're all set! You can connect to %s using:", info.network);
	printf("$ catgirl %s\n\n", info.config);
	bool run;
	pop = yesno(
		&run, true,
		"Connect now?\n"
		"There's nothing left to do. "
		"This will start catgirl. "
	);
	if (pop || !run) return EX_OK;

	execlp("catgirl", "catgirl", info.config, NULL);
	err(EX_UNAVAILABLE, "catgirl");
}