/*
 * 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 <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <assert.h>
#include <syslog.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/sendfile.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <pwd.h>
#include <grp.h>
#include <stdarg.h>


#include "sockutil.h"
#include "qarsh_packet.h"

#define QARSHD_BUFSIZE QARSH_MAX_PACKET_SIZE/2
/*
 * QA Remote Shell Daemon
 */

int debug = 0;
int dopause = 0;

/* Globals */
const int qinfd = 0; /* qarshd in file descriptor */
const int qoutfd = 1; /* qarshd out file descriptor */
pid_t child_pid;
sigset_t orig_sigmask;
int childfds[3] = { -1, -1, -1 }; /* pipes to child for stdin/stdout/stderr */
int receivefd = -1;
int sendfd = -1;

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

	va_start(ap, format);
	vsnprintf(buf, 4096, format, ap);
	va_end(ap);

	syslog(priority, "%s", buf);
}

int
setup_user(char *user, char *group)
{
	struct passwd *pw;
	struct group *grp;

	if ((pw = getpwnam(user)) == NULL) {
		syslog(LOG_WARNING, "User \"%s\" not found\n", user);
		return 0;
	}

	if (group && (grp = getgrnam(group))) {
		setgid(grp->gr_gid);
		initgroups(pw->pw_name, grp->gr_gid);
	} else {
		setgid(pw->pw_gid);
		initgroups(pw->pw_name, pw->pw_gid);
	}

	/* setuid is last so we can still do setgid and initgroups */
	setuid(pw->pw_uid);

	if (pw->pw_dir) chdir(pw->pw_dir);
	
	return 1;
}

void
sig_handler(int sig)
{
}

pid_t 
run_cmd(const char *cmd)
{
	pid_t pid;
	int pipefds[2], parentfds[3];

	pipe(pipefds);
	childfds[0] = pipefds[1]; parentfds[0] = pipefds[0]; /* stdin */
	pipe(pipefds);
	childfds[1] = pipefds[0]; parentfds[1] = pipefds[1]; /* stdout */
	pipe(pipefds);
	childfds[2] = pipefds[0]; parentfds[2] = pipefds[1]; /* stderr */


	syslog(LOG_INFO, "Running cmdline: %s\n", cmd);
	if ((pid = fork()) < 0) {
		syslog(LOG_WARNING, 
			"Could not fork() in run_cmd(): %s\n", strerror(errno));
		exit(1);
	}

	if (pid == 0) { /* child */
		setpgrp();

		/* Connect stdin, stdout, and stderr to parent's pipes */
		dup2(parentfds[0], fileno(stdin));
		dup2(parentfds[1], fileno(stdout));
		dup2(parentfds[2], fileno(stderr));
		/* close end of pipes we're not using */
		close(parentfds[0]);
		close(parentfds[1]);
		close(parentfds[2]);
		close(childfds[0]);
		close(childfds[1]);
		close(childfds[2]);

		execlp("sh", "sh", "-c", cmd, NULL);
		fprintf(stderr, "exec of %s failed: %d, %s\n", cmd, errno, strerror(errno));
		exit(127);
	} 
	close(parentfds[0]);
	close(parentfds[1]);
	close(parentfds[2]);
	return pid;
}

struct qa_packet *
prepare_recvfile(struct qa_packet *qp)
{
	int fd;

	syslog(LOG_INFO, "Receive file %s, size = %lld, mode = %o\n",
		qp->qp_recvfile.qp_path,
		(long long int)qp->qp_recvfile.qp_count,
		qp->qp_recvfile.qp_mode);

	if ((fd = open(qp->qp_recvfile.qp_path, O_TRUNC|O_CREAT|O_WRONLY,
					qp->qp_recvfile.qp_mode)) < 0) {
		syslog(LOG_WARNING, "Could not open %s to receive file: %s\n",
				qp->qp_recvfile.qp_path, strerror(errno));
		return make_qp_returncode(-1, errno, strerror(errno));
	}

	/* Set permissions again to override umask */
	fchmod(fd, qp->qp_recvfile.qp_mode);

	/* Store fd to check data packets against */
	receivefd = fd;

	return make_qp_data_allow(0, QARSHD_BUFSIZE);
}

