/* Simple file transfer protocol over a tcp socket */

#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <limits.h>
#include <fcntl.h>
#include <sys/signal.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>

#include "sftp.h"

#define check_connected() \
	if (connected == 0) {\
		printf("Not connected\n");\
		return -1;\
	}

int do_help (char **args, int nargs);
int do_command(char **args, int nargs);
int do_delete(char **args, int nargs);
int do_open (char **args, int nargs);
int do_close (char **args, int nargs);
int do_get (char **args, int nargs);
int do_put (char **args, int nargs);
int do_cd (char **args, int nargs);
int do_pwd (char **args, int nargs);
int do_lcd (char **args, int nargs);
int do_quit (char **args, int nargs);

map action[] = {
	{"?",		do_help},
	{"help",	do_help},
	{"open",	do_open},
	{"close",	do_close},
	{"get",		do_get},
	{"put",		do_put},
	{"cd",		do_cd},
	{"pwd",		do_pwd},
	{"ls",		do_command},
	{"rm",		do_command},
	{"delete",	do_delete},
	{"lcd",		do_lcd},
	{"exec",	do_command},
	{"quit",	do_quit},
	{"exit",	do_quit},
	{NULL,		NULL},
};

int connected = 0;
pid_t child;
int sock;
char *shell;

void
usage(char *prog) {
	printf ("Usage: %s [hostname]\n", prog);
	exit(-1);
}

int
do_help (char **args, int nargs) {
	printf ("Supported commands are:\n");
	printf ("? help open close get put cd pwd ls delete rm lcd exec quit exit\n");
	return 0;
}

int
do_command(char **args, int nargs) {
	char buf[BUFSIZE];
	int i, len = 0;
	message m;

	check_connected();
	for (i=0; i<nargs; i++) {
		if (i==0 && strcmp(args[i], "exec") == 0)
			continue;
		strcpy(buf+len, args[i]);
		len += strlen(args[i]);
		buf[len++] = 0;
	}
	send_message(sock, _message(EXEC, buf, len));
	while (1) {
		recv_message(sock, &m);
		switch (m.command) {
			case ERROR:
				puts((char *)m.data);
				free(m.data);
				return -1;
			case SUCCESS:
				return 0;
			case STREAM:
				break;
			case DATA:
				fwrite(m.data, 1, m.len, stdout);
				free(m.data);
				break;
			case ENDDATA:
				return 0;
		}
	}
}

int
do_delete(char **args, int nargs) {
	args[0] = "rm";
	return do_command(args, nargs);
}

int
do_open (char **args, int nargs) {
	int s[2];

	if (nargs < 2) {
		printf ("Usage: open hostname\n");
		return -1;
	}
	if (socketpair(PF_UNIX, SOCK_STREAM, 0, s) < 0) {
		perror("socketpair");
		return -1;
	}
	connected = 1;
	if ((child = fork())) {
		int n, has_pass = 0;
		char buffer[BUFSIZE];

		close(s[1]);
		sock = s[0];
		n = read(sock, buffer, BUFSIZE);
		if (n <= 0)
			return -1;
		if (strcmp(buffer, "sftp")) {
			fputs (buffer, stdout);
			has_pass = 1;
		}
		if (has_pass) {
			n = read(sock, buffer, BUFSIZE);
			printf("\n");
		}
		if (strcmp(&(buffer[n-4]), "sftp"))
			return 0;
	}
	else {
		close(s[0]);
		dup2(s[1], 0);
		dup2(s[1], 1);
/*		dup2(s[1], 2);*/
		execlp(shell, shell, args[1], "sftpserv", NULL);
	}
	return 0;
}

int
do_close (char **args, int nargs) {
	check_connected();
	if (connected)
		send_message(sock, _message(CLOSE, NULL, 0));
	connected = 0;
	return 0;
}

