/*
 * QA Remote Copy - client side
 *
 * Copy a file to, or from, server.
 * 
 * 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 <assert.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 <sys/sendfile.h>
#include <libgen.h>
#include <stdarg.h>


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

#define QARSHD_CONTROL_PORT 5016

/* Globals */
int qacp_fd = -1;		/* The control connection to qacpd */
unsigned short qarsh_ss_family;  /* AF_INET/AF_INET6, set on connect */
short quiet = 0;

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

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


void
usage()
{
/* 	printf("usage: qacp: [-r] [[user@]host1:]file1 [...] " */
	       /* "[[user@]host2:]file2\n"); */

	/* For now, only simple cmdlines */
	fprintf(stderr, "usage:\n\t%s\n\t%s\n",
	        "qacp <localpath> [<loacalpath> ...] [user@]host:<rmtpath>",
	        "qacp [user@]host:<rmtpath> [[user@]host:<rmtpath> ...] <localpath> ");
}


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

	qp = make_qp_setuser(user, group);
	send_packet(qacp_fd, qp);
	qpfree(qp);
	qp = recv_packet(qacp_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(qacp_fd);
		qpfree(qp);
		exit(125);
	}

	qpfree(qp);
}


struct qp_rstat_pkt *
qacp_rstat(const char *rmtpath, int *rstaterrno)
{
	struct qa_packet *qp;
	struct qp_rstat_pkt *rstatp;

	qp = make_qp_rstat(rmtpath, NULL);
	send_packet(qacp_fd, qp);
	qpfree(qp);

	qp = recv_packet(qacp_fd);
	if (qp) {
		if (qp->qp_type == QP_RSTAT) {
			rstatp = malloc(sizeof *rstatp);
			rstatp->qp_path = strdup(qp->qp_rstat.qp_path);
			rstatp->qp_st_mode = qp->qp_rstat.qp_st_mode;
			rstatp->qp_st_uid = qp->qp_rstat.qp_st_uid;
			rstatp->qp_st_gid = qp->qp_rstat.qp_st_gid;
			rstatp->qp_st_size = qp->qp_rstat.qp_st_size;
			*rstaterrno = 0;
		} else if (qp->qp_type == QP_RETURNCODE) {
			rstatp = NULL;
			*rstaterrno = qp->qp_returncode.qp_errno;
		}

		qpfree(qp);
	} 

	return rstatp;
}

void
free_rstat(struct qp_rstat_pkt *p)
{
	free(p->qp_path);
	free(p);
}

void
qacp_sendonefile(const char *host, const char *srcfile, const char *destfile)
{
	struct qa_packet *qp;
	int fd;
	ssize_t nbytes;
	struct stat sb;
	const int bufsize = QARSH_MAX_PACKET_SIZE/2;
	char buf[bufsize];
	fd_set rfds, wfds;
	int nfd;
	struct timeval timeout;
	int dallow;
	int complete;


	if ((fd = open(srcfile, O_RDONLY)) <0) {
		fprintf(stderr, "Could not open %s: %s\n", srcfile, strerror(errno));
		close(qacp_fd);
		exit(errno);
	}

	if (fstat(fd, &sb) < 0) {
		fprintf(stderr, "Could not stat %s: %s\n", srcfile, strerror(errno));
		close(qacp_fd);
		exit(errno);
	}

	if (S_ISDIR(sb.st_mode)) {
		fprintf(stderr, "Skipping directory %s, recursive not supported, use rsync.\n", srcfile);
		return;
	}

	/* Packet types are qarshd-centric, so if we want to send a file to the
	 * host running qarshd we have to tell qarshd to recv a file. */

	qp = make_qp_recvfile(destfile, sb.st_size, sb.st_mode);
	send_packet(qacp_fd, qp);
	qpfree(qp);

	dallow = 0;
	complete = 0;
	for (;;) {
		FD_ZERO(&rfds);
		FD_SET(qacp_fd, &rfds);
		FD_ZERO(&wfds);
		if (!complete && dallow > 0) FD_SET(qacp_fd, &wfds);
		timeout.tv_sec = 10;
		timeout.tv_usec = 0;

		nfd = select(qacp_fd+1, &rfds, &wfds, NULL, &timeout);
		if (nfd < 0) {
			if (errno == EINTR) {
				continue;
			} else {
				fprintf(stderr, "select errno %d, %s\n", errno, strerror(errno));
			}
		} else if (nfd > 0) {
			if (nfd && FD_ISSET(qacp_fd, &rfds)) {
				qp = recv_packet(qacp_fd);
				if (qp == NULL) {
					fprintf(stderr, "No packet\n");
					break;
				}
				if (qp->qp_type == QP_DALLOW) {
					if (qp->qp_dallow.qp_count > 0) {
						dallow += qp->qp_dallow.qp_count;
					} else {
						dallow = 0;
					}
				} else if (qp->qp_type == QP_RETURNCODE) {
					if (qp->qp_returncode.qp_rc == 0) { /* success */
						break;
					} else if (qp->qp_returncode.qp_rc == -1) {
						fprintf(stderr, "Transfer failed: %s\n", 
								qp->qp_returncode.qp_strerror);
						exit(125);
					}
				} else {
					fprintf(stderr, "Received unexpected packet\n");
					dump_qp(qp);
					break;
				}
				qpfree(qp);
			}
			if (nfd && FD_ISSET(qacp_fd, &wfds)) {
				/* Create a data packet to send */
				nbytes = read(fd, buf, dallow);
				if (nbytes < 0) {
					fprintf(stderr, "read() error: %s\n",
							strerror(errno));
					qp = make_qp_returncode(-1, errno, strerror(errno));
				} else {
					qp = make_qp_data(0, nbytes, buf);
					if (nbytes == 0) { /* Transfer is complete */
						complete = 1;
					} else {
						dallow -= nbytes;
					}
				}
				send_packet(qacp_fd, qp);
				qpfree(qp);
			}
		}
	}
	close(fd);

	if (!quiet) {
		printf("%-30.30s  -> %s:%s\n", srcfile, host, destfile);
	}

	return;
}


