/*
 * QA Remote Shell - client side
 *
 * Run a command on the server with lots of wizz-bang options
 *
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <netdb.h>
#include <syslog.h>
#include <pwd.h>
#include <time.h>

#include "sockutil.h"
#include "qarsh_packet.h"
#include "btime.h"
#include "hbeat.h"

#define QARSH_MINPORT 5010


/* Globals */
int qarsh_fd = -1;		/* The control connection to qarshd */
hbeat_t qarsh_hb;       /* Heartbeat handle */
int signal_to_send = 0;
int sigs_to_propogate[] = { SIGINT, SIGTERM, SIGHUP, SIGUSR1, SIGUSR2 };
sigset_t pselect_sigmask;

void
usage(const char *pname)
{
	fprintf(stderr, "%s [options] [user[.group]@]hostname cmdline ...\n"
		     "-l user         Run cmdline using this user name.\n"
		     "-g group        Run cmdline using this group name.\n"
		     "-p port         Use this port to contact qarshd.\n"
			 "-t timeout      Number of seconds a remote host can be\n"
			 "                silent before we give up and exit\n"
		     "                A value of 0 disables heartbeating.\n"
		     "                Default is value is 120.\n"

	,pname);

	return;
}

char *
copyargs(char **argv)
{
	int cc;
	char **ap, *p;
	char *args;

	cc = 0;
	for (ap = argv; *ap; ++ap)
		cc += strlen(*ap) + 1;
	if (cc == 0) return NULL;
	args = malloc(cc);
	if (!args) {
		perror("qarsh: malloc failed in copyargs");
		exit(1);
	}
	for (p = args, ap = argv; *ap; ++ap) {
		for (p = strcpy(p, *ap); *p; ++p);
		if (ap[1])
			*p++ = ' ';
	}
	return args;
}

void
sig_handler(int sig)
{
	signal_to_send = sig;
}

void 
sig_alrm_handler(int sig)
{
	fprintf(stderr, "No heartbeat from remote host\n");
	exit(1);
}

void
setup_signals(void)
{
	struct sigaction sa;
	sigset_t sigmask;
	int i, n;

	n = sizeof sigs_to_propogate / sizeof *sigs_to_propogate;
	sigemptyset(&sigmask);
	for (i = 0; i < n; i++) {
		sigaddset(&sigmask, sigs_to_propogate[i]);
	}
	sigprocmask(SIG_BLOCK, &sigmask, &pselect_sigmask);
	sa.sa_handler = sig_handler;
	sa.sa_mask = sigmask;
	sa.sa_flags = SA_RESTART;
	for (i = 0; i < n; i++) {
		sigaction(sigs_to_propogate[i], &sa, NULL);
	}
}

void
set_remote_user(char *user, char *group)
{
	struct qa_packet *qp;

	qp = make_qp_setuser(user, group);
	qp->qp_seq = 1;
	send_packet(qarsh_fd, qp);
	qpfree(qp);
	qp = recv_packet(qarsh_fd);
	if (qp && qp->qp_type == QP_RETURNCODE 
			&& qp->qp_returncode.qp_rc == -1) {
		
		fprintf(stderr, "Remote side failed, %s\n",
				qp->qp_returncode.qp_strerror);
		close(qarsh_fd);
		exit(125);
	}
}