int
do_get (char **args, int nargs) {
	message m;
	int i;

	check_connected();
	for (i = 1; i < nargs; i++) {
		int gotten = 0;
		send_message(sock,
			     _message(REQUEST, args[i], strlen(args[i])+1));
		while (1) {
			recv_message(sock, &m);
			if (m.command == NOFILEMATCH) {
				if (gotten == 0)
					puts((char*)m.data);
				break;
			}
			else if (m.command != SENDFILE)
				continue;
			if (m.data) {
				struct timeval tv1, tv2, tv;
				int size;

				gettimeofday(&tv1, NULL);
				size = do_recvfile(sock, *(u_int8_t *)m.data);
				gettimeofday(&tv2, NULL);
				if (size < 0) {
					puts("failed to get file");
					return -1;
				}
				tv = timediff(tv2, tv1);
				printf ("%d bytes sent in %ld.%02ld secs, "
					"%8.2f K/s\n",
					size, (long)tv.tv_sec,
					(long)tv.tv_usec/10000,
					size/(1024*(tv.tv_sec + (double)tv.tv_usec/1000000)));
				free(m.data);
				gotten++;
			}
		}
	}
	return 0;
}

int
do_put (char **args, int nargs) {
	int i, n, total = 0;

	check_connected();
	for (i = 1; i < nargs; i++) {
		n = do_sendfile(sock, args[i]);
		if (n < 0) {
			printf("error sending %s\n", args[i]);
			return -1;
		}
		total += n;
	}
	if (total == 0)
		printf("No files specified\n");
	else
		printf("%d file%s sent\n", total, (total == 1) ? "" : "s");
	return 0;
}

int
do_cd (char **args, int nargs) {
	message m;

	check_connected();
	if (nargs > 1)
		send_message(sock, _message(CHDIR, args[1], strlen(args[1])+1));
	else
		send_message(sock, _message(CHDIR, NULL, 0));
	recv_message(sock, &m);
	switch (m.command) {
		case ERROR:
			puts((char*)m.data);
			return -1;
		case SUCCESS:
			return 0;
	}
	return 0;
}

int
do_pwd (char **args, int nargs) {
	message m;

	check_connected();
	send_message(sock, _message(GETDIR, NULL, 0));
	recv_message(sock, &m);
	switch (m.command) {
		case ERROR:
			puts((char*)m.data);
			return -1;
		case TELLDIR:
			puts((char*)m.data);
			return 0;
	}
	return 0;
}

int
do_lcd (char **args, int nargs) {
	int ret;
	if (nargs > 1)
		ret = chdir (args[1]);
	else
		ret = chdir (getenv("HOME"));
	if (ret < 0)
		printf ("chdir failed\n");
	return ret;
	
}

int
do_quit (char **args, int nargs) {
	if (connected)
		do_close(NULL, 0);
	exit(0);
}

void
reap() {
	int status;
	int pid;
	while ((pid = wait3(&status, WNOHANG, NULL)) > 0) {
		if (pid != child)
			continue;
		if (WIFSIGNALED(status))
			printf ("child died with signal %d\n",
					WTERMSIG(status));
		printf ("Not connected\n");
		connected = 0;
	}
}

int
main(int argc, char **argv) {
	char *base, buffer[LINESIZE];
	int ret;
	struct sigaction act;
	sigset_t set;

	base = findbasename(argv[0]);
	if (base[0] == 'r')
		shell = "rsh";
	else
		shell = "ssh";

	switch (argc) {
		case 1: break;
		case 2: do_open(argv, 2); break;
		default: usage(argv[0]);
	}

	sigemptyset(&set);
	act.sa_handler = reap;
	act.sa_mask = set;
	act.sa_flags = 0;
	sigaction (SIGCHLD, &act, NULL);

	while (1) {
		printf ("sftp> ");
		if (fgets (buffer, LINESIZE, stdin) == NULL) {
			do_quit(NULL, 0);
			exit(0);
		}
		ret = do_action(buffer, action);
	}
}
