#include <sys/types.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include "includes.h"
#include "xmalloc.h"
#include "sftp.h"
#include "sftp-common.h"
#include "sftp-client.h"
#include "log.h"
#include "misc.h"

#define MAX_STR_LEN 256
#define DEFAULT_COPY_BUFLEN	32768	/* Size of buffer for up/download */
#define DEFAULT_NUM_REQUESTS	64	/* # concurrent outstanding requests */

#define SSH_PROGRAM "dbclient"
#define DROPBEAR_PASSWORD_ENV "DROPBEAR_PASSWORD"
#define DEFAULT_DOWNLOAD_PATH "/tmp"
char *__progname;

/* SIGINT received during command processing */
volatile sig_atomic_t interrupted = 0;

/* PID of ssh transport process */
static pid_t sshpid = -1;
/* Suppress diagnositic messages */
int quiet = 0;

typedef enum
{
    SFTP_OP_INVALID,
    SFTP_OP_GET,
    SFTP_OP_PUT
}SFTP_OP;


/* ARGSUSED */
static void
killchild(int signo)
{
	if (sshpid > 1) {
		kill(sshpid, SIGTERM);
		waitpid(sshpid, NULL, 0);
	}

	_exit(1);
}

/* ARGSUSED */
static void
suspchild(int signo)
{
	if (sshpid > 1) {
		kill(sshpid, signo);
		while (waitpid(sshpid, NULL, WUNTRACED) == -1 && errno == EINTR)
			continue;
	}
	kill(getpid(), SIGSTOP);
}


static void
connect_to_server(char *path, char **args, int *in, int *out, char* password)
{
	int c_in, c_out;

#ifdef USE_PIPES
	int pin[2], pout[2];

	if ((pipe(pin) == -1) || (pipe(pout) == -1))
		fatal("pipe: %s", strerror(errno));
	*in = pin[0];
	*out = pout[1];
	c_in = pout[0];
	c_out = pin[1];
#else /* USE_PIPES */
	int inout[2];

	if (socketpair(AF_UNIX, SOCK_STREAM, 0, inout) == -1)
		fatal("socketpair: %s", strerror(errno));
	*in = *out = inout[0];
	c_in = c_out = inout[1];
#endif /* USE_PIPES */

	if ((sshpid = fork()) == -1)
		fatal("fork: %s", strerror(errno));
	else if (sshpid == 0) {
		if ((dup2(c_in, STDIN_FILENO) == -1) ||
		    (dup2(c_out, STDOUT_FILENO) == -1)) {
			fprintf(stderr, "dup2: %s\n", strerror(errno));
			_exit(1);
		}
		close(*in);
		close(*out);
		close(c_in);
		close(c_out);

		/*
		 * The underlying ssh is in the same process group, so we must
		 * ignore SIGINT if we want to gracefully abort commands,
		 * otherwise the signal will make it to the ssh process and
		 * kill it too.  Contrawise, since sftp sends SIGTERMs to the
		 * underlying ssh, it must *not* ignore that signal.
		 */
		signal(SIGINT, SIG_IGN);
		signal(SIGTERM, SIG_DFL);
        setenv(DROPBEAR_PASSWORD_ENV, password, 1);
		execvp(path, args);
		fprintf(stderr, "exec: %s: %s\n", path, strerror(errno));
		_exit(1);
	}

	signal(SIGTERM, killchild);
	signal(SIGINT, killchild);
	signal(SIGHUP, killchild);
	signal(SIGTSTP, suspchild);
	signal(SIGTTIN, suspchild);
	signal(SIGTTOU, suspchild);
	close(c_in);
	close(c_out);
}

static void
usage(char* progname)
{
	fprintf(stderr,
	    "usage: %s -[pg] -e passwd -l localfile -r remotefile username@hostname\n"
	    "-a                   When this option is set, we resume download or upload if possible\n"
	    "-p                   put the local file to server\n"
	    "-g                   get the remote file from server\n"
	    "-y                   Always accept remote host key if unknown\n"
	    "-e passwd            password of user\n"
	    "-l file              local file name\n"
	    "-r file              remote file name\n"
	    "-S program           The ssh program\n"
	    "username@hostname    user name and host name\n",
	    progname);
	exit(1);
}

