#include <errno.h>
#include <limits.h>
#include <pwd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>

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

#include <tls.h>

#include "common.h"
#include "io.h"

#define TLS_OFF 0
#define TLS_ON  1
#define TLS_PEM 2

struct pem {
	char path[PATH_MAX];
	char *dir;
	char *cert;
	size_t certsz;
};

int tls;

static struct pem pem = { .dir = ".share/sacc/cert" };

static int
mkpath(char *path, mode_t mode)
{
	char *s;
	int r;

	for (s = path+1; (s = strchr(s, '/')) != NULL; ++s) {
		s[0] = '\0';
		errno = 0;
		r = mkdir(path, mode);
		s[0] = '/';
		if (r == -1 && errno != EEXIST)
			return -1;
	};
	if (mkdir(path, S_IRWXU) == -1 && errno != EEXIST)
		return -1;
	return 0;
}

static int
setup_tls(void)
{
	struct passwd *pw;
	char *p;
	int n;

	if ((p = getenv("SACC_CERT_DIR")) != NULL) {
		n = snprintf(pem.path, sizeof(pem.path), "%s/", p);
		if (n < 0 || (unsigned)n >= sizeof(pem.path)) {
			diag("PEM path too long: %s/", p);
			return -1;
		}
	} else {
		if ((pw = getpwuid(geteuid())) == NULL)
			return -1;
		n = snprintf(pem.path, sizeof(pem.path), "%s/%s/",
		             pw->pw_dir, pem.dir);
		if (n < 0 || (unsigned)n >= sizeof(pem.path)) {
			diag("PEM path too long: %s/%s/", pw->pw_dir, pem.dir);
			return -1;
		}
	}

	if (mkpath(pem.path, S_IRWXU) == -1) {
		diag("Can't create cert dir: %s: %s",
		     pem.path, strerror(errno));
	} else {
		pem.cert = pem.path + n;
		pem.certsz = sizeof(pem.path) - n;
	}

	return 0;
}

static int
close_tls(struct cnx *c)
{
	int r;

	if (tls != TLS_OFF && c->tls) {
		do {
			r = tls_close(c->tls);
		} while (r == TLS_WANT_POLLIN || r == TLS_WANT_POLLOUT);

		tls_free(c->tls);
	}

	return close(c->sock);
}

static int
savepem(struct tls *t, char *path)
{
	FILE *f;
	const char *s;
	size_t ln;

	if (path == NULL)
		return -1;
	if ((s = tls_peer_cert_chain_pem(t, &ln)) == NULL)
		return -1;
	if ((f = fopen(path, "w")) == NULL)
		return -1;
	fprintf(f, "%.*s\n", ln, s);
	if (fclose(f) != 0)
		return -1;

	return 0;
}

static char *
conftls(struct tls *t, const char *host)
{
	struct tls_config *tc;
	char *p;
	int n;

	tc = NULL;
	p = NULL;

	if (pem.cert == NULL)
		return NULL;

	n = snprintf(pem.cert, pem.certsz, "%s", host);
	if (n < 0 || (unsigned)n >= pem.certsz) {
		diag("PEM path too long: %s/%s", pem.cert, host);
		return NULL;
	}

	switch (tls) {
	case TLS_ON:
		/* check if there is a local certificate for target */
		if (access(pem.path, R_OK) == 0) {
			if ((tc = tls_config_new()) == NULL)
				return NULL;
			if (tls_config_set_ca_file(tc, pem.path) == -1)
				goto end;
			if (tls_configure(t, tc) == -1)
				goto end;
			p = pem.path;
		}
		break;
	case TLS_PEM:
		/* save target certificate to file */
		if ((tc = tls_config_new()) == NULL)
			return NULL;
		tls_config_insecure_noverifycert(tc);
		if (tls_configure(t, tc) == -1)
			goto end;
		p = pem.path;
		break;
	}
end:
	tls_config_free(tc);
	return p;
}

