#define _POSIX_C_SOURCE 200809L

#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <string.h>
#include <strings.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <time.h>
#include <sys/stat.h>
#include <fcntl.h>

#include "tlse.h"
#define MARK fprintf(stderr, "%s %s:%d\n", __FILE__, __func__, __LINE__)

struct tls_uri {
	char *scheme;
	char *userinfo;
	char *host;
	char *port;
	char *path;
	char *query;
	char *fragment;
};
int tls_parse_uri(char *, struct tls_uri *);
void tls_free_uri(struct tls_uri *);

int open_tcp_connection(char *host, int port);

/* if trustpolicy is 0, we just accept anything */
int verify_trust(struct TLSContext *context, struct TLSCertificate **chain, int
		len) {
	/* suppress unused */
	if (context || chain || len) {
		return 0;
	}

	return 0;
}

static char hexchars[] = "0123456789abcdefABCDEF";

static void hex(char *dst, uint8_t *src, size_t len) {
	while (len--) {
		dst[0] = hexchars[(src[0]>>4)&0xf];
		dst[1] = hexchars[src[0]&0xf];
		dst+=2;
		src++;
	}
}

#if 0
static void hexbin(uint8_t *dst, unsigned char *src, size_t len) {
	size_t i;
	int x;

	for (i=0; i<len; i+=2) {
		sscanf((const char *)src+i, "%02x", &x);
		dst[i/2] = x;
	}
}
#endif

char *my_getline(struct tls_buffer *b, int fd, size_t *size) {
	char *loc = 0;
	char buf[4096];
	ssize_t rv = 0;

	while (b->error == 0) {
		loc = memchr(b->buffer, '\n', b->len);
		if (loc) {
			*size = loc - b->buffer + 1;
			return b->buffer;
		} else {
			rv = read(fd, buf, sizeof buf);
			if (rv == -1) {
				return 0;
			}
			if (rv == 0) {
				break;
			}
			tls_buffer_append(b, buf, rv);
		}
	}
	if (rv == 0) {
		*size = b->len;
		return b->buffer;
	}
	return 0;
}

/*
 * We use a trust on first use policy.  The trust DB is a simple
 * file in /var/lib/zpm/known_hosts, or ~/.zpm/known_hosts, or ZPM_KNOWNHOSTS.
 * if -k is given, no verification is done
 */
int verify_first(struct TLSContext *context, struct TLSCertificate **chain, int
		certs) {
	int err;
	char *trustfile, *homedir = 0, *host, *fp;
	unsigned char certhash[65];
	int trustdb;
	struct tls_buffer tbuf;

	char *line = 0;
	size_t len = 0;

	if (certs == 0 || chain == 0) {
		return 1;
	}

	err = tls_certificate_is_valid(chain[0]);
	if (err) {
		return err;
	}

	if (context->sni) {
		err = tls_certificate_valid_subject(chain[0], context->sni);
		if (err) {
			return err;
		}
	}

	hex(certhash, chain[0]->fp, 32);
	certhash[64] = 0;

	trustfile = getenv("ZPM_KNOWNHOSTS");
	if (!trustfile) {
		if (geteuid() == 0) {
			trustfile = "/var/lib/zpm/known_hosts";
		} else {
			/* we could do this with a series of
			 * openat() calls instead of building
			 * up a string 
			 */
			trustfile = getenv("HOME");
			if (!trustfile) {
				fprintf(stderr, "home = %s\n", trustfile);
				return 1;
			}
			len = snprintf(homedir, 0, "%s/.zpm/known_hosts", trustfile);
			homedir = malloc(len+1);
			if (!homedir) {
				return 1;
			}
			len = snprintf(homedir, len+1, "%s/.zpm/known_hosts", trustfile);
			trustfile = homedir;
		}
	}
	/* cert is valid on its face, so check against the trust db */
	trustdb = open(trustfile, O_RDWR|O_CREAT, 0600);
	if (trustdb == -1) {
		fprintf(stderr, "cannot open trustdb %s: %s\n", trustfile, strerror(errno));
		if (homedir) {
			free(homedir);
		}
		return 1;
	}

	if (homedir) {
		free(homedir);
	}

	len = 0;
	tls_buffer_init(&tbuf, 128);
	do {
		char *off;

		tls_buffer_shift(&tbuf, len);
		line = my_getline(&tbuf, trustdb, &len);

		if (!line || !len) {
			break;
		}

		fp = line;
		while (isspace(*fp)) {
			fp++;
		}
		if (*fp == '#') {
			continue;
		}
		off = strchr(line, ':');
		if (!off) {
			continue;
		}
		host = off + 1;
		*off = 0;
		if (line[len-1] == '\n') {
			line[len-1] = 0;
		}

		if (strlen(fp) != 64) {
			continue;
		}

		if (len && line[len-1] == '\n') {
			line[len-1] = 0;
		}

		if (strcmp(context->sni, host) != 0) {
			continue;
		}

		int match = (memcmp(certhash, fp, 64) == 0); 

		close(trustdb);
		tls_buffer_free(&tbuf);
		return match ? no_error : bad_certificate;
	} while (!tbuf.error);

	/* got here, so we should be at EOF, so add this host to trust db */
	lseek(trustdb, 0, SEEK_END);

	/* re-use the buffer so we only do one write */
	/* ignore errors, the cert is fine, we just can't update
	 * the trustdb if there's errors here
	 */
	tbuf.len = 0;
	tls_buffer_append(&tbuf, certhash, 64);
	tls_buffer_append_byte(&tbuf, ':');
	tls_buffer_append(&tbuf, context->sni, strlen(context->sni));
	tls_buffer_append_byte(&tbuf, '\n');
	write(trustdb, tbuf.buffer, tbuf.len);
	close(trustdb);
	tls_buffer_free(&tbuf);

	return no_error;
}

