#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>

#define INIT "export EDITOR=ed PAGER=ul\n"
#define HINIT "export TERM=tty43 EDITOR=ed PAGER=cat\n"

#define MAXBUF 2048
#define MAXEARG 32
#define MAXWD 255

char *getpw(char *);
void killall();
void noop();

int
main(int argc, char *argv[])
{
	char	*address, bufin[MAXBUF], bufout[MAXBUF], dir[30],
		*eargv[MAXEARG], in[30], **nargv, out[30], *p, pw[255], *q,
		wd[MAXWD];
	fd_set	rfds0, rfds1;
	int	child, dumb, escape, fdin, fdout, hardcopy, i, n, offset, s,
		sshpass;
	struct timeval	tv;

	signal(SIGINT, noop);
	signal(SIGQUIT, killall);
	signal(SIGCHLD, killall);

	/* Create named pipes. */
	sprintf(in, "/var/tmp/rtty.in.%d", getpid());
	sprintf(out, "/var/tmp/rtty.out.%d", getpid());
	if(mkfifo(in, 0644) == -1 || mkfifo(out, 0664) == -1)
		if(errno != EEXIST)
			err(1, "mkfifo");

	/* Process rtty-specific flags. */
	escape = dumb = hardcopy = sshpass = 0;
	for(;argv[1][0] == '+';){
		if(strchr(argv[1], 'd')) dumb = 1;
		if(strchr(argv[1], 'h')) hardcopy = 1;
		if(strchr(argv[1], 'x')) escape = 1;
		if(strchr(argv[1], 'p')) sshpass = 1;

		/* Remove flag argument from argv. */
		argv[1] = argv[0];
		argv++; argc--;
	}

	/* Save pointer to address argument. */
	address = NULL;
	for(i = 1; i < argc; i++)
		if(argv[i][0] != '-'){
			address = argv[i];
			break;
		}

	/* If +p is given, immediately ask for password. */
	if(sshpass){
		printf("password: ");
		getpw(pw);
		setenv("SSHPASS", pw, 1);
		sshpass = 1;
	}

	/*
	 * A child ssh process is created, with standard in connected to
	 * the input pipe (fdin) and standard out connected to the output
	 * pipe (fdout).
	 */
	if(fork() == 0){
		if(!(fdin = open(in, O_RDONLY)))
			err(1, "open");
		if(!(fdout = open(out, O_WRONLY)))
			err(1, "open");
		dup2(fdin, 0);
		dup2(fdout, 1);

		/* Create new argument vector. */
		if(!(nargv = malloc(sizeof(char *)*(argc+10))))
			err(1, "malloc");
		offset = -1;
		if(sshpass){
			nargv[++offset] = "sshpass";
			nargv[++offset] = "-ePass";
		}
		nargv[++offset] = "ssh";
		nargv[++offset] = "-tt";
		if(!sshpass)
			nargv[++offset] = "-oBatchMode=yes";
		for(i = 1; i < argc; i++)
			nargv[i+offset] = argv[i];
		nargv[argc+offset] = NULL;

		signal(SIGINT, SIG_IGN);

		/* Exec into ssh. */
		execvp(nargv[0], nargv);
		err(1, "%s", nargv[0]);
	}

	if(!(fdin = open(in, O_WRONLY)))
		err(1, "open");
	if(!(fdout = open(out, O_RDONLY)))
		err(1, "open");

	/* Print initialization command. */
	if(hardcopy)
		dprintf(fdin, HINIT);
	else
		dprintf(fdin, INIT);

	FD_ZERO(&rfds0);
	FD_ZERO(&rfds1);

	tv.tv_sec = 0;
	tv.tv_usec = 1;

	for(;;){
		/*
		 * User input is read from standard in and copied to the
		 * input pipe. It is saved in bufin for later use.
		 */
		FD_SET(0, &rfds0);
		if(select(0+1, &rfds0, NULL, NULL, &tv) > 0){
			n = read(0, bufin, MAXBUF);
			bufin[n] = 0;

			/*
			 * As +x is given, escaped commands are enabled,
			 * prefixed with !. Escaped commands are run
			 * locally, but operate on remote files. Any
			 * specified remote files will be downloaded and
			 * re-uploaded when the escaped command finishes.
			 * (Escaped commands are recognized only after
			 * shell prompt.)
			 */
			if(!(escape && bufin[0] == '!'
			&& (strcmp(bufout+strlen(bufout)-2, "$ ") == 0
			|| strcmp(bufout+strlen(bufout)-2, "% ") == 0
			|| strcmp(bufout+strlen(bufout)-2, "# ") == 0))){
				dprintf(fdin, "%s", bufin);
				continue;
			}

			/*
			 * The remote server is queried for the current
			 * working directory. The response contains the
			 * query ("pwd\n"), which is skipped, followed by
			 * the working directory ("/some/path\n") and then
			 * a shell prompt, which is skipped.
			 */
			dprintf(fdin, "pwd\n");
			for(i = 0; i < MAXWD-1; i++){
				read(fdout, wd, 1);
				if(*wd == '\n') break;
			}
			n = read(fdout, wd, MAXWD);
			wd[strcspn(wd, "\n\015")] = 0;

			/* Construct corresponding argument vector. */
			bufin[strcspn(bufin, "\n")] = 0;
			p = bufin+1;
			i = 0;
			for(q = strtok(p, " "); q;
			(q = strtok(NULL, " ")), i++)
				if(i < MAXEARG-1)
					eargv[i] = q;
			eargv[i] = NULL;

			for(i = 1; eargv[i]; i++)
				if(eargv[i][0] != '-'
				&& eargv[i][0] != '+')
					goto found;
			fprintf(stderr, "no file specified\n");
			goto done;
found:
			if(eargv[i+1]){
				fprintf(stderr,
					"more than one file given\n");
				goto done;
			}

			/* Create a temporary local directory. */
			strcpy(dir, "/tmp/rtty.XXXXXX");
			if(!mkdtemp(dir)){
				warn("mkdtemp");
				continue;
			}
			chdir(dir);

			/* Construct argument to scp. */
			if(!(p = malloc(strlen(address)-1+strlen(wd)-1
			+strlen(eargv[i])-1+3)))
				err(1, "malloc");
			if(eargv[i][0] == '/'){
				sprintf(p, "%s:%s", address, eargv[i]);
			}else
				sprintf(p, "%s:%s/%s", address, wd, eargv[i]);
			q = strrchr(eargv[i], '/');
			if(q) eargv[i] = q+1;

			/*
			 * During the execution of the external commands,
			 * the SIGCHLD handler is temporarily disabled.
			 */
			signal(SIGCHLD, NULL);

			/* Download file. */
			if((child = fork()) == 0){
				execlp("scp", "scp", p, ".", NULL);
				err(1, "scp");
			}
			waitpid(child, &s, WALLSIG);
			if(WIFEXITED(s) && WEXITSTATUS(s)){
				fprintf(stderr, "scp terminated abnormally\n");
				fprintf(stderr, "file saved in %s\n", dir);
				goto done;
			}

			/* Run escaped command. */
			if((child = fork()) == 0){
				signal(SIGINT, SIG_IGN);
				execvp(eargv[0], eargv);
				err(1, "%s", eargv[0]);
			}
			waitpid(child, &s, WALLSIG);
			if(WIFEXITED(s) && WEXITSTATUS(s)){
				fprintf(stderr, "%s terminated abnormally\n",
					eargv[0]);
				fprintf(stderr, "file saved in %s\n", dir);
				goto done;
			}

			/* Upload file. */
			if((child = fork()) == 0){
				execlp("scp", "scp", eargv[i], p, NULL);
				err(1, "scp");
			}
			waitpid(child, &s, WALLSIG);
			if(WIFEXITED(s) && WEXITSTATUS(s)){
				fprintf(stderr, "scp terminated abnormally\n");
				fprintf(stderr, "file saved in %s\n", dir);
				goto done;
			}

			/* Remove temporary directory. */
			unlink(eargv[i]);
			rmdir(dir);

			/* Clean up. */
done:
			signal(SIGCHLD, killall);
			free(p);

			/*
			 * After an escaped command, an empty line is sent
			 * to the remote server, in order for a new shell
			 * prompt to be triggered.
			 */
			dprintf(fdin, "\n");
		}

		/*
		 * System output is read from the output pipe and copied
		 * to standard out for the user to see.
		 */
		FD_SET(fdout, &rfds1);
		if(select(fdout+1, &rfds1, NULL, NULL, &tv) > 0){
			n = read(fdout, bufout, MAXBUF);
			bufout[n] = 0;

			/*
			 * Bufin and bufout are copied to the temporary
			 * pointers p and q. p is incremented in order to
			 * skip any potential repetition of the command
			 * given by the user on standard in. (It is assumed
			 * that the typed command fits in bufin, i.e. is
			 * not larger than MAXBUF.)
			 */
			p = bufin;
			q = bufout;
			for(;;){
				if(*q == '\n'){
					q++;
					break;
				}
				if(*q == 13){
					q++;
					continue;
				}
				if(*p != *q){
					q = bufout;
					break;
				}
				p++; q++;
			}

			/*
			 * Bo is printed and bufin is cleared, so as not
			 * to perform the skip again, incorrectly.
			 */
			printf("%s", q);
			fflush(stdout);
			bufin[0] = 0;

			/*
			 * Unless the +d flag is given, rtty tries to detect
			 * password prompts and hide the user's input.
			 */
			if(dumb)
				continue;
			if(q[strlen(q)-1] == ':'
			|| strcmp(q+strlen(q)-2, ": ") == 0){
				if(!(q = strrchr(bufout, '\n')))
					q = bufout;
				if(strstr(q, "password")
				|| strstr(q, "Password")
				|| strstr(q, "passphrase")
				|| strstr(q, "Passphrase")
				|| strstr(q, "pass phrase")
				|| strstr(q, "Pass phrase")){
					getpw(pw);
					dprintf(fdin, "%s\n", pw);
				}
			}
		}
	}

}

char *
getpw(char *pw)
{
	struct termios orig, term;

	tcgetattr(0, &orig);
	tcgetattr(0, &term);
	term.c_lflag &= ~ECHO;
	tcsetattr(0, TCSANOW, &term);

	fgets(pw, 255, stdin);
	pw[strcspn(pw, "\n")] = 0;

	printf("\n");
	tcsetattr(0, TCSAFLUSH, &orig);

	return pw;
}

void
noop()
{
}

void
killall(int s)
{
	kill(0, SIGTERM);
	exit(1);
}