struct qa_packet *
receive_data(struct qa_packet *qp)
{
	ssize_t ret, nwrote;

	if (qp->qp_data.qp_count == 0) { /* EOF */
		close(receivefd);
		receivefd = -1;
		if (debug) syslog(LOG_DEBUG, "Transfer complete\n");
		return make_qp_returncode(0, 0, "Transfer complete");
	}
	if (debug) syslog(LOG_DEBUG, "Data for %d, %d\n", receivefd, qp->qp_data.qp_count);
	nwrote = 0;
	do {
		ret = write(receivefd, qp->qp_data.qp_blob + nwrote, qp->qp_data.qp_count - nwrote);
		if (ret < 0) {
			return make_qp_returncode(-1, errno, strerror(errno));
		} 
		nwrote += ret;
	} while (nwrote < qp->qp_data.qp_count);
	return make_qp_data_allow(0, nwrote);
}

struct qa_packet *
prepare_sendfile(struct qa_packet *qp)
{
	int fd;

	syslog(LOG_INFO, "Sending file %s\n", qp->qp_sendfile.qp_path);

	if ((fd = open(qp->qp_sendfile.qp_path, O_RDONLY)) == -1) {
		syslog(LOG_WARNING, "Could not open %s to send file: %s\n",
				qp->qp_sendfile.qp_path, strerror(errno));
		return make_qp_returncode(-1, errno, strerror(errno));
	}

	sendfd = fd;
	return NULL;
}

struct qa_packet *
rstat(const char *path)
{
	struct stat sb;
	struct qa_packet *rp;

	if (stat(path, &sb) < 0) {
		rp = make_qp_returncode(-1, errno, strerror(errno));
	} else {
		rp =  make_qp_rstat(path, &sb);
	}

	return rp;
}