void
qacp_recvonefile(const char *host, const char *srcfile, const char *destfile)
{
	struct qa_packet *qp;
	struct qa_packet *rp;
	int outfd;
	ssize_t nwrote;
	struct qp_rstat_pkt *rstatp;
	int rstaterrno;
	fd_set rfds;
	int nfd;
	struct timeval timeout;

	rstatp = qacp_rstat(srcfile, &rstaterrno);
	if (!rstatp) {
		fprintf(stderr, "%s: %s\n", srcfile, strerror(rstaterrno));
		return;
	}

	if ((outfd = open(destfile, O_TRUNC|O_CREAT|O_WRONLY)) < 0) {
		fprintf(stderr, "Could not open %s: %s\n", destfile, strerror(errno));
		close(qacp_fd);
		exit(errno);
	}

	fchmod(outfd, rstatp->qp_st_mode);

	/* Recall that the packet types are qarshd-centric, so if we want
	 * to recv a file from the host running qarshd we have to tell 
	 * qarshd to send a file. */

	qp = make_qp_sendfile(srcfile);
	send_packet(qacp_fd, qp);
	qpfree(qp);

	qp = make_qp_data_allow(0, QARSH_MAX_PACKET_SIZE/2);
	send_packet(qacp_fd, qp);
	qpfree(qp);

	for (;;) {
		FD_ZERO(&rfds);
		FD_SET(qacp_fd, &rfds);

		timeout.tv_sec = 10;
		timeout.tv_usec = 0;

		nfd = select(qacp_fd+1, &rfds, NULL, NULL, &timeout);
		if (nfd < 0) {
			if (errno == EINTR) {
				continue;
			} else {
				fprintf(stderr, "select errno %d, %s\n", errno, strerror(errno));
			}
		} else if (nfd > 0) {
			if (nfd && FD_ISSET(qacp_fd, &rfds)) {
				qp = recv_packet(qacp_fd);
				if (qp == NULL) {
					fprintf(stderr, "No packet\n");
					break;
				}
				if (qp->qp_type == QP_DATA) {
					if (qp->qp_data.qp_count == 0) {
						break;
					}
					if ((nwrote = write(outfd, qp->qp_data.qp_blob, qp->qp_data.qp_count)) == -1) {
						fprintf(stderr, "write() error: %s\n", strerror(errno));
						rp = make_qp_returncode(-1, errno, strerror(errno));
						send_packet(qacp_fd, rp);
						qpfree(rp);
						break;
					} else if (nwrote == qp->qp_data.qp_count) {
						rp = make_qp_data_allow(0, nwrote);
						send_packet(qacp_fd, rp);
						qpfree(rp);
					} else {
						fprintf(stderr, "short write\n");
						break;
					}
				} else if (qp->qp_type == QP_RETURNCODE) {
					if (qp->qp_returncode.qp_rc == 0) { /* success */
						break;
					} else if (qp->qp_returncode.qp_rc == -1) {
						fprintf(stderr, "Transfer failed: %s\n", 
								qp->qp_returncode.qp_strerror);
						exit(125);
					}
				} else {
					fprintf(stderr, "Received unexpected packet\n");
					dump_qp(qp);
					break;
				}
				qpfree(qp);
			}
		}
	}
	close(outfd);

	if (!quiet) {
		printf("%s:%s  -> %30.30s\n", host, srcfile, destfile);
	}
	
	return;
}


