#include <limits.h>
#include <netdb.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <arpa/inet.h>

#include <netinet/in.h>

#include <sys/socket.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>

#include "ftptossh.h"

struct fd_ll *fds;

/* ftptossh - Ben Winslow <rain@bluecherry.net> */

/* I highly recommend using auth or firewalling the port you listen on.  An
   attacker could do all kinds of not-so-neat things if he/she had access to
   spawn ssh to connect to an arbitary host/user */
/* A final note--the ssh subprocesses can not take any input from the user,
   as the process is daemonized.  Because of this, you must have passwordless
   RSA authentication set up for the hosts you want to FTP to, and their host
   key must be known */

/* This is temporary -- I'm too lazy to write a config file parser right now */
#define AUTH_USER "rain"
#define AUTH_PASS "foo"
#define USE_AUTH
#define PORT 4318
#define USE_OPENSSH

int main(int argc, char *argv[])
{
	fd_set fds_in;
	struct fd_ll *f;
	int highfd;
	
	srand(time(NULL));
	
	if (!make_listener(PORT, handle_incoming))
		return 1;
		
	for (;;) {
		highfd = 0;
		FD_ZERO(&fds_in);
		for (f = fds; f; f = f->next) {
			if (f->fd > highfd)
				highfd = f->fd;
			FD_SET(f->fd, &fds_in);
		}
		if (select(highfd + 1, &fds_in, NULL, NULL, NULL) < 0) {
			perror("select()");
			return 1;
		}
		for (f = fds; f; f = f->next) {
			if (FD_ISSET(f->fd, &fds_in) && f->handle)
				f->handle(f);
		}
		
	}
	return 0;
}

struct fd_ll *find_fd_by_num(int fd)
{
	struct fd_ll *f;
	for (f = fds; f && f->fd != fd; f = f->next);
	return f;
}
struct fd_ll *add_fd(int fd, void *handler)
{
	struct fd_ll *new;

	new = malloc(sizeof(struct fd_ll));
	new->fd = fd;
	new->handle = handler;
	new->data = NULL;
	new->buf = NULL;
	new->next = fds;
	fds = new;

	return new;
}
void rem_fd(struct fd_ll *fd)
{
	struct fd_ll *prev, *f;
	
	if (!fd)
		return;
	close(fd->fd);

	prev = fds;
	while (prev) {
		if (prev->next == fd)
			break;
		prev = prev->next;
	}
	if (prev)
		prev->next = fd->next;
	else
		fds = fd->next;

	for (f = fds; f; f = f->next) {
		if (f->data == fd->data)
			f->data = NULL;
	}
	
	if (fd->data)
		free(fd->data);
	free(fd);
}

struct fd_ll *make_listener(int port, void *handler)
{
	struct sockaddr_in sin;
	int fd;
	
	sin.sin_family = AF_INET;
	sin.sin_port = htons(port);
	sin.sin_addr.s_addr = htons(INADDR_ANY);
	
	if ((fd = socket(AF_INET, SOCK_STREAM, IPPROTO_IP)) < 0) {
		perror("socket()");
		return NULL;
	}
	if (bind(fd, (struct sockaddr *)&sin, sizeof(sin)) < 0) {
		perror("bind()");
		return NULL;
	}
	if (listen(fd, 5) < 0) {
		perror("listen()");
		return NULL;
	}

	return add_fd(fd, handler);
}