/* Handle only qarsh related packets */
void
handle_qarsh()
{
	fd_set rfds, wfds;
	int nfd, maxfd;
	struct timeval timeout;
	int child_status;
	struct qa_packet *qp = NULL, *rp = NULL;
	int allowed_out = 0, allowed_err = 0; /* number of bytes we can send to client */
	char buf[QARSHD_BUFSIZE];		/* short term buffer for stdout and stderr */
	char buf_in[QARSHD_BUFSIZE];	/* long term buffer for stdin */
	int z_in = 0;		/* Number of bytes in stdin buffer */
	int eof_in = 0;		/* Have we seen EOF on stdin yet? */
	int nbytes;
	struct sigaction chldact;

	signal(SIGPIPE, SIG_IGN);
	chldact.sa_handler = sig_handler;
	chldact.sa_flags = 0; /* Make sure SA_RESTART is NOT set so select doesn't timeout on SIGCHLD */
	sigaction(SIGCHLD, &chldact, NULL);

	qp = make_qp_data_allow(0, QARSHD_BUFSIZE);
	send_packet(qoutfd, qp);
	qpfree(qp);
	
	for (;;) {

		timeout.tv_sec = 3;
		timeout.tv_usec = 0;
		FD_ZERO(&rfds); FD_ZERO(&wfds);
		FD_SET(qinfd, &rfds);
		maxfd = qinfd;
		if (childfds[0] > -1 && z_in) {
			FD_SET(childfds[0], &wfds);
			maxfd = childfds[0] > maxfd ? childfds[0] : maxfd;
		}
		if (childfds[1] > -1 && allowed_out) {
			FD_SET(childfds[1], &rfds);
			maxfd = childfds[1] > maxfd ? childfds[1] : maxfd;
		}
		if (childfds[2] > -1 && allowed_err) {
			FD_SET(childfds[2], &rfds);
			maxfd = childfds[2] > maxfd ? childfds[2] : maxfd;
		}

		/* Check the child just before we do the select so there's no time
		 * for the child to exit before the select, causing the select to hang */
		if (child_pid && waitpid(child_pid, &child_status, WNOHANG) == child_pid) {
			rp = make_qp_cmdexit(child_pid, child_status);
			send_packet(qoutfd, rp);
			qpfree(rp);
			child_pid = 0;
		}

		nfd = select(maxfd+1, &rfds, &wfds, NULL, &timeout);
		if (nfd < 0) {
			if (errno == EINTR) {
				/* signals handled above here */
				continue;
			} else {
				syslog(LOG_ERR, "select errno %d, %s\n", 
				       errno, strerror(errno));
			}
		} else if (nfd > 0) {
			if (nfd && FD_ISSET(qinfd, &rfds)) {
				qp = recv_packet(qinfd);
				if (qp == NULL) {
					if (debug) syslog(LOG_DEBUG, "That's enough\n");
					break;
				}
				if (debug) dump_qp(qp);
				switch(qp->qp_type) {
				case QP_KILL:
					if (child_pid) {
						syslog(LOG_INFO, "Sending child %d signal %d", 
							child_pid, qp->qp_kill.qp_sig);
						kill(child_pid * -1, qp->qp_kill.qp_sig);
					}
					break;
				case QP_RUNCMD:
					syslog(LOG_ERR, "Only one command allowed per connection");
					break;
				case QP_DATA:
					if (qp->qp_data.qp_remfd != 0) {
						syslog(LOG_ERR, "Received data for fd %d\n", qp->qp_data.qp_remfd);
						break;
					}
					if (qp->qp_data.qp_count > (QARSHD_BUFSIZE - z_in)) {
						syslog(LOG_ERR, "Received too much data for fd %d, %d > %d\n",
								qp->qp_data.qp_remfd, qp->qp_data.qp_count, QARSHD_BUFSIZE - z_in);
						break;
					}
					if (eof_in) {
						syslog(LOG_ERR, "Received data on stdin after EOF\n");
						break;
					}
					if (qp->qp_data.qp_count == 0) {
						eof_in = 1;
						if (z_in == 0) {
							close(childfds[0]);
							childfds[0] = -1;
						}
					} else {
						memcpy(buf_in+z_in, qp->qp_data.qp_blob, qp->qp_data.qp_count);
						z_in += qp->qp_data.qp_count;
					}
					break;
				case QP_DALLOW:
					if (qp->qp_dallow.qp_remfd == 1) {
						allowed_out += qp->qp_dallow.qp_count;
					} else if (qp->qp_dallow.qp_remfd == 2) {
						allowed_err += qp->qp_dallow.qp_count;
					} else {
						syslog(LOG_ERR, "Received data allow for fd %d\n", qp->qp_dallow.qp_remfd);
					}
					break;
				default:
					syslog(LOG_ERR, "Unexpected message type %s", qp_packet_type(qp->qp_type));
					break;
				}
				nfd--;
				qpfree(qp);
			}
			if (nfd && childfds[0] != -1 && FD_ISSET(childfds[0], &wfds)) {
				/* Child is ready for data on stdin */
				nbytes = write(childfds[0], buf_in, z_in);
				if (nbytes == -1) {  /* Don't try writing to the child again */
					close(childfds[0]);
					childfds[0] = -1;
					qp = make_qp_returncode(-1, errno, strerror(errno));
					send_packet(qoutfd, qp);
					qpfree(qp);
				} else if (nbytes == z_in) {
					z_in = 0;
					if (eof_in) {
						close(childfds[0]);
						childfds[0] = -1;
					}
				} else {
					memmove(buf_in, buf_in+nbytes, z_in - nbytes);
					z_in -= nbytes;
				}
				if (!eof_in && nbytes > 0) {
					qp = make_qp_data_allow(0, nbytes);
					send_packet(qoutfd, qp);
					qpfree(qp);
				}
				nfd--;
			}
			if (nfd && childfds[1] != -1 && FD_ISSET(childfds[1], &rfds)) {
				/* Child has something to send to stdout */
				nbytes = read(childfds[1], buf, allowed_out);
				qp = make_qp_data(1, nbytes, buf);
				send_packet(qoutfd, qp);
				qpfree(qp);
				allowed_out -= nbytes;
				if (nbytes == 0) {
					close(childfds[1]);
					childfds[1] = -1;
					allowed_out = 0;
				}
				nfd--;
			}
			if (nfd && childfds[2] != -1 && FD_ISSET(childfds[2], &rfds)) {
				/* Child has something to send to stderr */
				nbytes = read(childfds[2], buf, allowed_err);
				qp = make_qp_data(2, nbytes, buf);
				send_packet(qoutfd, qp);
				qpfree(qp);
				allowed_err -= nbytes;
				if (nbytes == 0) {
					close(childfds[2]);
					childfds[2] = -1;
					allowed_err = 0;
				}
				nfd--;
			}

		} else {
			if (debug) syslog(LOG_DEBUG, "Nothing to do in handle_qarsh\n");
		}
	}
}

