#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 DPUTS(d, s) write(d, s, strlen(s))

#define INIT "export TERM=tty43 EDITOR=ed PAGER=cat\n"
#define IINIT "export EDITOR=ed PAGER=ul\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, i, interactive, n, offset,
		s, sshpass;
	struct timeval	tv;

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

	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");

	escape = dumb = interactive = sshpass = 0;
	for(;argv[1][0] == '+';){
		if(strchr(argv[1], 'd')) dumb = 1;
		if(strchr(argv[1], 'i')) interactive = 1;
		if(strchr(argv[1], 'x')) escape = 1;
		if(strchr(argv[1], 'p')) sshpass = 1;

		argv[1] = argv[0];
		argv++; argc--;
	}

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

	if(sshpass){
		printf("password: ");
		getpw(pw);
		setenv("SSHPASS", pw, 1);
	}

	/* Start ssh. */
	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);

		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);

		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");

	/* Initialize remote environment. */
	if(interactive)
		DPUTS(fdin, IINIT);
	else
		DPUTS(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;

			/*
			 * Escaped commands are recognized only after a
			 * remote shell prompt.
			 */
			if(escape && bufin[0] == '!'
			&& (strcmp(bufout+strlen(bufout)-2, "$ ") == 0
			|| strcmp(bufout+strlen(bufout)-2, "% ") == 0
			|| strcmp(bufout+strlen(bufout)-2, "# ") == 0))
				goto escape;

			DPUTS(fdin, bufin);
			continue;

escape:
			/*
			 * Get remote working directory. (The response
			 * looks like "pwd\n(directory)\n(prompt)".)
			 */
			DPUTS(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;

			bufin[strcspn(bufin, "\n")] = 0;
			i = 0;
			p = bufin+1;
			q = strtok(p, " ");
			for(; 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;
			}

			strcpy(dir, "/tmp/rtty.XXXXXX");
			if(!mkdtemp(dir)){
				warn("mkdtemp");
				continue;
			}
			chdir(dir);

			/* Construct remote path for 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;

			signal(SIGCHLD, NULL);

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

			unlink(eargv[i]);
			rmdir(dir);

done:
			signal(SIGCHLD, killall);
			free(p);

			/* Trigger new remote shell prompt. */
			DPUTS(fdin, "\n");
		}

		/*
		 * Remote 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;

			/*
			 * The server's echo of the client's command is
			 * skipped, assuming that bufin contains the
			 * command in its entirety.
			 */
			i = 0;
			p = bufin;
			q = bufout;
			for(;;){
				if(*q == '\n'){
					q++;
					break;
				}
				if(*q == '\r'){
					q++;
					continue;
				}
				if(*p == '\t' && strspn(q, " ") >= 8-i%8){
					p++; q += 8-i%8; i += 8-i%8;
					continue;
				}
				if(*p != *q){
					q = bufout;
					break;
				}
				p++; q++; i++;
			}

			/* Don't skip the echo again. */
			printf("%s", q);
			fflush(stdout);
			bufin[0] = 0;

			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);
					DPUTS(fdin, pw);
					DPUTS(fdin, "\n");
				}
			}
		}
	}
}

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);
}