static int
connect_tls(struct cnx *c, struct addrinfo *ai, const char *host)
{
	struct tls *t;
	char *s, *pempath;
	int r;

	c->tls = NULL;
	s = NULL;
	r = CONN_ERROR;

	if (connect(c->sock, ai->ai_addr, ai->ai_addrlen) == -1)
		return r;

	if (!tls)
		return CONN_VALID;

	if ((t = tls_client()) == NULL)
		return r;

	pempath = conftls(t, host);

	if (tls_connect_socket(t, c->sock, host) == -1)
		goto end;

	do {
		r = tls_handshake(t);
	} while (r == TLS_WANT_POLLIN || r == TLS_WANT_POLLOUT);

	if (r == 0) {
		switch (tls) {
		case TLS_ON:
			c->tls = t;
			break;
		case TLS_PEM:
			r = savepem(t, pempath) == 0 ? CONN_RETRY : CONN_ERROR;
			tls = TLS_ON;
			break;
		}
	} else {
		diag("Can't establish TLS with \"%s\": %s",
		     host, tls_error(t));

		if (!interactive) {
			r = CONN_ABORT;
			goto end;
		}

		if (pem.cert) {
			s = uiprompt("Save certificate locally and retry? [yN]: ");
			switch (*s) {
			case 'Y':
			case 'y':
				tls = TLS_PEM;
				r = CONN_RETRY;
				goto end;
			}
		}

		s = uiprompt("Retry on cleartext? [Yn]: ");
		switch (*s) {
		case 'Y':
		case 'y':
		case '\0':
			tls = TLS_OFF;
			r = CONN_RETRY;
			break;
		default:
			r = CONN_ABORT;
		}
	}
end:
	free(s);
	if (r != CONN_VALID)
		tls_free(t);

	return r;
}

static void
connerr_tls(struct cnx *c, const char *host, const char *port, int err)
{
	if (c->sock == -1) {
		diag("Can't open socket: %s", strerror(err));
	} else {
		if (tls != TLS_OFF && c->tls) {
			diag("Can't establish TLS with \"%s\": %s",
			     host, tls_error(c->tls));
		} else {
			diag("Can't connect to: %s:%s: %s", host, port,
			     strerror(err));
		}
	}
}

static char *
parseurl_tls(char *url)
{
	char *p;

	if (p = strstr(url, "://")) {
		if (!strncmp(url, "gopher", p - url)) {
			if (tls)
				diag("Switching from gophers to gopher");
			tls = TLS_OFF;
		} else if (!strncmp(url, "gophers", p - url)) {
			tls = TLS_ON;
		} else {
			die("Protocol not supported: %.*s", p - url, url);
		}
		url = p + 3;
	}

	return url;
}

static ssize_t
read_tls(struct cnx *c, void *buf, size_t bs)
{
	ssize_t n;

	if (tls != TLS_OFF && c->tls) {
		do {
			n = tls_read(c->tls, buf, bs);
		} while (n == TLS_WANT_POLLIN || n == TLS_WANT_POLLOUT);
	} else {
		n = read(c->sock, buf, bs);
	}

	return n;
}

static ssize_t
write_tls(struct cnx *c, void *buf, size_t bs)
{
	ssize_t n;

	if (tls) {
		do {
			n = tls_write(c->tls, buf, bs);
		} while (n == TLS_WANT_POLLIN || n == TLS_WANT_POLLOUT);
	} else {
		n = write(c->sock, buf, bs);
	}

	return n;
}

int (*iosetup)(void) = setup_tls;
int (*ioclose)(struct cnx *) = close_tls;
int (*ioconnect)(struct cnx *, struct addrinfo *, const char *) = connect_tls;
void (*ioconnerr)(struct cnx *, const char *, const char *, int) = connerr_tls;
char *(*ioparseurl)(char *) = parseurl_tls;
ssize_t (*ioread)(struct cnx *, void *, size_t) = read_tls;
ssize_t (*iowrite)(struct cnx *, void *, size_t) = write_tls;