#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <netinet/in.h>
#include <netdb.h>

typedef struct _session
{
	char buffer[8096];
  int fd;
	int read_mark;
	int write_mark;
} session_t;

typedef struct _file
{
  char *name;
  struct _file *next;
} file_t;

file_t *files = NULL;

char current_command[5];
session_t command_session;
session_t data_session;
char *username;
char *password;
char *directory;
char *target_dir;

int first_run = 1;

int connect_to_server(char *hostname, int port);

file_t *file_new(char *name, int length)
{
  file_t *f;

  f = malloc(sizeof(file_t));
  memset(f, 0, sizeof(file_t));
  
  f->name = malloc(length + 1);
  memset(f->name, 0, length + 1);
  memcpy(f->name, name, length);
  return f;
}

void file_push(file_t *f)
{
  if (files == NULL)
    files = f;
  else
    {
      f->next = files;
      files = f;
    }
}

file_t *file_peek()
{
  return files;
}

file_t *file_pop()
{
  file_t *f;
  
  if (files == NULL)
    return NULL;

  f = files;
  files = files->next;
  return f;
}

void file_free(file_t *f)
{
  if (f == NULL)
    return;

  free(f->name);
  free(f);
}

void send_command(int fd, char *command, int length)
{
	send(fd, command, length, 0);
	strncpy(current_command, command, 4);
}

// end of line
int find_eol(char *buffer, int length)
{
	int i;

	for (i = 0; i < length - 1; i++)
		if (buffer[i] == '\r' && buffer[i + 1] == '\n')
			return i + 1;

	//printf("eol source len: %d\n", length);

	return -1;
}

void welcome_callback(int fd, int code, char *resp, int length)
{
	char command[1024];

	printf("[-] Sending username: %s\n", username);
	memset(command, 0, sizeof(command));
	snprintf(command, 1023, "USER %s\r\n", username);
	send_command(fd, command, strlen(command));
}

void password_callback(int fd, int code, char *resp, int length) 
{
	char command[1024];
	printf("[-] Sending password\n");

	memset(command, 0, sizeof(command));
	snprintf(command, 1023, "PASS %s\r\n", password);
	send_command(fd, command, strlen(command));
}

void login_success_callback(int fd, int code, char *resp, int length)
{
	char command[1024];

	printf("[-] Set image mode\n");
	memset(command, 0, sizeof(command));
	snprintf(command, 1023, "TYPE I\r\n");
	send_command(fd, command, strlen(command));
}

void type_callback(int fd, int code, char *resp, int length)
{
	char command[1024];
	
	printf("[-] Changing directory to %s\n", directory);
	memset(command, 0, sizeof(command));
	snprintf(command, 1023, "CWD %s\r\n", directory);
	send_command(fd, command, strlen(command));
}

void cwd_callback(int fd, int code, char *resp, int length)
{
	char command[1024];

	memset(command, 0, sizeof(command));
	snprintf(command, 1023, "PASV\r\n");
	send_command(fd, command, strlen(command));
}

void parse_pasv_array(char *resp, int arr[])
{
  int i;
  int j;
  int count = 0;
  int begin;
  
  for (i = 0; i < strlen(resp); i++)
    if (resp[i] == '(')
      break;

  begin = i + 1;

  for (j = 0; j < 6; j++)
    {
      arr[count] = atoi(resp + begin);
      count++;
      
      for (i = begin; i < strlen(resp); i++)
	if (resp[i] == ',')
	  break;

      begin = i + 1;
    } 
}

int connect_using_pasv(char *resp)
{
  int arr[6];
  char host[1024];
  int port;

  parse_pasv_array(resp, arr);
  sprintf(host, "%d.%d.%d.%d", arr[0], arr[1], arr[2], arr[3]);
  port = arr[4] * 256 + arr[5];

  printf("[-] Connecting %s:%d for ftp-data\n", host, port);
  return connect_to_server(host, port);
}