int verify_roots(struct TLSContext *context, struct TLSCertificate **chain, int
		len) {
	int i, err;

	if (chain) {
		for (i = 0; i < len; i++) {
			struct TLSCertificate *certificate = chain[i];
			err = tls_certificate_is_valid(certificate);
			if (err) {
				return err;
			}
		}
	}
	
	err = tls_certificate_chain_is_valid(chain, len);
	if (err) {
		return err;
	}

	if (len > 0 && context->sni) {
		err = tls_certificate_valid_subject(chain[0], context->sni);
		if (err) {
			return err;
		}
	}

	/* Perform certificate validation against ROOT CA */
	err = tls_certificate_chain_is_valid_root(context, chain, len);
	if (err) {
		return err;
	}

	return no_error;
}

struct io {
	struct tls_buffer response;
	struct TLSContext *tls;
	int socket;
	int status_code;
	time_t last_modified;
	time_t date;
	size_t content_length;
	char *redirect;
};

int month(char *m) {
	char *months[] = {"Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul",
		"Aug", "Sep", "Oct", "Nov", "Dec" };
	int i;

	for (i=0; i < 12; i++) {
		if (!strncasecmp(m, months[i], 3)) {
			return i+1;
		}
	}
	return 0;
}

/* Wed, 06 Feb 2019 10:06:05 GMT */
time_t parse_date(char *d) {
	struct tm tm = { 0 };
	int rv;

	int h, m, s, dom, Y;
	char M[4];
	rv = sscanf(d, "%*[a-zA-Z,] %d %s %d %d:%d:%d", &dom, M, &Y, &h, &m, &s);
	if (rv == 6) {
		tm.tm_year = Y - 1900;
		tm.tm_hour = h;
		tm.tm_min = m;
		tm.tm_sec = s;
		tm.tm_mon = month(M)-1;
		tm.tm_mday = dom;

		return mktime(&tm);
	}
	return -1;
}

char *find_header(struct io *io, char *header, size_t *len) {
	char *eoh, *soh;
	size_t hlen;

	*len = 0;

	hlen = strlen(header);
	eoh = strstr(io->response.buffer, "\r\n\r\n");
	if (!eoh) {
		return 0;
	}
	soh = io->response.buffer;
	do {
		soh = strstr(soh, "\r\n");
		if (soh == eoh) {
			break;
		}
		soh += 2;
		if (!memcmp(soh, header, hlen)) {
			break;
		}
	} while (soh < eoh);

	if (soh >= eoh) {
		return 0;
	}
	eoh = strstr(soh, "\r\n");
	soh += hlen;
	while (soh < eoh && isspace(*soh)) {
		soh++;
	}
	*len = (eoh - soh);
	return soh;
}

