/*
 * QA Remote Shell - client side
 *
 * Run a command on the server with lots of wizz-bang options
 *
 * Copyright © 2005-2008  Red Hat, Inc. All rights reserved.
 *
 * This copyrighted material is made available to anyone wishing to use,
 * modify, copy, or redistribute it subject to the terms and conditions of the
 * GNU General Public License v.2.  This program is distributed in the hope
 * that it will be useful, but WITHOUT ANY WARRANTY expressed or implied,
 * including the implied warranties of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE. See the GNU General Public License for more details.
 * You should have received a copy of the GNU General Public License along with
 * this program; if not, write to the Free Software Foundation, Inc., 51
 * Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. Any Red Hat
 * trademarks that are incorporated in the source code or documentation are not
 * subject to the GNU General Public License and may only be used or replicated
 * with the express permission of Red Hat, Inc.
 *
 * Red Hat Author(s): Nathan Straz <nstraz@redhat.com>
 *                    Dean Jansa <djansa@redhat.com>
 */

#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 <pwd.h>
#include <time.h>
#include <stdarg.h>

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

#define QARSH_MINPORT 5010
#define QARSH_BUFSIZE QARSH_MAX_PACKET_SIZE/2
#define CONNECT_TIMEOUT 30


/* Globals */
int qarsh_fd = -1;		/* The control connection to qarshd */
unsigned short qarsh_ss_family;  /* AF_INET/AF_INET6, set on connect */
hbeat_t qarsh_hb;       /* Heartbeat handle */
int signal_to_send = 0;
int sigs_to_propogate[] = { SIGINT, SIGTERM, SIGHUP, SIGUSR1, SIGUSR2 };
sigset_t pselect_sigmask;
int connection_timeout = 0;


void
lprintf(int priority, const char *format, ...)
{
	va_list ap;

	va_start(ap, format);
	vfprintf(stderr, format, ap);
	va_end(ap);
}

void
usage()
{
	fprintf(stderr, "qarsh [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"
		     "                Env var QARSH_TIMEOUT can also be set.\n"
	);

	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)
{
	connection_timeout = 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
reset_signals(void)
{
	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_UNBLOCK, &sigmask, &pselect_sigmask);
	for (i = 0; i < n; i++) {
		signal(sigs_to_propogate[i], SIG_DFL);
	}
}

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

	qp = make_qp_setuser(user, group);
	ret = send_packet(qarsh_fd, qp);
	qpfree(qp);
	if (ret <= 0) {
		fprintf(stderr, "Failed to send set user message: %s\n", strerror(errno));
		exit(125);
	}
	qp = recv_packet(qarsh_fd);
	if (!qp) {
		fprintf(stderr, "Failed to receive response to set user\n");
		close(qarsh_fd);
		exit(125);
	} else if (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);
	}
	qpfree(qp);
}