int
recvfiles(char **argv, int argc, int fileidx, short recursive)
{
	char *cp;
	char *rhost = NULL;
	char *ruser = NULL;
	char *rgrp = NULL;
	struct passwd *pw;
	int file;
	char *rmtpath;
	char *destpath;
	int destisdir;
	char *tmpstr;
	char *rbnp;  /* remote path basename */
	struct qp_rstat_pkt *rstatp;
	struct stat sb;
	int rstaterrno;

	if (strchr(argv[argc-1], ':') != NULL) {
		fprintf(stderr, "%s is not a valid local file string\n",
		        argv[argc-1]);
		return -1; 
	}

	destpath = strdup(argv[argc-1]);

	/* Figure out if our destpath is just a dir. 
	 * If so we need to set the flag to tack on the remote
	 * path basename onto our destpath
	 */
	
	if (stat(destpath, &sb) < 0) {
		/* stat couldn't file the file, check if the dirname exists, 
		 * if so we need to reset our destpath to the local dirname
		 * if that directory exists. 
		 */
		tmpstr = strdup(destpath);
		cp = strdup(dirname(tmpstr));
		free(tmpstr);
		if (stat(cp, &sb) < 0) {
			fprintf(stderr, "%s: %s\n", cp, strerror(errno));
			free(cp);
			goto error_free_destpath;
		}

		if (S_ISDIR(sb.st_mode)) {
			/* No, we don't want to set destisdir, as our dest is really
			 * a full path, we striped off the filename above to check
			 * the rest of the path exists. 
			 */
			destisdir = 0;
		} else {
			fprintf(stderr, "%s not a directory\n", cp);
			free(cp);
			goto error_free_destpath;
		}
		free(cp);
	} else {
		if (S_ISDIR(sb.st_mode)) {
			destisdir = 1;
		} else if (S_ISREG(sb.st_mode)) {
			destisdir = 0;
		} else {
			fprintf(stderr, "%s not a file or directory\n", destpath);
			goto error_free_destpath;
		}
	}

	if (!destisdir && argc - fileidx > 2) {
		fprintf(stderr, "Will not copy multiple remote files to the same local file\n");
		goto error_free_destpath;
	}
	

	for (file = fileidx; file < (argc - 1); file++) {

	/* For each file in remote file list (fileidx to argc-2) */
	/* Rstat the file, if it is a dir, error */
	/* Otherwise copy it to local location */

		if ((rmtpath = strchr(argv[file], ':')) == NULL) {
			fprintf(stderr, "%s is not a valid remote host:file string\n",
		        	argv[file]);
			goto error_free_destpath;
		}

		/* Split the remote path from the rest */
		*(rmtpath++) = '\0';
		rhost = argv[file];

		/* Grab any user/group info */
		if ((cp = strchr(argv[file], '@'))) {
			ruser = argv[file];
			rhost = cp+1;
			*cp = '\0';
		}

		if (ruser && (cp = strchr(ruser, '.'))) {
			rgrp = cp+1;
			*cp = '\0';
		}


		if (ruser == NULL) {
			if (!(pw = getpwuid(getuid()))) {
				fprintf(stderr, "qacp: unknown user id.\n");
				goto error_free_destpath;
			}
			ruser = pw->pw_name;
		}

		qacp_fd = connect_to_host(rhost, QARSHD_CONTROL_PORT, &qarsh_ss_family);
		if (qacp_fd == -1) {
			if (errno == 0) {
				fprintf(stderr, "Could not connect to %s:%d, %d: %s\n",
					rhost, QARSHD_CONTROL_PORT, h_errno, hstrerror(h_errno));
			} else {
				fprintf(stderr, "Could not connect to %s:%d, %d: %s\n",
					rhost, QARSHD_CONTROL_PORT, errno, strerror(errno));
			}
			goto error_free_destpath;
		}

		set_remote_user(ruser, rgrp);

		rstatp = qacp_rstat(rmtpath, &rstaterrno);
		if (rstatp) {
			if (S_ISREG(rstatp->qp_st_mode)) {
				if (destisdir) {
					rbnp = basename(rmtpath);

					cp = malloc(strlen(destpath) + strlen(rbnp) + 2);
					strcpy(cp, destpath);
					strcat(cp, "/");
					strcat(cp, rbnp);

					qacp_recvonefile(rhost, rmtpath, cp);
					free(cp);
				} else {
					qacp_recvonefile(rhost, rmtpath, destpath);
				}
			}
			if (S_ISDIR(rstatp->qp_st_mode)) {
				fprintf(stderr, "%s: Not a regular file\n", argv[file]);
				free(destpath);
				free(rmtpath);
				free_rstat(rstatp);
				return -1;
			}
			free_rstat(rstatp);
		} else {
			fprintf(stderr, "%s: %s\n", argv[file], strerror(rstaterrno));
			goto error_free_destpath;
		}

		close(qacp_fd);
	} /* for file */
	free(destpath);

	return 0;

error_free_destpath:
	free(destpath);
	return -1;
}