static inline char* getFileName(char* filename)
{
    int len = 0;
    int cnt = 0;
    int isFind = 0;
    char* retStr = NULL;

    len = strlen(filename);
    if (len ==0 || filename[len-1]=='/')
    {
        return NULL;
    }
    retStr = malloc(len+1);
    if (NULL == retStr)
    {
        return NULL;
    }
    memset(retStr, 0, len+1);
    
    cnt = len;
    while(cnt-- >0)
    {
        if(filename[cnt] == '/')
        {
            isFind = 1;
            break;
        }
    }
    if (isFind)
    {
        strncpy(retStr, &filename[cnt+1], len);
    }
    else
    {
        strncpy(retStr, filename, len);
    }
    return retStr;
}
char *get_progname(char *argv0)
{
	char *p, *q;

	if (argv0 == NULL)
		return ("unknown");	/* XXX */
	p = strrchr(argv0, '/');
	if (p == NULL)
		p = argv0;
	else
		p++;

	if ((q = strdup(p)) == NULL) {
		perror("strdup");
		exit(1);
	}
	return q;
}


int main(int argc, char **argv)
{
    char *sshargs[20] = {NULL};
    int sshargc = 0;
    int in = 0, out = 0, err = 0;
    int rlen = 0;
    int llen = 0;
    int ch = 0;
    char ssh_password[MAX_STR_LEN] = {0};
    char ssh_program[MAX_STR_LEN] = {0};
    char lFile[MAX_STR_LEN] = {0};
    char rFile[MAX_STR_LEN] = {0};
    struct sftp_conn *conn;
    size_t copy_buffer_len = DEFAULT_COPY_BUFLEN;
	size_t num_requests = DEFAULT_NUM_REQUESTS;
    LogLevel ll = SYSLOG_LEVEL_INFO;
    SFTP_OP sftp_op = SFTP_OP_INVALID;
	long long limit_kbps = 0;
    char *remote_path = NULL;
    char *userhost = NULL;
    char *host = NULL;
    char *filename=NULL;
    int isresume = 0;
    int isAccept = 0;
    
    /* Ensure that fds 0, 1 and 2 are open or directed to /dev/null */
	sanitise_stdfd();

    __progname=get_progname(argv[0]);

    strncpy(ssh_program, SSH_PROGRAM, sizeof(ssh_program));
    while ((ch = getopt(argc, argv, "apgye:l:r:S:")) != -1)
    {
        switch (ch)
        {
            case 'p':
            {
                sftp_op = SFTP_OP_PUT;
                break;
            }
            case 'g':
            {
                sftp_op = SFTP_OP_GET;
                break;
            }
            case 'e':
            {
                snprintf(ssh_password, sizeof(ssh_password), "%s", optarg);
                break;
            }
            case 'l':
            {
                snprintf(lFile, sizeof(lFile), "%s", optarg);
                break;
            }
            case 'r':
            {
                snprintf(rFile, sizeof(rFile), "%s", optarg);
                break;
            }
            case 'a':
            {
                isresume=1;
                break;
            }
            case 'y':
            {
                isAccept=1;
                break;
            }
            case 'S':
            {
                snprintf(ssh_program, sizeof(ssh_program), "%s", optarg);
                break;
            }
            default:
            {
                usage(argv[0]);
            }
        }
    }

    if (optind+1 != argc)
        usage(argv[0]);

    rlen = strlen(rFile);
    llen = strlen(lFile);
    if ((strlen(ssh_password) == 0)||(sftp_op==SFTP_OP_INVALID)||
        ((sftp_op==SFTP_OP_PUT)&&((llen==0)||(lFile[llen-1] == '/')))||
        ((sftp_op==SFTP_OP_GET)&&((rlen==0)||(rFile[rlen-1] == '/'))))
    {
        usage(argv[0]);
    }

    userhost = xstrdup(argv[optind]);
    if ((host = strrchr(userhost, '@')) == NULL)
    {
        usage(argv[0]);
    }

    *host++ = '\0';
	if (!userhost[0]) {
		fprintf(stderr, "Missing username\n");
		usage(argv[0]);
	}
    if (colon(host) != NULL) 
    {
        usage(argv[0]);
    }

    host = cleanhostname(host);
    if (!*host) 
    {
        fprintf(stderr, "Missing hostname\n");
        usage(argv[0]);
    }
    
    log_init(argv[0], ll, SYSLOG_FACILITY_LOCAL0, 0);
    
    sshargs[sshargc++] = ssh_program;
    if (isAccept)
    {
        sshargs[sshargc++] = "-y";
    }
    sshargs[sshargc++] = "-l";
    sshargs[sshargc++] = userhost;
    sshargs[sshargc++] = "-s";
    sshargs[sshargc++] = host;
    sshargs[sshargc++] = "sftp";

    connect_to_server(ssh_program, sshargs, &in, &out, ssh_password);
    
    conn = do_init(in, out, copy_buffer_len, num_requests, limit_kbps);
    if (conn == NULL)
        fatal("Couldn't initialise connection to server");

    if(sftp_op==SFTP_OP_PUT)
    {
        if (access(lFile, F_OK) != 0)
        {
            fatal("File \"%s\" not exist.", lFile);
        }

        //we get the file name from local file if the remote file is empty or path
        if (rlen == 0 ||(rFile[rlen-1]=='/'))
        {
            filename = getFileName(lFile);
            if (NULL == filename)
            {
                fatal("Get file name from \"%s\" failed.", lFile);
            }
            if (rlen == 0)
            {
                strncpy(rFile, filename, sizeof(rFile)-1);
            }
            else if (rFile[rlen-1]=='/')
            {
                strncat(rFile, filename, sizeof(rFile)-rlen-1);
            }
            free(filename);
            filename=NULL;
        }
        //we need to get the full path
        if ((rFile[0] != '/'))
        {
            remote_path = do_realpath(conn, ".");
        	if (remote_path == NULL)
        		fatal("Get remote path failed");

            filename = path_append(remote_path, rFile);
            if (filename == NULL)
                fatal("append path [%s] file[%s] failed", remote_path, rFile);

            strncpy(rFile, filename, sizeof(rFile)-1);
            free(filename);
            filename = NULL;
            free(remote_path);
            remote_path=NULL;
        }
        
        if(do_upload(conn, lFile, rFile,
			    1, isresume,1) == -1)
        {
            fatal("Upload fail");
        }
    }
    else if (sftp_op==SFTP_OP_GET)
    {
        if (rFile[0] != '/')
        {
            remote_path = do_realpath(conn, ".");
        	if (remote_path == NULL)
        		fatal("Need cwd");

            filename = path_append(remote_path, rFile);
            if (filename == NULL)
                fatal("append path [%s] file[%s] failed", remote_path, rFile);

            strncpy(rFile, filename, sizeof(rFile)-1);
            free(filename);
            filename=NULL;
            free(remote_path);
            remote_path=NULL;
        }

         //we get the file name from local file if the remote file is empty or path
        if (llen == 0 ||(lFile[llen-1]=='/'))
        {
            filename = getFileName(rFile);
            if (NULL == filename)
            {
                fatal("Get file name from \"%s\" failed.", rFile);
            }
            if (llen == 0)
            {
                strncpy(lFile, filename, sizeof(lFile)-1);
            }
            else if (lFile[llen-1]=='/')
            {
                strncat(lFile, filename, sizeof(lFile)-llen-1);
            }
            free(filename);
            filename = NULL;
        }

        if (lFile[0]!='/')
        {
            filename = path_append(DEFAULT_DOWNLOAD_PATH, lFile);
            if (filename == NULL)
                fatal("append path [%s] file[%s] failed", DEFAULT_DOWNLOAD_PATH, lFile);

            strncpy(lFile, filename, sizeof(lFile)-1);
            free(filename);
            filename=NULL;
        }
        
        if (do_download(conn, rFile, lFile, NULL,
	        1, isresume, 1) == -1)
        {
            fatal("Download fail");
        }
    }

    free(conn);

#if !defined(USE_PIPES)
    shutdown(in, SHUT_RDWR);
    shutdown(out, SHUT_RDWR);
#endif

	close(in);
	close(out);

	while (waitpid(sshpid, NULL, 0) == -1)
		if (errno != EINTR)
			fatal("Couldn't wait for ssh process: %s",
			    strerror(errno));

	exit(err == 0 ? 0 : 1);
    
}