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


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

/*
 * QA Remote Shell Daemon
 */


/* Globals */
struct sockaddr_in peername;
int child_exitted = 0;

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

void
sig_handler(int sig)
{
	if (sig == SIGCHLD) child_exitted++;
}

pid_t 
run_cmd(const char *cmd, int p_in, int p_out, int p_err)
{
	pid_t pid;
	int new_in, new_out, new_err;

	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 qarsh */
		new_in = connect_to_peer(&peername, p_in);
		if (new_in == -1) syslog(LOG_WARNING, "connect to new_in failed");
		dup2(new_in, fileno(stdin));
		new_out = connect_to_peer(&peername, p_out);
		if (new_out == -1) syslog(LOG_WARNING, "connect to new_out failed");
		dup2(new_out, fileno(stdout));
		new_err = connect_to_peer(&peername, p_err);
		if (new_err == -1) syslog(LOG_WARNING, "connect to new_err failed");
		dup2(new_err, fileno(stderr));

		execlp("sh", "sh", "-c", cmd, NULL);
		printf("exec of %s failed: %d, %s\n", cmd, errno, strerror(errno));
		exit(127);
	} 
	return pid;
}

ssize_t 
recvfile(const char *path, int if_port, size_t count, mode_t mode)
{
	int sd;
	int ofd;
	char buf[BUFSIZ];
	ssize_t nread;
	ssize_t nwrote;
	ssize_t nleft;

	/* Read count bytes from ifd (sd after we connect), 
	 * write into file @ path 
	 */

	sd = connect_to_peer(&peername, if_port);
	if (sd == -1) {
		syslog(LOG_WARNING, "connect to if_port failed\n");
		return -1;
	}

	if ((ofd = open(path, O_TRUNC|O_CREAT|O_WRONLY, mode)) < 0) {
		syslog(LOG_WARNING, "Could not open %s to recv file: %s\n",
		       path, strerror(errno));
		return -1;
	}
	
	fchmod(ofd, mode);

	nleft = count;
	while (nleft > 0) {
		nread = read(sd, buf, BUFSIZ);
		if (nread < 0) {
			return nread;
		} else if (nread == 0) {  /* EOF */
			break;
		}

		nwrote = write(ofd, buf, nread);
		nleft -= nread;
	}

	if (nleft != 0) {
		unlink(path);
		syslog(LOG_WARNING, "Short file transfer in recvfile(), "
		       "%zd bytes lost, wanted %zd\n", nleft, count);
	}

	close(sd);
	close(ofd);

	return count - nleft;
}

ssize_t
pushfile(const char *path, int of_port)
{
	int outsd;
	int infd;
	off_t offset = 0;
	ssize_t nbytes;
	struct stat sb;

	outsd = connect_to_peer(&peername, of_port);
	if (outsd == -1) {
		syslog(LOG_WARNING, "Connect to of_port (%d) failed: %s\n",
		       of_port, strerror(errno));
		return -1;
	}

	infd = open(path, O_RDONLY);
	if (infd == -1) {
		syslog(LOG_WARNING, "Could not open %s: %s\n",
		       path, strerror(errno));
		close(outsd);
		return -1;
	}

	if (fstat(infd, &sb) < 0) {
		syslog(LOG_WARNING, "Could not stat %s: %s\n",
		       path, strerror(errno));
		close(infd);
		close(outsd);
		return -1;
	}

	nbytes = sendfile(outsd, infd, &offset, sb.st_size);

	close(infd);
	close(outsd);

	return nbytes;
}

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