int
sendfiles(char **argv, int argc, int fileidx, short recursive)
{
	char *cp;
	char *rhost = NULL;
	char *ruser = NULL;
	char *rgrp = NULL;
	struct passwd *pw;
	int file;
	char *rmtpath;
	char *destpath;
	char *tmpstr;
	char *lbnp;  /* local file basename */
	char *rdnp;  /* remote path dirname */
	struct qp_rstat_pkt *rstatp;
	int rstaterrno;


	if ((rmtpath = strchr(argv[argc-1], ':')) == NULL) {
		fprintf(stderr, "%s is not a valid remote host:file string\n",
		        argv[argc-1]);
		return -1; 
	}

	*(rmtpath++) = '\0';
	rhost = argv[argc-1];

	/* Grab any user/group info */
	if ((cp = strchr(rhost, '@'))) {
		ruser = argv[argc-1];
		rhost = cp+1;
		*cp = '\0';
	}

	if (ruser && (cp = strchr(ruser, '.'))) {
		rgrp = cp+1;
		*cp = '\0';
	}

	if (ruser == NULL) {
		if (!(pw = getpwuid(getuid()))) {
			fprintf(stderr, "qacp: unknown user id.\n");
			return -1;
		}
		ruser = pw->pw_name;
	}

	qacp_fd = connect_to_host(rhost, QARSHD_CONTROL_PORT, &qarsh_ss_family);
	if (qacp_fd == -1) {
		if (errno == 0) {
			fprintf(stderr, "Could not connect to %s:%d, %d: %s\n",
				rhost, QARSHD_CONTROL_PORT, h_errno, hstrerror(h_errno));
		} else {
			fprintf(stderr, "Could not connect to %s:%d, %d: %s\n",
				rhost, QARSHD_CONTROL_PORT, errno, strerror(errno));
		}

		return -1;
	}

	set_remote_user(ruser, rgrp);

	rstatp = qacp_rstat(rmtpath, &rstaterrno);

	for (file = fileidx; file < (argc - 1); file++) {
		lbnp = basename(argv[file]);

		if (rstatp) {
			/* 
			 * If rmtpath is a dir, then we tack on the 
			 * local files basename.
		 	 */
			if (S_ISDIR(rstatp->qp_st_mode)) {
				destpath = malloc(strlen(rmtpath) + strlen(lbnp) + 2);
				strcpy(destpath, rmtpath);
				strcat(destpath, "/");
				strcat(destpath, lbnp);
			}

			/* It rmtpath is a file, then we leave it as is */
			if (S_ISREG(rstatp->qp_st_mode)) {
				destpath = strdup(rmtpath);
			}
		} else {
			struct qp_rstat_pkt *tmprstatp;

			/* rmtpath does not exist, check if the dirname does */
			tmpstr = strdup(rmtpath);
			rdnp = strdup(dirname(tmpstr));
			free(tmpstr);

			tmprstatp = qacp_rstat(rdnp, &rstaterrno);
			free(rdnp);

			if (tmprstatp) {
				/* Ok, the dir exists, i.e. rmtpath is /tmp/foo
			 	* and the first rstat was an error.  Now we know
			 	* that /tmp exists, we leave destpath == rmtpath.
				*/
				destpath = strdup(rmtpath);
				free_rstat(tmprstatp);
			} else {
				fprintf(stderr, "%s:%s - %s\n", rhost, rmtpath, 
				        strerror(rstaterrno));
				return -1;
			}
		}

		qacp_sendonefile(rhost, argv[file], destpath);

		free(destpath);
	}

	if (rstatp) {
		free_rstat(rstatp);
	}

	close(qacp_fd);
	return 0;
}


int 
main(int argc, char *argv[])
{
	int c;
	short recursive=0;
	int eflag;

	while ((c = getopt(argc, argv, "+rq")) != -1) {
		switch (c) {
			case 'r':
				fprintf(stderr, "-r not supported\n");
				exit(2);
				recursive = 1;
				break;
			case 'q':
				quiet = 1;
				break;
			case '?':
			default:
				printf("Unknown option %c\n", (char)optopt);
				usage();
				exit(1);
		} 
	}

	if (argv[optind] == NULL || (argc - optind < 2)) {
		usage();
		exit(2);
	}

	eflag = 0;
	/* If last arg has a ':' then we are sending files */
	if (strchr(argv[argc-1], ':') != NULL) {
		if (sendfiles(argv, argc, optind, recursive) < 0) {
			eflag++;
		}
	} else {
		if (recvfiles(argv, argc, optind, recursive) < 0) {
			eflag++;
		}
	}

	if (eflag)  {
		exit(1);
	}

	exit(0);
}