int
run_remote_cmd(char *cmdline)
{
	struct qa_packet *qp;
	int ret;
	int rc;
	int allowed_in = 0; /* bytes we can send to qarshd */
	char b_out[QARSH_BUFSIZE], b_err[QARSH_BUFSIZE]; /* Buffers */
	char buf[QARSH_BUFSIZE];
	int z_out = 0, z_err = 0; /* size in buffer */
	int eof_out = 0, eof_err = 0;
	fd_set rfds, wfds;
	int nset;
	int nbytes;
	struct timespec timeout;
	short cmd_finished;
	short stdin_isatty = isatty(fileno(stdin));

	qp = make_qp_runcmd(cmdline);
	ret = send_packet(qarsh_fd, qp);
	qpfree(qp);
	if (ret == -1) {
		fprintf(stderr, "Failed to send runcmd packet: %s\n", strerror(errno));
		exit(125);
	}

	/* Setup signal handling stuff so we can propogate signals */
	setup_signals();

	/* Tell qarshd how much data it can send on stdout and stderr */
	qp = make_qp_data_allow(1, QARSH_BUFSIZE);
	ret = send_packet(qarsh_fd, qp);
	qpfree(qp);
	if (ret == -1) {
		fprintf(stderr, "Failed to send data allow packet: %s\n", strerror(errno));
		exit(125);
	}
	qp = make_qp_data_allow(2, QARSH_BUFSIZE);
	ret = send_packet(qarsh_fd, qp);
	qpfree(qp);
	if (ret == -1) {
		fprintf(stderr, "Failed to send data allow packet: %s\n", strerror(errno));
		exit(125);
	}

	hbeat(qarsh_hb);
	cmd_finished = 0;
	for (;;) {
		if (cmd_finished) {
			timeout.tv_sec = 1;
			timeout.tv_nsec = 0;
		} else {
			timeout.tv_sec = 5;
			timeout.tv_nsec = 0;
		}
		FD_ZERO(&rfds);
		FD_ZERO(&wfds);
		FD_SET(qarsh_fd, &rfds);
		/* allowed_in is set to -1 when stdin is closed */
		if (allowed_in != -1) {
			/* close stdin if it is a tty */
			if (stdin_isatty) {
				qp = make_qp_data(0, 0, 0, NULL);
				ret = send_packet(qarsh_fd, qp);
				if (ret == -1) {
					fprintf(stderr, "Failed to send data packet: %s\n", strerror(errno));
					exit(125);
				}
				qpfree(qp);
				close(fileno(stdin));
				allowed_in = -1;
			} else {
				/* Only set stdin if we're allowed to send some data, otherwise we
				 * could read 0 bytes and close stdin before we actually get eof */
				if (allowed_in > 0) FD_SET(fileno(stdin), &rfds);
			}
		}
		if (z_out) { FD_SET(fileno(stdout), &wfds); }
		if (z_err) { FD_SET(fileno(stderr), &wfds); }

		nset = pselect(qarsh_fd+1, &rfds, &wfds, NULL, &timeout, 
				&pselect_sigmask);

		/* Timeout hit, send a heartbeat */
		if (nset == 0) {
			if (!hbeat(qarsh_hb)) {
				/* If the heartbeat fails, we should exit now.
				 * The hbeat state will take precedence over the
				 * exit status of the command. */
				break;
			}
		}

		if (nset == -1 && errno == EINTR) {
			/* Only test signals */
			if (signal_to_send) {
				qp = make_qp_kill(signal_to_send);
				ret = send_packet(qarsh_fd, qp);
				if (ret == -1) {
					fprintf(stderr, "Failed to send kill packet: %s\n", strerror(errno));
					exit(125);
				}
				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), &rfds)) {
				nbytes = read(fileno(stdin), buf, allowed_in);
				if (nbytes >= 0) {
					qp = make_qp_data(0, 0, nbytes, buf);
					ret = send_packet(qarsh_fd, qp);
					if (ret == -1) {
						fprintf(stderr, "Failed to send data packet: %s\n", strerror(errno));
						exit(125);
					}
					qpfree(qp);
					allowed_in -= nbytes;
					if (nbytes == 0) {
						close(fileno(stdin));
						allowed_in = -1;
					}
				}
				nset--;
			}
			if (nset && FD_ISSET(fileno(stdout), &wfds)) {
				nbytes = write(fileno(stdout), b_out, z_out);
				if (nbytes == z_out) {
					z_out = 0;
					if (eof_out) close(fileno(stdout));
				} else if (nbytes == -1) {
					fprintf(stderr, "Error writing to stdout: %s\n", strerror(errno));
				} else {
					memmove(b_out, b_out+nbytes, z_out - nbytes);
					z_out -= nbytes;
				}
				if (!eof_out) {
					qp = make_qp_data_allow(1, nbytes);
					ret = send_packet(qarsh_fd, qp);
					if (ret == -1) {
						fprintf(stderr, "Failed to send data allow packet: %s\n", strerror(errno));
						exit(125);
					}
					qpfree(qp);
				}
				nset--;
			}
			if (nset && FD_ISSET(fileno(stderr), &wfds)) {
				nbytes = write(fileno(stderr), b_err, z_err);
				if (nbytes == z_err) {
					z_err = 0;
					if (eof_err) close(fileno(stderr));
				} else {
					memmove(b_err, b_err+nbytes, z_err - nbytes);
					z_err -= nbytes;
				}
				if (!eof_err) {
					qp = make_qp_data_allow(2, nbytes);
					ret = send_packet(qarsh_fd, qp);
					if (ret == -1) {
						fprintf(stderr, "Failed to send data allow packet: %s\n", strerror(errno));
						exit(125);
					}
					qpfree(qp);
				}
				nset--;
			}
			if (nset && FD_ISSET(qarsh_fd, &rfds)) {
				qp = recv_packet(qarsh_fd);
				if (qp == NULL) {
					fprintf(stderr, "recv_packet() returned NULL!\n:");
					break;
				}
			
				if (qp->qp_type == QP_CMDEXIT) {
					cmd_finished = 1;
					rc = qp->qp_cmdexit.qp_status;
					/* Don't break yet, we need to make
					 * sure all output is read. */
				} else if (qp->qp_type == QP_DATA) {
					if (qp->qp_data.qp_remfd == 1 && qp->qp_data.qp_count <= (QARSH_BUFSIZE - z_out)) {
						if (qp->qp_data.qp_count == 0) eof_out = 1;
						memcpy(b_out+z_out, qp->qp_data.qp_blob, qp->qp_data.qp_count);
						z_out += qp->qp_data.qp_count;
					} else if (qp->qp_data.qp_remfd == 2 && qp->qp_data.qp_count <= (QARSH_BUFSIZE - z_err)) {
						if (qp->qp_data.qp_count == 0) eof_err = 1;
						memcpy(b_err+z_err, qp->qp_data.qp_blob, qp->qp_data.qp_count);
						z_err += qp->qp_data.qp_count;
					} else {
						fprintf(stderr, "ERROR: Bad data packet: fd %d, cnt: %d\n, bufleft: %d, %d",
								qp->qp_data.qp_remfd, qp->qp_data.qp_count,
								QARSH_BUFSIZE - z_out, QARSH_BUFSIZE - z_err);
					}
				} else if (qp->qp_type == QP_DALLOW) {
					if (qp->qp_dallow.qp_remfd == 0) {
						/* If we already closed stdin, don't change allowed_in */
						if (allowed_in != -1) allowed_in += qp->qp_dallow.qp_count;
					} else {
						fprintf(stderr, "ERROR: Received data allow for fd %d\n",
								qp->qp_dallow.qp_remfd);
					}
				} else if (qp->qp_type == QP_RETURNCODE) { /* qarshd hit an error while writing to stdin */
					if (qp->qp_returncode.qp_rc == -1 && qp->qp_returncode.qp_errno == EPIPE) {
						close(fileno(stdin));
						allowed_in = -1;
					} else {
						fprintf(stderr, "Remote command hit I/O error: %s\n", qp->qp_returncode.qp_strerror);
					}
				} else {
					fprintf(stderr, "Unexpected packet type %s\n", qp_packet_type(qp->qp_type));
				}
				qpfree(qp);
				nset--;
			}
		}
		if (cmd_finished 
				&& eof_out && z_out == 0 
				&& eof_err && z_err == 0) {
			/* If the command is complete, we've seen EOF
			 * on outputs and both output buffers are empty we can
			 * exit now.  We need to test all conditions at once so
			 * if none are true, we'll still check for heartbeat. */
			break;
		}
	}
	if (hbeat_getstate(qarsh_hb) == HOST_TIMEOUT) {
		fprintf(stderr, "Didn't receive heartbeat for %d seconds\n", 
				hbeat_getmaxtimeout(qarsh_hb));
		return W_EXITCODE(127, 0);
	} else if (hbeat_getstate(qarsh_hb) == HOST_REBOOT) {
		fprintf(stderr, "Remote host rebooted\n");
		return W_EXITCODE(127, 0);
	}
	if (!cmd_finished) {
		fprintf(stderr, "Remote command exited with unknown state\n");
		return W_EXITCODE(127, 0);
	}
	return rc;
}