int
run_remote_cmd(char *cmdline)
{
	struct qa_packet *qp;
	char *buf;
	int bufsize;
	int rc;
	int p_in, p_out, p_err;  /* Port numbers */
	int l_in, l_out, l_err;  /* listening sockets */
	int c_in, c_out, c_err;  /* client sockets */
	fd_set readfds, testfds;
	int nset;
	struct sockaddr_in caddr;
	socklen_t clen;
	struct timespec timeout;
	short cmd_finished;

	l_in = bind_any(QARSH_MINPORT);
	p_in = getsockport(l_in);
	l_out = bind_any(QARSH_MINPORT);
	p_out = getsockport(l_out);
	l_err = bind_any(QARSH_MINPORT);
	p_err = getsockport(l_err);
	
	qp = make_qp_runcmd(cmdline, p_in, p_out, p_err);
	qp->qp_seq = 1;
	send_packet(qarsh_fd, qp);
	qpfree(qp);

	/* Get the stdin, stdout, and stderr connections up before we do work */
	FD_ZERO(&readfds);
	FD_SET(l_in, &readfds);
	FD_SET(l_out, &readfds);
	FD_SET(l_err, &readfds);
	c_in = c_out = c_err = 0;

	do {
		testfds = readfds;
		nset = select(FD_SETSIZE, &testfds, NULL, NULL, NULL);

		if (FD_ISSET(l_in, &testfds)) {
			clen = sizeof caddr;
			c_in = accept(l_in, (struct sockaddr *)&caddr, &clen);
			if (c_in == -1) {
				fprintf(stderr, 
					"accept of l_in failed, %d: %s\n",
				 		errno, strerror(errno));
				continue;
			}
		}
		if (FD_ISSET(l_out, &testfds)) {
			clen = sizeof caddr;
			c_out = accept(l_out, (struct sockaddr *)&caddr, &clen);
			if (c_out == -1) {
				fprintf(stderr, 
					"accept of l_out failed, %d: %s\n",
				 		errno, strerror(errno));
				continue;
			}
		}
		if (FD_ISSET(l_err, &testfds)) {
			clen = sizeof caddr;
			c_err = accept(l_err, (struct sockaddr *)&caddr, &clen);
			if (c_err == -1) {
				fprintf(stderr, 
					"accept of l_err failed, %d: %s\n",
				 		errno, strerror(errno));
				continue;
			}
		}
	} while (c_in == 0 || c_out == 0 || c_err == 0);
	close(l_in);
	close(l_out);
	close(l_err);
	l_in = l_out = l_err = -1;

	/* Now we can start doing some real work */
	
	FD_ZERO(&readfds);
	FD_SET(qarsh_fd, &readfds);
	FD_SET(c_out, &readfds);
	FD_SET(c_err, &readfds);
	FD_SET(fileno(stdin), &readfds);
	/* Setup signal handling stuff so we can propogate signals */
	setup_signals();

	if (fcntl(fileno(stdin), F_SETFL, O_NONBLOCK) != 0) {
		fprintf(stderr, 
			"fcntl stdin O_NONBLOCK failed, %d: %s\n", 
			errno, strerror(errno));
	}
	buf = malloc(1024);
	memset(buf, 0, 1024);

	cmd_finished = 0;
	for (;;) {
		if (cmd_finished) {
			timeout.tv_sec = 0;
			timeout.tv_nsec = 0;
		} else {
			timeout.tv_sec = 5;
			timeout.tv_nsec = 0;
		}
		testfds = readfds;
		memset(buf, 0, 1024);

		nset = pselect(FD_SETSIZE, &testfds, NULL, NULL, &timeout, 
				&pselect_sigmask);

		if (nset == 0) {
			if (!hbeat(qarsh_hb)) {
				fprintf(stderr, "No heartbeat from remote host");
				/* Set our return packet as NULL so we exit 
				 * with unknown error. */
				qp = NULL;
				break;
			}	

			if (cmd_finished) {
				break;
			}
	
			continue;
		}

		if (nset == -1 && errno == EINTR) {
			/* Only test signals */
			if (signal_to_send) {
				qp = make_qp_kill(signal_to_send);
				send_packet(qarsh_fd, qp);
				qpfree(qp);
				signal_to_send = 0;
			}
		} else if (nset > 0) {
			/* We got traffic, poke the state into the hbeat because
			 * we may not have gotten the chance to call hbeat() above
			 * which would normally reset the state for us.
			 */
			hbeat_setstate(qarsh_hb, HOST_ALIVE);
			
			if (nset && FD_ISSET(fileno(stdin), &testfds)) {
				bufsize = read(fileno(stdin), buf, 1024);
				if (bufsize > 0) {
					write(c_in, buf, bufsize);
				} else if (bufsize == 0) {
					FD_CLR(fileno(stdin), &readfds);
					close(fileno(stdin));
					close(c_in);
					c_in = 0;
				}
				nset--;
			}
			if (nset && c_out && FD_ISSET(c_out, &testfds)) {
				bufsize = read(c_out, buf, 1024);
				if (bufsize > 0) {
					write(fileno(stdout), buf, bufsize);
				} else if (bufsize == 0) {
					FD_CLR(c_out, &readfds);
					close(c_out);
					c_out = 0;
				}
				nset--;
			}
			if (nset && c_err && FD_ISSET(c_err, &testfds)) {
				bufsize = read(c_err, buf, 1024);
				if (bufsize > 0) {
					write(fileno(stderr), buf, bufsize);
				} else if (bufsize == 0) {
					FD_CLR(c_err, &readfds);
					close(c_err);
					c_err = 0;
				}
				nset--;
			}
			if (nset && FD_ISSET(qarsh_fd, &testfds)) {
				qp = recv_packet(qarsh_fd);
				if (qp == NULL) {
					fprintf(stderr, "recv_packet() returned NULL!\n:");
					break;
				}
			
				/* dump_qp(qp); */
				if (qp && qp->qp_type == QP_CMDEXIT) {
					cmd_finished = 1;
					/* break; */
				}
				nset--;
			}
		}
	}
	if (c_out) close(c_out);
	if (c_err) close(c_err);
	if (qp == NULL) {
		fprintf(stderr, "Remote command exited with unknown state\n");
		free(buf);
		return 127;
	}
	if (WIFSIGNALED(qp->qp_cmdexit.qp_status)) {
		rc = 128 + WTERMSIG(qp->qp_cmdexit.qp_status);
	} else {
		rc = WEXITSTATUS(qp->qp_cmdexit.qp_status);
	}
	qpfree(qp);
	free(buf);
	return rc;
}

