/*  Next File Transfer Server
    Copyright (C) 2020  Ian Cowburn

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty 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, see <http://www.gnu.org/licenses/>.

*/
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>

#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <netdb.h>

static const char *name;


static void *Malloc(size_t len)
{
    void *p;

    p = malloc(len);

    if (!p)
    {
    	perror(name);
	exit(EXIT_FAILURE);
    }

    return p;
}


static char *GetLine(char *buff, size_t len)
{
    size_t l;

    if (feof(stdin))
    {
	return NULL;
    }

    printf("NFT> ");

    if (!fgets(buff, len, stdin))
    {
	return NULL;
    }

    l = strlen(buff);

    if (l > 0 && buff[l-1]=='\n')
    {
	buff[l-1]=0;
    }

    if (l)
    {
	return buff;
    }
    else
    {
	return GetLine(buff, len);
    }
}


static int Socket(int p)
{
    struct sockaddr_in addr;
    int sock;

    if ((sock = socket(AF_INET,SOCK_STREAM,0))==-1)
    {
	perror("socket");
	exit(EXIT_FAILURE);
    }

    /* Bind port to address
    */
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(p);

    if (bind(sock, (void*)&addr, sizeof(addr))==-1)
    {
	perror("bind");
	exit(EXIT_FAILURE);
    }

    if (listen(sock, 5)==-1)
    {
	perror("listen");
	exit(EXIT_FAILURE);
    }

    return sock;
}


static int Read(int sock, char *buff, size_t len)
{
    while(len > 0)
    {
	size_t rd;

	rd = read(sock, buff, len);

	if (rd == -1)
	{
	    perror(name);
	    return 0;
	}
	else if (rd == 0)
	{
	    return 0;
	}
	else
	{
	    len -= rd;
	    buff += rd;
	}
    }

    return 1;
}


static int Write(int sock, const char *buff, size_t len)
{
    while(len > 0)
    {
	size_t wr;

	wr = write(sock, buff, len);

	if (wr == -1)
	{
	    perror(name);
	    return 0;
	}
	else if (wr == 0)
	{
	    return 0;
	}
	else
	{
	    len -= wr;
	    buff += wr;
	}
    }

    return 1;
}


static int WriteBlock(int sock, const char *block, size_t len)
{
    char buff[32];

    if (len > 999999)
    {
    	fprintf(stderr, "%s: length of block %zu too large\n", name, len);
	return 0;
    }

    snprintf(buff, sizeof buff, "%.6zu", len); 

    if (!Write(sock, buff, 6))
    {
    	return 0;
    }

    if (!Write(sock, block, len))
    {
    	return 0;
    }

    return 1;
}


static size_t ReadSize(int sock)
{
    char buff[32];

    if (!Read(sock, buff, 6))
    {
    	return 0;
    }

    return atoi(buff);
}


static size_t FileSize(const char *path)
{
    struct stat st;

    if (stat(path, &st) == 0)
    {
    	return st.st_size;
    }
    else
    {
    	return 0;
    }
}


static int Put(int sock, const char *path)
{
    int quit = 0;
    char status[3] = {0};
    size_t len;
    int fd;

    if ((len = FileSize(path)) == 0)
    {
    	perror(path);
    }
    else
    {
	char *buff;
	int fd;

	buff = Malloc(len);

	fd = open(path, O_RDONLY);
	read(fd, buff, len);
	close(fd);

	if (!Write(sock, "PF", 2))
	{
	    quit = 1;
	}

	if (!quit && !WriteBlock(sock, path, strlen(path)))
	{
	    quit = 1;
	}

	if (!quit && !WriteBlock(sock, buff, len))
	{
	    quit = 1;
	}

	if (!quit && Read(sock, status, 2))
	{
	    if (strcmp(status, "!E") == 0)
	    {
		printf("Command failed\n");
	    }
	    else if (strcmp(status, "OK") == 0)
	    {
		printf("OK\n");
	    }
	    else
	    {
	    	printf("Unknown status '%s'\n", status);
	    }
	}
	else
	{
	    quit = 1;
	}

	free(buff);
    }

    return quit;
}


static int Get(int sock, const char *path)
{
    int quit = 0;
    char status[3] = {0};
    int fd;

    if ((fd = open(path, O_WRONLY|O_CREAT|O_TRUNC, 0666)) == -1)
    {
    	perror(path);
	return 0;
    }

    if (!Write(sock, "GF", 2))
    {
	quit = 1;
    }

    if (!quit && !WriteBlock(sock, path, strlen(path)))
    {
	quit = 1;
    }

    if (!quit && Read(sock, status, 2))
    {
	if (strcmp(status, "!E") == 0)
	{
	    printf("Command failed\n");
	}
	else if (strcmp(status, "OK") == 0)
	{
	    size_t len;

	    printf("OK\n");

	    if ((len = ReadSize(sock)) > 0)
	    {
		char *buff;

		buff = Malloc(len);

		if (Read(sock, buff, len))
		{
		    write(fd, buff, len);
		}
		else
		{
		    quit = 1;
		}
	    }
	}
	else
	{
	    printf("Unknown status '%s'\n", status);
	}
    }
    else
    {
	quit = 1;
    }

    close(fd);

    return quit;
}