void recv_list_data(int fd)
{
  int read_mark = 0;
  int write_mark = 0;
  int begin = 0;
  int end;
  int read_bytes;
  int i;
  char buffer[8096];
  file_t *f;

  memset(buffer, 0, sizeof(buffer));

  while (1)
    {
      read_bytes = read(fd, buffer, sizeof(buffer) - 1);
      if (read_bytes <= 0)
	break;

      write_mark += read_bytes;

      for (i = read_mark; i < write_mark - 1; i++)
	{
	  if (buffer[i] == '\r' && buffer[i + 1] == '\n')
	    {
	      end = i + 1;
	      f = file_new(buffer + begin, end - begin - 1);
	      file_push(f);
	      begin = end + 1;
	      printf("%s\n", f->name);
	    }
	    
	}

    }
  
  close(fd);
}

void nlst_ready_callback(int fd, int code, char *resp, int length)
{
  recv_list_data(data_session.fd);
}

void nlst_callback(int fd, int code, char *resp, int length)
{
  char command[1024];
  printf("[-] Listing completed.\n");
    
  sprintf(command, "PASV\r\n");
  send_command(fd, command, strlen(command));
}

void pasv_callback(int fd, int code, char *resp, int length)
{
  char command[1024];
  file_t *f;
  
  //  printf("%s\n", resp);
  
  if (files == NULL && first_run)
    {
      printf("[-] Retrieving file list\n");
      memset(command, 0, sizeof(command));
      snprintf(command, 1023, "NLST\r\n");
      send_command(fd, command, strlen(command));
      
      data_session.fd = connect_using_pasv(resp);
    }
  else if (files == NULL && !first_run)
    {
      printf("[-] Completed.\n");
      exit(0);
    }
  else
    {
      f = file_peek();
      printf("[-] Fetch file: %s\n", f->name);
      memset(command, 0, sizeof(command));
      snprintf(command, 1023, "RETR %s\r\n", f->name);
      send_command(fd, command, strlen(command));

      data_session.fd = connect_using_pasv(resp);
      first_run = 0;
    }
}

void recv_file(int fd, file_t *f)
{
  FILE *fp;
  int read_bytes;
  int total_bytes = 0;
  char buffer[8096];
  char path[1024];

  memset(path, 0, sizeof(path));
  sprintf(path, "%s/%s", target_dir, f->name);

  fp = fopen(path, "wb");
  if (fp == NULL)
    perror("file open error");

  while(1)
    {
      read_bytes = read(fd, buffer, sizeof(buffer));
      if (read_bytes <= 0)
	break;

      total_bytes += read_bytes;

      fwrite(buffer, read_bytes, 1, fp);
    }

  if (fp != NULL)
    fclose(fp);

  file_free(f);
  close(fd);

  printf("[-] OK: %d bytes copied\n", total_bytes);
}

void retr_callback(int fd, int code, char *resp, int length)
{
  file_t *f;

  f = file_pop();
  if (code == 150) {
    recv_file(data_session.fd, f);
  }
  else {
    file_free(f);
    send_command(fd, "PASV\r\n", 6);
  }
}

void retr_complete_callback(int fd, int code, char *resp, int length)
{
  send_command(fd, "PASV\r\n", 6);
}

int match_command(char *command, int actual_code, int expected_code)
{
//	printf("actual: %d, expected: %d\n", actual_code, expected_code);
	return (!strncmp(current_command, command, 4) &&
			actual_code == expected_code) ? 1 : 0;
}