int 
main(int argc, char *argv[])
{
	int c;
	int port = 5008;
	char *host;
	char *remuser = NULL;
	char *remgroup = NULL;
	char *args;
	struct passwd *pw;
	int ret;
	struct sigaction sa;
	sigset_t sigmask;
	char *cp;
	int max_timeout = 120;

	openlog("qarsh", LOG_PID, LOG_DAEMON);

	if ((cp = getenv("QARSH_TIMEOUT")) != NULL) {
		max_timeout = atoi(cp);
	}

	while ((c = getopt(argc, argv, "+p:l:g:t:")) != -1) {
	switch (c) {
		case 'l':
			remuser = strdup(optarg);
			break;
		case 'g':
			remgroup = strdup(optarg);
			break;
		case 'p':
			port = atoi(optarg);
			break;
		case 't':
			max_timeout = atoi(optarg);
			break;
		case '?':
		default:
			printf("Unknown option %c\n", (char)optopt);
			usage(argv[0]);
			exit(1);
	} }

	if ((host = argv[optind++]) == NULL) {
		usage(argv[0]);
		exit(1);
	}

	/* check for user and group in form [user[.group]@]hostname */
	{
		char *sp;

		if ((sp = strchr(host, '@'))) {
			remuser = host;
			host = sp+1;
			*sp = '\0';
		}
		if (remuser && (sp = strchr(remuser, '.'))) {
			remgroup = sp+1;
			*sp = '\0';
		}
	}
	if (!(pw = getpwuid(getuid()))) {
		fprintf(stderr, "qarsh: unknown user id.\n");
		exit(1);
	}
	if (remuser == NULL) {
		remuser = strdup(pw->pw_name);
	}

	argc -= optind;
	argv += optind;
	if ((args = copyargs(argv)) == NULL) {
		usage(argv[0]);
		exit(1);
	}

	memset(&sa, 0, sizeof sa);
	sigemptyset(&sigmask);
	sa.sa_mask = sigmask;
	sa.sa_handler = SIG_IGN;
	sigaction(SIGTTIN, &sa, NULL);

	memset(&sa, 0, sizeof sa);
	sigemptyset(&sigmask);
	sa.sa_mask = sigmask;
	sa.sa_flags = SA_ONESHOT;
	sa.sa_handler = sig_alrm_handler;
	sigaction(SIGALRM, &sa, NULL);

	alarm(max_timeout);
	qarsh_fd = connect_to_host(host, port);
	alarm(0);

	if (qarsh_fd == -1) {
		if (errno == 0) {
			fprintf(stderr, "Could not connect to %s:%d, %d: %s\n",
				host, port, h_errno, hstrerror(h_errno));
		} else {
			fprintf(stderr, "Could not connect to %s:%d, %d: %s\n",
				host, port, errno, strerror(errno));
		}
		return 127;
	}

	qarsh_hb = hbeat_init(host, max_timeout);
	if (!qarsh_hb) {
		fprintf(stderr, "Could not init heartbeat to %s\n", host);
		return 127;
	}

	set_remote_user(remuser, remgroup);

	ret = run_remote_cmd(args);
	close(qarsh_fd);
	free(args);
	return ret;
}