void
handle_packets(int infd)
{
	fd_set rfds;
	int nfd;
	struct timespec timeout;
	struct qa_packet *qp = NULL, *rp = NULL;
	sigset_t sigmask, orig_sigmask;
	struct sigaction sa;

	pid_t child_pid = 0;
	int child_status;
	ssize_t nbytes;	

	sigemptyset(&sigmask);
	sigaddset(&sigmask, SIGCHLD);
	sigprocmask(SIG_BLOCK, &sigmask, &orig_sigmask);
	sa.sa_handler = sig_handler;
	sa.sa_mask = sigmask;
	sa.sa_flags = SA_RESTART;
	sigaction(SIGCHLD, &sa, NULL);

	for (;;) {
		FD_SET(infd, &rfds);
		timeout.tv_sec = 3;
		timeout.tv_nsec = 0;

		nfd = pselect(infd+1, &rfds, NULL, NULL, &timeout, &orig_sigmask);

		if (child_exitted) {
			waitpid(child_pid, &child_status, 0);
			child_exitted--;
			rp = make_qp_cmdexit(child_pid, child_status);
			send_packet(fileno(stdout), rp);
			qpfree(rp);
		}

		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(infd);
			if (qp == NULL) {
				syslog(LOG_INFO, "That's enough\n");
				break;
			}
			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_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(fileno(stdout), rp);
				qpfree(rp);
				break;
			case QP_RUNCMD:
				child_pid = run_cmd(qp->qp_runcmd.qp_cmdline,
						qp->qp_runcmd.qp_stdin_port,
						qp->qp_runcmd.qp_stdout_port,
						qp->qp_runcmd.qp_stderr_port);
				break;
			case QP_RECVFILE: 
				syslog(LOG_INFO, "Got a QP_RECVFILE with path = %s, "
				       "ifd = %d, count = %zd, mode = %o\n",
				        qp->qp_recvfile.qp_path,
				        qp->qp_recvfile.qp_if_port,
				        qp->qp_recvfile.qp_count,
				        qp->qp_recvfile.qp_mode);
				nbytes = recvfile(qp->qp_recvfile.qp_path, 
				                  qp->qp_recvfile.qp_if_port, 
				                  qp->qp_recvfile.qp_count, 
 				                  qp->qp_recvfile.qp_mode);
				if (nbytes < 0) {
					rp = make_qp_returncode(-1, errno, strerror(errno));
				} else if (nbytes < qp->qp_recvfile.qp_count) {
					char tmpstr[512];
					sprintf(tmpstr, "Excpected %zd, wrote %zd\n", 
					        qp->qp_recvfile.qp_count, nbytes);
					rp = make_qp_returncode(-1, 0, tmpstr);
				} else {
					rp = make_qp_returncode(0, 0, "Transfer Complete");
				}
				send_packet(fileno(stdout), rp);
				qpfree(rp);
				break;
			case QP_SENDFILE:
				syslog(LOG_INFO, "Got a QP_SENDFILE with path = %s, "
				       "ofd = %d\n",
				        qp->qp_sendfile.qp_path,
				        qp->qp_sendfile.qp_of_port);
				nbytes = pushfile(qp->qp_sendfile.qp_path, 
				                  qp->qp_sendfile.qp_of_port);
				if (nbytes < 0) {
					rp = make_qp_returncode(-1, errno, strerror(errno));
				} else {
					rp = make_qp_returncode(0, 0, "Transfer Complete");
				}
				send_packet(fileno(stdout), rp);
				qpfree(rp);
				break;
			case QP_RSTAT:
				syslog(LOG_INFO, "Got a QP_RSTAT with path = %s\n",
				       qp->qp_rstat.qp_path);
				rp = rstat(qp->qp_rstat.qp_path);
				send_packet(fileno(stdout), rp);
				qpfree(rp);
				break;
			default:
				syslog(LOG_WARNING, 
					"Packet type %s unimplemented",
					qp_packet_type(qp->qp_type));
			}
		} else {
			syslog(LOG_DEBUG, "Nothing to do\n");
		}
		
	}
}

int
main(int argc, char *argv[])
{
	int ch;
	socklen_t peernamelen;

	openlog("qarshd", LOG_PID, LOG_DAEMON);

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

	/* daemon initialization */
	peernamelen = sizeof peername;
	getpeername(0, (struct sockaddr *)&peername, &peernamelen);
	syslog(LOG_INFO, "Talking to peer %s:%d", 
		inet_ntoa(peername.sin_addr), ntohs(peername.sin_port));
	/* Start reading packets from stdin */
	handle_packets(0);

	return 0;
}