int 
main(int argc, char *argv[])
{
	int c;
	int port = 5016;
	char *host = NULL;
	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;

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

again:
	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();
			exit(1);
	} }
	/* Some programs (rsync) put the hostname before some qarsh options
	 * We need to pull the hostname out and if there are still args
	 * keep trying to parse them.  (code from OpenSSH)
	 */
	argc -= optind;
	argv += optind;

	if (argc > 0 && !host && **argv != '-') {
		host = *argv;
		if (argc > 1) {
			optind = 1;
			goto again;
		}
		argc--;
		argv++;
	}

	if (!host) {
		usage();
		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 (remuser == NULL) {
		if (!(pw = getpwuid(getuid()))) {
			fprintf(stderr, "qarsh: can not look up local username.\n");
			exit(1);
		}
		remuser = pw->pw_name;
	}

	if ((args = copyargs(argv)) == NULL) {
		usage();
		exit(1);
	}
	/* close fd above stderr which parent may have left open */
	for (c = 3; c < 30; c++) {
		close(c);
	}

	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(CONNECT_TIMEOUT);
	qarsh_fd = connect_to_host(host, port, &qarsh_ss_family);
	alarm(0);
	if (connection_timeout) {
		fprintf(stderr, "Could not connect to remote host\n");
		exit(127);
	}

	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);
	hbeat_free(qarsh_hb);
	free(args);
	/* If the remote cmd was killed, we need to be killed too */
	if (WIFSIGNALED(ret)) {
		reset_signals();
		raise(WTERMSIG(ret));
	}
	/* Otherwise we need to exit with the same exit status */
	return WEXITSTATUS(ret);
}