void handle_proxy(struct fd_ll *fd)
{
	struct ftpproxy_t *ftpinfo;
	struct fd_ll *proxy = NULL;
	struct hostent *he;
	pid_t pid;
	static int doing_auth = 0;
	static char *auth_user = NULL;
	int x, status, secure = 1;
	char buf[BUF_SIZE];
	
	ftpinfo = fd->data;

	if (readln(fd->fd, buf) <= 0) {
		rem_fd(fd);
		return;
	}
	if (buf[strlen(buf) - 1] == '\n')
		buf[strlen(buf) - 1] = 0;
	if (buf[strlen(buf) - 1] == '\r')
		buf[strlen(buf) - 1] = 0;

	
	if (!strncasecmp(buf, "USER ", 5)) {
		split(NULL, buf, ' ');
#ifdef USE_AUTH
		if (!strchr(buf, '@')) {
			doing_auth = 1;
			auth_user = strdup(buf);
			writeln(fd->fd, "331 Password required for %s\r\n", auth_user);
			return;
		} else if (!ftpinfo->auth_ok) {
			writeln(fd->fd, "530 Please login with USER and PASS.\r\n");
			return;
		}
#endif
		ftpinfo->port = 21;
		ftpinfo->user = malloc(strlen(buf) + 1);
		split(ftpinfo->user, buf, '@');
		ftpinfo->host = malloc(strlen(buf) + 1);
		split(ftpinfo->host, buf, ':');
		ftpinfo->loccontrolport = (int)((65535.0 - 1024.0) * rand() / RAND_MAX) + 1024;
		if (!(he = gethostbyname(ftpinfo->host))) {
			writeln(fd->fd, "421 Unable to resolve hostname\r\n");
			return;
		}
		memcpy(&ftpinfo->host_ip, he->h_addr, he->h_length);
		if ((pid = get_secure_port(ftpinfo, ftpinfo->host, ftpinfo->loccontrolport, ftpinfo->port, DIR_OUTBOUND)) < 0) {
			writeln(fd->fd, "421 Unable to create new process\r\n");
			return;
		}
		waitpid(pid, &status, 0);
		if (WIFEXITED(status)) {
			if (WEXITSTATUS(status) > 0)
				secure = 0;
		}
		if (secure) {
			if (!(proxy = open_socket("localhost", ftpinfo->loccontrolport, handle_remote))) {
				writeln(fd->fd, "421 Unable to connect to remote host\r\n");
				return;
			} else
				writeln(fd->fd, "220-Secure connection established to %s:%d\r\n", ftpinfo->host, ftpinfo->port);
		} else {
			writeln(fd->fd, "421 Secure connection to %s:%d failed.\r\n", ftpinfo->host, ftpinfo->port);
			/* Straight FTP proxy code here */
		}
		writeln(proxy->fd, "USER %s\r\n", ftpinfo->user);
		ftpinfo->remfd = proxy->fd;
		proxy->data = ftpinfo;
	} else if (!strncasecmp(buf, "PORT ", 5)) {
		char *splitbuf, ipbuf[16];
		unsigned char a, b, c, d, portlow, porthigh;
		short port;
		pid_t pid;
		int status;
		
		if (!ftpinfo->auth_ok) {
			writeln(fd->fd, "530 Please login with USER and PASS.\r\n");
			return;
		}
		
		split(NULL, buf, ' ');
		
		splitbuf = malloc(strlen(buf) + 1);
		
		split(splitbuf, buf, ',');
		a = atoi(splitbuf);
		split(splitbuf, buf, ',');
		b = atoi(splitbuf);
		split(splitbuf, buf, ',');
		c = atoi(splitbuf);
		split(splitbuf, buf, ',');
		d = atoi(splitbuf);
		split(splitbuf, buf, ',');
		porthigh = atoi(splitbuf);
		split(splitbuf, buf, ',');
		portlow = atoi(splitbuf);
		port = (porthigh << 8) | portlow;
		
		free(splitbuf);
		
		ftpinfo->remdataport = (int)((65535.0-1024.0)*rand()/RAND_MAX)+1024;
		sprintf(ipbuf, "%d.%d.%d.%d", a, b, c, d);
		if ((pid = get_secure_port(ftpinfo, ipbuf, port, ftpinfo->remdataport, DIR_INBOUND)) < 0) {
			perror("oi");
		}
		waitpid(pid, &status, 0);
		if (WIFEXITED(status)) {
			if (WEXITSTATUS(status) > 0)
				writeln(fd->fd, "500 PORT command failed: unable to establish encrypted connection\n");
			else
				writeln(fd->fd, "200 PORT command successful.\r\n");
		}
		strcpy(ipbuf, inet_ntoa(ftpinfo->host_ip));
		for (x = 0; x < strlen(ipbuf); x++) {
			if (ipbuf[x] == '.')
				ipbuf[x] = ',';
		}
		writeln(ftpinfo->remfd, "PORT %s,%d,%d\r\n", ipbuf, (ftpinfo->remdataport >> 8) & 0xff, ftpinfo->remdataport & 0xff);
#ifdef USE_AUTH
	} else if (doing_auth && !strncasecmp(buf, "PASS ", 5)) {
		split(NULL, buf, ' ');
		if (!auth_user) {
			writeln(fd->fd, "530 Login with USER first.\r\n");
			return;
		}
		if (strcmp(buf, AUTH_PASS) || strcmp(auth_user, AUTH_USER)) {
			writeln(fd->fd, "530 Login incorrect.\r\n");
			return;
		}
		ftpinfo->auth_ok = 1;
		doing_auth = 0;
		if (auth_user)
			free(auth_user);
#endif
	} else {
		if (ftpinfo->remfd != -1 && ftpinfo->auth_ok)
			writeln(ftpinfo->remfd, "%s\r\n", buf);
#ifdef USE_AUTH
		else if (!auth_user)
			writeln(fd->fd, "530 Login with USER first.\r\n");
		else
			writeln(fd->fd, "530 Please login with USER and PASS.\r\n");
#endif
	}
}