void command_callback(int fd)
{
	int eol;
	int code;
	int read_bytes;
	char *buffer = command_session.buffer;
	char *resp;
	char status[4096];
	int buffer_size = sizeof(command_session.buffer);
	
	read_bytes = read(fd,
					  buffer + command_session.write_mark,
					  buffer_size - command_session.write_mark);
	
	if (read_bytes < 0)
	{
		printf("[*] connection error occurred\n");
		exit(0);
	}
	
	if (read_bytes == 0)
		return;

	//printf("read: %d bytes\n", read_bytes);
	
	command_session.write_mark += read_bytes;
	resp = buffer + command_session.read_mark;

	eol = find_eol(resp, command_session.write_mark - command_session.read_mark);
	if (eol > 0)
	{
	  memset(status, 0, sizeof(status));
	  strncpy(status, resp, eol);
		code = atoi(resp);
		printf("[*] %s\n", status);
		if (match_command("", code, 220))
			welcome_callback(fd, code, resp, eol);
		else if (match_command("USER", code, 530))
		  {
		    printf("login failed.\n");
		    exit(0);
		  }
		else if (match_command("USER", code, 331))
			password_callback(fd, code, resp, eol);
		else if (match_command("PASS", code, 230))
			login_success_callback(fd, code, resp, eol);
		else if (match_command("PASS", code, 530))
		  {
		    printf("[-] Login failed.\n");
		    exit(0);
		  }
		else if (match_command("TYPE", code, 200))
			type_callback(fd, code, resp, eol);
		else if (match_command("CWD ", code, 250))
			cwd_callback(fd, code, resp, eol);
		else if (match_command("CWD ", code, 550))
		  {
		    printf("[-] Directory not found.\n");
		    exit(0);
		  }
		else if (match_command("PASV", code, 227))
			pasv_callback(fd, code, resp, eol);
		else if (match_command("NLST", code, 150))
		  nlst_ready_callback(fd, code, resp, eol);
		else if (match_command("NLST", code, 226))
		  nlst_callback(fd, code, resp, eol);
		else if (match_command("RETR", code, 150))
		  retr_callback(fd, code, resp, eol);
		else if (match_command("RETR", code, 550))
		  retr_callback(fd, code, resp, eol);
		else if (match_command("RETR", code, 226))
		  retr_complete_callback(fd, code, resp, eol);
		else if (code == 421)
		{
			printf("timeout. disconnected.\n");
			exit(0);
		}

		command_session.read_mark += (eol + 1);
	}

	//printf("write mark: %d, read mark: %d\n",
	//	   command_session.write_mark, command_session.read_mark);

	if (command_session.write_mark == command_session.read_mark)
	{
		command_session.write_mark = 0;
		command_session.read_mark = 0;
	}

}


int connect_to_server(char *hostname, int port) 
{
	int fd;
	int ret;
	struct sockaddr_in server_addr;
	struct in_addr in;
	struct hostent *entry;

	entry = gethostbyname(hostname);
	if (entry == NULL)
	{
		perror("domain resolver failed");
		return -1;
	}

	if (*entry->h_addr_list == NULL)
		return -1;

	in.s_addr = *((long int *)*entry->h_addr_list);

	fd = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (fd < 0)
	{
		perror("create socket");
		return -1;
	}

	memset(&server_addr, 0, sizeof(struct sockaddr_in));
	server_addr.sin_family = AF_INET;
	server_addr.sin_port = htons(port);
	server_addr.sin_addr = in;

	ret = connect(fd, (struct sockaddr*)&server_addr, sizeof(server_addr));
	if (ret < 0)
	{
		printf("connect failed\n");
		return ret;
	}

	//printf("[-] connecting to %s\n", hostname);
	return fd;
}


int main(int argc, char *argv[])
{
	int fd;
	memset(&command_session, 0, sizeof(session_t));
	memset(&data_session, 0, sizeof(session_t));
	memset(current_command, 0, sizeof(current_command));
	
	if (argc < 4)
	{
		printf("Usage: sftp [host] [id] [password] [remote dir] [local dir]\n");
		return 0;
	}

	username = argv[2];
	password = argv[3];
	directory = "./";
	target_dir = "./";
	
	if (argc > 4)
	  {
		directory = argv[4];
		target_dir = directory;
	  }

	if (argc > 5)
	  target_dir = argv[5];

	fd = connect_to_server(argv[1], 21);
	if (fd < 0)
		return -1;

	mkdir(target_dir, 0700);
	
	while(1)
	{
		command_callback(fd);
	}
	
	return 0;
}