#define _XOPEN_SOURCE 600

#include <sys/stat.h>
#include <sys/types.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <sys/ioctl.h>
#include <fcntl.h>
#include <unistd.h>
#include <termios.h>
#include <signal.h>
#include <errno.h>
#include <stdbool.h>
#include <stdarg.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

struct termios oldterm;
int master, slave;
int pipeout;
pid_t child;

bool encode_resize;

void
usage(void)
{
	fprintf(stderr, "Usage: vtmitm [--encode-resize] -- CMD [ARG]..\n");
	exit(1);
}

void
err(const char *fmtstr, ...)
{
	va_list ap;

	va_start(ap, fmtstr);
	vfprintf(stderr, fmtstr, ap);
	fprintf(stderr, ": %s\r\n", strerror(errno));
	va_end(ap);

	exit(1);
}

int
writeall(int fd, char *buffer, size_t size)
{
	ssize_t nbytes;

	while (size) {
		nbytes = write(fd, buffer, size);
		if (nbytes <= 0) return 1;
		buffer += nbytes;
		size -= nbytes;
	}

	return 0;
}

int
fwd(int in, int out, int out2)
{
	static char buffer[8194];
	ssize_t got, i, last;

	do {
		got = read(in, buffer, sizeof(buffer));
		if (got < 0) return 1;

		if (writeall(out, buffer, got))
			return 1;

		if (out2 >= 0) {
			if (writeall(out2, buffer, got))
				return 1;
		}
	} while (got == sizeof(buffer));

	return 0;
}

void
mitm(int pts, int pipe)
{
	fd_set fds;

	while (1) {
		FD_ZERO(&fds);
		FD_SET(0, &fds);
		FD_SET(pts, &fds);

		if (select(pts + 1, &fds, 0, 0, 0) == -1) {
			if (errno == EINTR) continue;
			return;
		}

		if (FD_ISSET(0, &fds)) { /* stdin -> ptty */
			if (fwd(0, pts, -1)) return;
		}

		if (FD_ISSET(pts, &fds)) { /* ptty -> stdout + pipe */
			if (fwd(pts, 1, pipe)) return;
		}
	}
}

void
rawterm(struct termios *term)
{
	term->c_iflag &= ~(IGNBRK | BRKINT | PARMRK | ISTRIP);
	term->c_iflag &= ~(INLCR | IGNCR | ICRNL | IXON);
	term->c_oflag &= ~OPOST;
	term->c_lflag &= ~(ECHO | ECHONL | ICANON | ISIG | IEXTEN);
	term->c_cflag &= ~(CSIZE | PARENB);
	term->c_cflag |= CS8;
}

void
sigwinch(int sig)
{
	char buf[64];
	struct winsize wsize;

	if (ioctl(0, TIOCGWINSZ, &wsize) != -1)
		ioctl(slave, TIOCSWINSZ, &wsize);
	kill(child, SIGWINCH);
	signal(SIGWINCH, sigwinch);

	/* transmit new size info in form of xterm resize escape sequence */
	snprintf(buf, sizeof(buf), "\x1b[8;%u;%ut", wsize.ws_row, wsize.ws_col);
	write(pipeout, buf, sizeof(buf));
}

void
sigchld(int sig)
{
	int exitcode;

	wait(&exitcode);
	exit(exitcode);
}

void
cleanup(void)
{
	tcsetattr(0, TCSANOW, &oldterm);
}

int
main(int argc, char *const *argv)
{
	struct termios newterm;
	struct winsize wsize;
	char *const *arg, *const *cmd;
	int exitcode;

	cmd = NULL;
	encode_resize = false;
	for (arg = argv + 1; *arg; arg++) {
		if (!strcmp(*arg, "--encode-resize")) {
			encode_resize = true;
		} else if (!strcmp(*arg, "--")) {
			cmd = arg + 1;
			break;
		} else {
			usage();
		}
	}

	if (!cmd) usage();

	pipeout = dup(1);
	if (!freopen("/dev/tty", "w", stdout))
		err("main: open /dev/tty");
	dup2(fileno(stdout), 1);

	master = posix_openpt(O_RDWR);
	if (master == -1) err("main: posix_openpt");
	if (grantpt(master)) err("main: grantpt");
	if (unlockpt(master)) err("main: unlockpt");

	slave = open(ptsname(master), O_RDWR);
	if (slave == -1) err("main: open pts");

	signal(SIGCHLD, sigchld);

	child = fork();
	if (child < 0) err("main: fork");

	if (child == 0) {
		close(master);

		signal(SIGCHLD, SIG_DFL);

		if (setsid() == -1)
			err("slave: setsid");

		if (ioctl(slave, TIOCSCTTY, 0) == -1)
			err("slave: ioctl TIOCSCTTY");

		if (dup2(slave, 0) == -1)
			err("slave: dup2 stdin");
		if (dup2(slave, 1) == -1)
			err("slave: dup2 stdout");
		if (dup2(slave, 2) == -1)
			err("slave: dup2 stderr");
		close(slave);

		execvp(*cmd, cmd);
		err("slave: execvp");
	} else {
		if (encode_resize) {
			signal(SIGWINCH, sigwinch);
			raise(SIGWINCH);
		}

		if (tcgetattr(0, &oldterm))
			err("master: tcgetattr");
		newterm = oldterm;
		rawterm(&newterm);
		if (tcsetattr(0, TCSANOW, &newterm))
			err("master: tcsetattr");
		atexit(cleanup);

		mitm(master, pipeout);

		waitpid(child, &exitcode, 0);
		close(master);
	}

	return exitcode;
}