void handle_incoming(struct fd_ll *fd)
{
	int newfd, sinlen = sizeof(struct sockaddr_in);
	struct sockaddr_in sin;
	struct fd_ll *new;
	struct ftpproxy_t *ftpinfo;
	
	if ((newfd = accept(fd->fd, (struct sockaddr *)&sin, &sinlen)) < 0)
		return;
	new = add_fd(newfd, handle_proxy);
	ftpinfo = new->data = malloc(sizeof(struct ftpproxy_t));
	ftpinfo->remfd = -1;
	ftpinfo->port = 0;
	ftpinfo->user = ftpinfo->pass = ftpinfo->host = NULL;
	ftpinfo->locfd = new->fd;
#ifdef USE_AUTH
	ftpinfo->auth_ok = 0;
#else
	ftpinfo->auth_ok = 1;
#endif
	ftpinfo->loccontrolport = ftpinfo->remdataport = -1;
	writeln(new->fd, "220 Connected to the FTP to SSH gateway proxy\r\n");
}

int writeln(int fd, const char *format,...)
{
	va_list args;
	char *buf;

	va_start(args, format);
	buf = alloca(BUF_SIZE * 2);

	vsprintf(buf, format, args);

	if (write(fd, buf, strlen(buf)) < 0)
		return -1;
#ifdef DEBUG
	printf("[OUT:%2d] %s", fd, buf);
#endif
	va_end(args);
	return 0;
}

int readln(int fd, char *buf)
{
	char *mybuf, *ptr;
	struct fd_ll *f;
	int result;

	mybuf = alloca(BUF_SIZE);

	if ((result = recv(fd, mybuf, BUF_SIZE - 1, MSG_PEEK)) < 0)
		return -1;
	mybuf[result] = 0;

	if (!result)
		return -2;

	if (!(f = find_fd_by_num(fd)))
		return -1;

	if (!(ptr = strchr(mybuf, '\n'))) {
		result = recv(fd, mybuf, result, 0);
		mybuf[result - 1] = 0;
		if (f->buf) {
			f->buf = realloc(f->buf, strlen(f->buf) + strlen(mybuf) + 1);
			strcpy(f->buf + strlen(f->buf), mybuf);
		} else {
			f->buf = malloc(result);
			strcpy(f->buf, mybuf);
		}
		return -1;
	}
	*ptr = 0;

	if (f->buf) {
		strcpy(buf, f->buf);
		recv(fd, buf + strlen(f->buf), strlen(mybuf) + 1, 0);
		buf[strlen(mybuf) + strlen(f->buf)] = 0;

		free(f->buf);
		f->buf = NULL;
	} else {
		recv(fd, buf, strlen(mybuf) + 1, 0);
		buf[strlen(mybuf)] = 0;
	}

#ifdef DEBUG
	printf("[ IN:%2d] %s\n", f->fd, buf);
#endif
	return strlen(buf);
}