void parse_header(struct io *io) {
	char *s = io->response.buffer;
	int code = 0;
	char *hval;
	size_t hlen;

	while (!isspace(*s)) {
		s++;
	}
	while (isspace(*s)) {
		s++;
	}
	code = strtol(s, 0, 10);
	io->status_code = code;

	hval = find_header(io, "Date:", &hlen);
	if (hval) {
		hval[hlen] = 0;
		io->date = parse_date(hval);
		hval[hlen] = '\r';
	}
	hval = find_header(io, "Last-Modified:", &hlen);
	if (hval) {
		hval[hlen] = 0;
		io->last_modified = parse_date(hval);
		hval[hlen] = '\r';
	}

	hval = find_header(io, "Content-Length:", &hlen);
	if (hval) {
		hval[hlen] = 0;
		io->content_length = strtoul(hval, 0, 10);
		hval[hlen] = '\r';
	}

	switch (code) {
		case 301:
		case 302:
		case 303:
		case 307:
			hval = find_header(io, "Location:", &hlen);
			if (hval) {
				io->redirect = strndup(hval, hlen);
			}
			break;
		default:
			break;
	}

}

ssize_t fill_buffer(struct io *io) {
	unsigned char buffer[4096];
	ssize_t ret;

	if (io->tls) {
		ret = tls_read(io->tls, buffer, sizeof buffer);
	} else {
		ret = read(io->socket, buffer, sizeof buffer);
	}

	if (ret > 0) {
		tls_buffer_append(&io->response, buffer, ret);
	}

	return ret;
}