void
handle_packets()
{
	fd_set rfds;
	int nfd;
	struct timeval timeout;
	struct qa_packet *qp = NULL, *rp = NULL;
	const int bufsize = QARSH_MAX_PACKET_SIZE/2;
	char buf[bufsize];
	off_t nbytes;	

	if (dopause) {
		signal(SIGALRM, sig_handler);
		pause();
		signal(SIGALRM, SIG_DFL);
	}

	for (;;) {
		FD_SET(qinfd, &rfds);
		timeout.tv_sec = 3;
		timeout.tv_usec = 0;

		nfd = select(qinfd+1, &rfds, NULL, NULL, &timeout);
		if (nfd < 0) {
			if (errno == EINTR) {
				/* signals handled above here */
				continue;
			} else {
				syslog(LOG_ERR, "select errno %d, %s\n", 
				       errno, strerror(errno));
			}
		} else if (nfd > 0) {
			qp = recv_packet(qinfd);
			if (qp == NULL) {
				if (debug) syslog(LOG_DEBUG, "That's enough\n");
				break;
			}
			if (debug) dump_qp(qp);
			switch (qp->qp_type) {
			case QP_SETUSER:
				if (setup_user(qp->qp_setuser.qp_user,
						qp->qp_setuser.qp_group) == 0) {
					rp = make_qp_returncode(-1, 0, "User not found");
				} else {
					rp = make_qp_ack(QP_SETUSER, 1);
				}
				send_packet(qoutfd, rp);
				qpfree(rp);
				break;
			case QP_RUNCMD:
				child_pid = run_cmd(qp->qp_runcmd.qp_cmdline);
				qpfree(qp);
				handle_qarsh(); /* continue processing qarsh separately */
				return;
				break;
			case QP_RECVFILE: /* Setup file descriptors to handle incoming data */
				rp = prepare_recvfile(qp);
				send_packet(qoutfd, rp);
				qpfree(rp);
				break;
			case QP_DATA: /* qacp sent some data */
				rp = receive_data(qp);
				if (rp) {
					send_packet(qoutfd, rp);
					qpfree(rp);
				}
				break;
			case QP_SENDFILE:
				rp = prepare_sendfile(qp);
				if (rp) {
					send_packet(qoutfd, rp);
					qpfree(rp);
				}
				break;
			case QP_DALLOW: /* qacp is ready to receive some data */
				if (sendfd != -1) {
					nbytes = read(sendfd, buf, bufsize < qp->qp_dallow.qp_count ? bufsize : qp->qp_dallow.qp_count);
					if (nbytes == -1) {
						qp = make_qp_returncode(-1, errno, strerror(errno));
					} else {
						rp = make_qp_data(0, nbytes, buf);
						if (nbytes == 0) {
							close(sendfd);
							sendfd = -1;
						}
					}
					send_packet(qoutfd, rp);
					qpfree(rp);
				}
				break;
			case QP_RETURNCODE: /* qacp either completed or hit an error */
				break;
			case QP_RSTAT:
				if (debug) syslog(LOG_DEBUG, "Got a QP_RSTAT with path = %s\n",
				       qp->qp_rstat.qp_path);
				rp = rstat(qp->qp_rstat.qp_path);
				send_packet(qoutfd, rp);
				qpfree(rp);
				break;
			default:
				syslog(LOG_WARNING, 
					"Packet type %s unimplemented",
					qp_packet_type(qp->qp_type));
			}
			qpfree(qp);
		} else {
			if (debug) syslog(LOG_DEBUG, "Nothing to do in handle_packets\n");
		}
		
	}
}

int
main(int argc, char *argv[])
{
	int ch;
	socklen_t peerlen;
	struct sockaddr_storage peername;
	char peer_hoststr[NI_MAXHOST];
	char peer_portstr[NI_MAXSERV];

	openlog("qarshd", LOG_PID, LOG_DAEMON);

	while ((ch = getopt(argc, argv, "pd")) != -1) {
		switch (ch) {
			case 'd':
				debug = 1;
				break;
			case 'p':
				dopause = 1;
				break;
			case '?':
			default:
				printf("unknown option '%c'\n", optopt);
				exit(1);
		}
	}

	/* daemon initialization */
	peerlen = sizeof peername;
	getpeername(0, (struct sockaddr *)&peername, &peerlen);
	getnameinfo((struct sockaddr *)&peername, peerlen, 
	            peer_hoststr, NI_MAXHOST, peer_portstr, 
	            NI_MAXSERV, NI_NUMERICHOST);
	syslog(LOG_INFO, "Talking to peer %s:%s (%s)", peer_hoststr, peer_portstr,
		peername.ss_family == AF_INET ? "IPv4" : "IPv6");

	/* Start reading packets from stdin */
	handle_packets();

	return 0;
}