void split(unsigned char *dest, unsigned char *source, int delim)
{
	char *ptr, *buf;

	ptr = strchr(source, delim);

	if (!ptr) {
		if (dest)
			strcpy(dest, source);

		*source = 0;
		return;
	}
	*ptr = 0;

	if (dest)
		strcpy(dest, source);

	buf = strdup(ptr + 1);
	strcpy(source, buf);
	free(buf);
}

struct fd_ll *open_socket(char *host, int port, void *handler)
{
	struct hostent *he;
	struct sockaddr_in sin;
	int fd;
	struct fd_ll *new;
	
	memset(&sin, 0, sizeof(sin));
	
	if (!(he = gethostbyname(host)))
		return NULL;
	sin.sin_family = AF_INET;
	sin.sin_port = htons(port);
	memcpy(&sin.sin_addr, he->h_addr, he->h_length);
	
	if ((fd = socket(AF_INET, SOCK_STREAM, IPPROTO_IP)) < 0)
		return NULL;
	if (connect(fd, (struct sockaddr *)&sin, sizeof(sin)))
		return NULL;
	new = add_fd(fd, handler);
	return new;
}

void handle_remote(struct fd_ll *fd)
{
	char buf[BUF_SIZE];
	struct ftpproxy_t *ftpinfo;
	
	ftpinfo = fd->data;
	
	if (readln(fd->fd, buf) <= 0 || !ftpinfo) {
		rem_fd(fd);
		return;
	}
	if (buf[strlen(buf) - 1] == '\n')
		buf[strlen(buf) - 1] = 0;
	if (buf[strlen(buf) - 1] == '\r')
		buf[strlen(buf) - 1] = 0;

	if (strcasecmp(buf, "200 PORT command successful.") && ftpinfo->locfd != -1)
		writeln(ftpinfo->locfd, "%s\r\n", buf);
}

int get_secure_port(struct ftpproxy_t *ftpinfo, char *addr, int local, int remote, int direction)
{
	int pid;
#ifndef DEBUG
	int x;
#endif
	char *redirarg, *portspec; 

	switch ((pid = fork())) {
		case -1:
			/* error */
			return -1;
		case 0:
			/* child */
			portspec = malloc(strlen(addr) + 6 + 6 + 1);
			if (direction == DIR_INBOUND) {
				redirarg = "-R";
				sprintf(portspec, "%d:%s:%d", remote, addr, local);
			} else if (direction == DIR_OUTBOUND) {
				redirarg = "-L";
				sprintf(portspec, "%d:%s:%d", local, addr, remote);
			} else
				return 1;
#ifdef DEBUG
			printf("Doing the following ssh port redirection: %s %s\n", redirarg, portspec);
#else
			for (x = 0; x < FD_SETSIZE; x++)
				close(x);
#endif
#ifdef USE_OPENSSH
			if (execlp("ssh", "ssh", "-f", "-l", ftpinfo->user, "-o", "Protocol 1", "-o", "PasswordAuthentication no", redirarg, portspec, ftpinfo->host, "sleep 60", NULL) < 0)
#else
#error Regular SSH support is not currently implimented, try OpenSSH instead.
#endif
				return 1;
		default:
			/* parent */
			return pid;
	}
	return -1;
}