#if 0
char *nextline(struct io *io) {
	char *eol = 0;;

	eol = memchr(io->response.buffer, '\n', io.response.size);
	while (eol == 0) {
		fill_buffer(io);
		eol = memchr(io->response.buffer, '\n', io.response.size);
	}
	if (eol) {

}
#endif

void append_header(struct tls_buffer *buf, char *header, char *val) {
	tls_buffer_append(buf, header, strlen(header));
	tls_buffer_append(buf, ": ", 2);
	tls_buffer_append(buf, val, strlen(val));
	tls_buffer_append(buf, "\r\n", 2);
}

void append_timeheader(struct tls_buffer *buf, char *header, time_t ts) {
	char timestr[80];
	struct tm *tm;

	tm = gmtime(&ts);

	strftime(timestr, sizeof timestr, "%a, %d %b %Y %H:%M:%S GMT", tm);
	append_header(buf, header, timestr);
}

static void pdots(int len, int ch, unsigned long was, unsigned long now,
		unsigned long total) {
	was = len * was / total;
	if (now > total) {
		now = total;
	}
	now = len * now / total;
	while (was++ < now) {
		putc(ch,stderr);
	}
}

static void fake_header(struct io *io, int fd) {
	struct stat st;
	int code = 200, rv;
	char *message, codestr[5], length[32];
	struct tls_buffer *hdr = &io->response;

	if (fd == -1) {
		switch (errno) {
			case EACCES: code = 403; break;
			case ENOENT: code = 404; break;
			default: code = 500; break;
		}
	} else {
		rv = fstat(fd, &st);
		if (rv == -1) {
			code = 500;
		}
	}

	if (io->last_modified >= st.st_mtime) {
		code = 304;
	}

	switch (code) {
		case 200: message = "OK"; break;
		case 304: message = "Not Modified"; break;
		case 403: message = "Forbidden"; break;
		case 404: message = "Not Found"; break;
		case 500: message = "Internal Server Error"; break;
		default: break;
	}
	sprintf(codestr, "%0.3d ", code);
	tls_buffer_append(hdr, "HTTP/1.1 ", 9);
	tls_buffer_append(hdr, codestr, 4);
	tls_buffer_append_str(hdr, message);
	tls_buffer_append(hdr, "\r\n", 2);

	append_timeheader(hdr, "Date", time(NULL));
	append_header(hdr, "Server", "zpm-fetchurl/0.9");
	if (code < 400) {
		append_timeheader(hdr, "Last-Modified", st.st_mtime);
		sprintf(length, "%zu", st.st_size);
		append_header(hdr, "Content-Length", length);
	}

	append_header(hdr, "Connection", "close");
	append_header(hdr, "Content-Type", "application/octet-stream");
	tls_buffer_append(hdr, "\r\n", 2);
}

int main(int ac, char *av[]) {
	int sockfd, port = -1, rv;
	ssize_t ret;
	int option;
#if 0
	char msg[] = "GET %s HTTP/1.1\r\nHost: %s:%i\r\nConnection: close\r\n\r\n";
	char msg2[] = "GET %s HTTP/1.1\r\nHost: %s:%i\r\nLast-Modified: %s\r\nConnection: close\r\n\r\n";
	char msg_buffer[1024];
#endif
	char *req_file = 0;
	char *host = 0;
	struct tls_uri uri;
	char *outfile = 0;
	int raw = 0, head = 0;
	int out = 1;
	int use_tls = 0;
	struct io io = { {0}, 0, -1, 0, 0, 0, 0, 0 };
	struct TLSContext *clientssl = 0;
	int failsilent = 0;
	char *lmfile = 0;
	int progressbar = 0;
	struct tls_buffer request;
	char lmtime[80];
	char *eoh = 0;
	size_t total = 0;
	size_t header_len;
	char *url = 0;
	int redirs = 0, redirlimit = 50, printstatus = 0;
	int verifypolicy = 1;

	ltc_mp = tfm_desc;

	while ((option = getopt(ac, av, "o:rIfz:#R:SkK")) != -1) {
		switch (option) {
			case 'o': outfile = optarg; break;
			case 'S': printstatus = 1; head = 1; break;
			case 'k': verifypolicy = 0; break;
			case 'K': verifypolicy = 2; break;
			case 'I': head = 1;
			case 'r': raw = 1; break;
			case 'f': failsilent = 1; break;
			case 'z': lmfile = optarg; break;
			case 'R': redirlimit = strtol(optarg, 0, 10); break;
			case '#': progressbar = 1; break;
			default:
				  exit(EXIT_FAILURE);
				  break;
		}
	}

	if (ac < optind) {
		fprintf(stderr, "Usage: %s uri\n", av[0]);
		exit(EXIT_FAILURE);
	}

	io.last_modified = 0;
	if (lmfile) {
		struct stat st;
		int rv;
		struct tm *mtime;
		time_t ts;

		rv = stat(lmfile, &st);
		if (rv == -1) {
			perror("stat failed:");
			exit(EXIT_FAILURE);
		}
		ts = st.st_mtime;
		io.last_modified = ts;
		mtime = gmtime(&ts);
		strftime(lmtime, sizeof lmtime, "%a, %d %b %Y %H:%M:%S GMT", mtime);
	}

	url = strdup(av[optind]);
	if (!url) {
		exit(EXIT_FAILURE);
	}

	if (outfile) {
		out = open(outfile, O_WRONLY|O_CREAT, 0600);
		if (out == -1) {
			perror("can't open output file:");
			exit(EXIT_FAILURE);
		}
	}

	signal(SIGPIPE, SIG_IGN);

	tls_buffer_init(&io.response, 0);
	tls_buffer_init(&request, 128);

	while (redirs++ <= redirlimit) {
		tls_free_uri(&uri);
		io.response.len = 0;
		request.len = 0;

		tls_parse_uri(url, &uri);
		host = uri.host;
		port = atoi(uri.port);
		req_file = uri.path;

		/* construct request */
		if (head) {
			tls_buffer_append(&request, "HEAD ", 5);
		} else {
			tls_buffer_append(&request, "GET ", 4);
		}
		tls_buffer_append(&request, uri.path, strlen(uri.path));
		tls_buffer_append(&request, " HTTP/1.1\r\n", 11);

		append_header(&request, "Host", host);
		append_header(&request, "Connection", "close");
		if (lmfile) {
			append_header(&request, "If-Modified-Since", lmtime);
		}
		tls_buffer_append(&request, "\r\n", 2);

		if (!strcmp(uri.scheme, "https")) {
			use_tls = 1;

			clientssl = tls_create_context(TLS_CLIENT, TLS_V12);

			/* optionally, we can set a certificate validation
			 * callback function if set_verify is not called, and
			 * root ca is set, `tls_default_verify` will be used
			 * (does exactly what `verify` does in this example)
			 */
			if (verifypolicy == 2) {
				char *cert_path = 0;
				cert_path = getenv("ZPM_CERTFILE");
				if (!cert_path) {
					cert_path = "/var/lib/zpm/roots.pem";
				}
				rv = tls_load_root_file(clientssl, cert_path);
				if (rv == -1) {
					fprintf(stderr, "Error loading root certs\n");
					return 1;
				}
				tls_set_verify(clientssl, verify_roots);
			} else if (verifypolicy == 1) {
				tls_set_verify(clientssl, verify_first);
			} else {
				tls_set_verify(clientssl, verify_trust);
			}

			if (!clientssl) {
				fprintf(stderr, "Error initializing client context\n");
				return -1;
			}
			tls_sni_set(clientssl, uri.host);
			clientssl->sync = 1;
			io.tls = clientssl;
			sockfd = open_tcp_connection(host, port);
			if (sockfd < 0) {
				perror("can't open connection");
				exit(EXIT_FAILURE);
			}
			tls_set_fd(clientssl, sockfd);
			if ((rv = tls_connect(clientssl)) != 1) {
				fprintf(stderr, "Handshake Error %i\n", rv);
				return 1;
			}
			ret = tls_write(clientssl, request.buffer, request.len);
		} else if (!strcmp(uri.scheme, "http")) {
			sockfd = open_tcp_connection(host, port);
			if (sockfd < 0) {
				perror("can't open connection");
				exit(EXIT_FAILURE);
			}
			ret = write(sockfd, request.buffer, request.len);
		} else if (!strcmp(uri.scheme, "file")) {
			sockfd = open(uri.path, O_RDONLY);
			fake_header(&io, sockfd);
			ret = 0;
		} else {
			fprintf(stderr, "scheme %s unknown\n", uri.scheme);
			exit(EXIT_FAILURE);
		}

		if (ret == -1) {
			fprintf(stderr, "unable to write http request: %s\n", strerror(errno));
			exit(EXIT_FAILURE);
		}

		io.socket = sockfd;

		do {
			if (io.response.len >= 4) {
				eoh = strstr(io.response.buffer, "\r\n\r\n");
			}
			if (!eoh) {
				ret = fill_buffer(&io);
				if (ret <= 0) {
					break;
				}
			}
		} while (!eoh);

		if (!eoh) {
			/* never got (complet) header */
			fprintf(stderr, "incomplete response to %s\n", av[optind]);
			exit(EXIT_FAILURE);
		}

		header_len = (size_t)(eoh - io.response.buffer) + 4;
		parse_header(&io);

		switch (io.status_code) {
			case 301:
			case 302:
			case 303:
			case 307:
				free(url);
				url = strdup(io.redirect);
				continue;
				break;
		}

		if (printstatus) {
			printf("%d\n", io.status_code);
			break;
		}

		if (!raw) {
			tls_buffer_shift(&io.response, header_len);
		}
		if (head) {
			io.response.len -= 2;
		}

		if (progressbar) {
			if (io.content_length) {
				fprintf(stderr, "(%lu) ", io.content_length);
			}
		}

		do {
			write(out, io.response.buffer, io.response.len);
			ret = io.response.len;
			io.response.len = 0;

			if (progressbar) {
				if (io.content_length) {
					pdots(50, '.', total, total+ret,
							io.content_length);
				} else {
					int old = total / 1000000;
					int new = (total+ret)/1000000;
					while (old < new) {
						putc('.',stderr);
					}
				}
				total += ret;
			}
			if (head) {
				break;
			}
			ret = fill_buffer(&io);
		} while (ret > 0);

		if (ret < 0) {
			fprintf(stderr, "%s read error %zd\n", uri.scheme, ret);
		}
		/* futimens(out, ...) */
		close(out);
		tls_buffer_free(&io.response);
		break;
	}

	if (use_tls) {
		tls_shutdown(clientssl);
		tls_free(clientssl);
	}

	close(sockfd);
	if (progressbar && io.status_code == 200) {
		fprintf(stderr, "(%lu)", total);
		putc('\n',stderr);
	}

	return io.status_code == 200 ? 0 : EXIT_FAILURE;
}