static int Cd(int sock, const char *path)
{
    int quit = 0;
    char status[3] = {0};
    size_t len;
    int fd;

    if (path)
    {
	if (!Write(sock, "CD", 2))
	{
	    quit = 1;
	}

	if (!quit && !WriteBlock(sock, path, strlen(path)))
	{
	    quit = 1;
	}

	if (!quit && Read(sock, status, 2))
	{
	    if (strcmp(status, "!E") == 0)
	    {
		printf("Command failed\n");
	    }
	    else if (strcmp(status, "OK") == 0)
	    {
		printf("OK\n");
	    }
	    else
	    {
	    	printf("Unknown status '%s'\n", status);
	    }
	}
	else
	{
	    quit = 1;
	}
    }
    else
    {
    	printf("Missing path\n");
    }

    return quit;
}


static int Mkdir(int sock, const char *path)
{
    int quit = 0;
    char status[3] = {0};

    if (path)
    {
	if (!Write(sock, "MD", 2))
	{
	    quit = 1;
	}

	if (!quit && !WriteBlock(sock, path, strlen(path)))
	{
	    quit = 1;
	}

	if (!quit && Read(sock, status, 2))
	{
	    if (strcmp(status, "!E") == 0)
	    {
		printf("Command failed\n");
	    }
	    else if (strcmp(status, "OK") == 0)
	    {
		printf("OK\n");
	    }
	    else
	    {
	    	printf("Unknown status '%s'\n", status);
	    }
	}
	else
	{
	    quit = 1;
	}
    }
    else
    {
    	printf("Missing path\n");
    }

    return quit;
}


static int Ls(int sock)
{
    int quit = 0;
    char status[3] = {0};

    if (!Write(sock, "LS", 2))
    {
	quit = 1;
    }

    if (!quit && Read(sock, status, 2))
    {
	if (strcmp(status, "!E") == 0)
	{
	    printf("Command failed\n");
	}
	else if (strcmp(status, "OK") == 0)
	{
	    int no;
	    int f;
	    char buff[1024];

	    no = ReadSize(sock);

	    for(f = 0; f < no; f++)
	    {
		int size;

		size = ReadSize(sock);
		Read(sock, buff, size);
		buff[size] = 0;

		printf("%s\n", buff);
	    }
	}
	else
	{
	    printf("Unknown status '%s'\n", status);
	}
    }
    else
    {
	quit = 1;
    }

    return quit;
}


int main(int argc, char *argv[])
{
    int sock;
    int port;
    int quit;

    if ((name = strrchr(argv[0], '/')))
    {
    	name++;
    }
    else
    {
    	name = argv[0];
    }

    if (argc<2)
    {
	fprintf(stderr, "%s: usage %s port\n", name, name);
	exit(EXIT_FAILURE);
    }

    port = atoi(argv[1]);
    sock = Socket(port);

    quit = 0;

    while(!quit)
    {
	char buff[1024];
	char *p;
	int closed;
	int fd;

	printf("%s: listening on port %d\n", name, port);

	if ((fd = accept(sock, NULL, 0)) == -1)
	{
	    perror("accept");
	    exit(EXIT_FAILURE);
	}

	printf("%s: accepted connection\n", name);

	closed = 0;

	while(!closed && (p = GetLine(buff, sizeof buff)))
	{
	    char *tok;

	    tok = strtok(buff, " \t");

	    if (strcmp(tok, "put") == 0)
	    {
		closed = Put(fd, strtok(NULL, " \t"));
	    }
	    else if (strcmp(tok, "get") == 0)
	    {
		closed = Get(fd, strtok(NULL, " \t"));
	    }
	    else if (strcmp(tok, "cd") == 0)
	    {
		closed = Cd(fd, strtok(NULL, " \t"));
	    }
	    else if (strcmp(tok, "mkdir") == 0)
	    {
		closed = Mkdir(fd, strtok(NULL, " \t"));
	    }
	    else if (strcmp(tok, "ls") == 0)
	    {
		closed = Ls(fd);
	    }
	    else if (strcmp(tok, "close") == 0)
	    {
		closed = 1;
	    }
	    else if (strcmp(tok, "exit") == 0)
	    {
		closed = 1;
		quit = 1;
	    }
	    else if (strcmp(tok, "help") == 0)
	    {
	    	printf("put <file>\n");
	    	printf("get <file>\n");
	    	printf("cd <path>\n");
	    	printf("mkdir <path>\n");
	    	printf("ls\n");
	    	printf("close\n");
	    	printf("exit\n");
	    }
	    else
	    {
		printf("Unrecognised command '%s'\n", tok);
	    }
	}

	close(fd);

	if (!p)
	{
	    quit = 1;
	}
    }

    close(sock);

    return EXIT_SUCCESS;
}