/*
 * Copyright (C) 2016  J. Ali Harlow <ali@juiblex.co.uk>
 *
 * 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 2 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, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include "config.h"

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/tftp.h>

/*
 * A simple TFTP server suitable only for use on the loopback interface
 * (it has no support for retransmitting lost packets).
 */

char inbuf[SEGSIZE+4];
char outbuf[SEGSIZE+4];

int send_error(int s, const struct sockaddr *client, socklen_t client_addr_len,
	       unsigned short code,const char *message)
{
	int r;
	char *buf;
	size_t len;

	len = 4 + strlen(message) + 1;
	buf = malloc(len);
	((unsigned short *)buf)[0] = htons(ERROR);
	((unsigned short *)buf)[1] = htons(code);
	strcpy(buf + 4, message);
	r = sendto(s, buf, len, 0, client, client_addr_len);
	free(buf);

	return r;
}

int send_file(int s, const struct sockaddr *client, socklen_t client_addr_len,
	       const char *path, const char *mode)
{
	const char *dotdot, *p;
	FILE *fp;
	size_t nb;
	int block = 0;

	if (strcasecmp(mode, "octet")) {
		send_error(s, client, client_addr_len, EBADOP, "Bad mode");
		return -1;
	}

	if (!*path || *path == '/') {
		send_error(s, client, client_addr_len, EACCESS,
			   "Access denied");
		return -1;
	}

	for (p = path; *p;) {
		dotdot = strstr(p, "..");
		if (!dotdot)
			break;
		if ((dotdot == path || dotdot[-1] == '/') &&
		    (dotdot[2] == '/' || dotdot[2] == '\0')) {
			send_error(s, client, client_addr_len, EACCESS,
				   "Access denied");
			return -1;
		}
		p = dotdot + 2;
	}

	fp = fopen(path, "rb");

	if (!fp) {
		if (errno == ENOENT)
			send_error(s, client, client_addr_len, ENOTFOUND,
				   strerror(errno));
		else if (errno == EACCES)
			send_error(s, client, client_addr_len, EACCESS,
				   strerror(errno));
		else
			send_error(s, client, client_addr_len, EUNDEF,
				   strerror(errno));
		return -1;
	}

	((unsigned short *)outbuf)[0] = htons(DATA);

	while((nb = fread(outbuf + 4, 1, SEGSIZE, fp)) >= 0)
	{
		((unsigned short *)outbuf)[1] = htons(++block);
		if (sendto(s, outbuf, nb + 4, 0, client, client_addr_len) < 0) {
			perror("sendto");
			fclose(fp);
			return -1;
		}

		/* Discard ACKs */
		(void)recvfrom(s, inbuf, sizeof(inbuf), 0, NULL, NULL);
		
		if (nb < SEGSIZE)
			break;
	}

	fclose(fp);

	return 0;
}

void serve(int s)
{
	char *filename, *mode;
	struct sockaddr_storage client;
	socklen_t client_addr_len;
	ssize_t nb;

	for(;;) {
		client_addr_len = sizeof(client);
		nb = recvfrom(s, inbuf, sizeof(inbuf), 0,
			      (struct sockaddr *)&client, &client_addr_len);

		if (nb < 0) {
			perror("recvfrom");
			exit(1);
		}

		if (nb >= 2) {
			switch (ntohs(*(unsigned short *)inbuf)) {
			case RRQ:
				filename = inbuf + 2;
				mode = memchr(filename, '\0', nb - 2);
				if (!mode) {
					send_error(s,
						   (struct sockaddr *)&client,
						   client_addr_len, EBADOP,
						   "Bad request");
					break;
				}
				mode++;
				if (!memchr(mode, '\0', nb - (mode - inbuf))) {
					send_error(s,
						   (struct sockaddr *)&client,
						   client_addr_len, EBADOP,
						   "Bad request");
					break;
				}
				send_file(s, (struct sockaddr *)&client,
					  client_addr_len, filename, mode);
				break;
			case WRQ:
				send_error(s, (struct sockaddr *)&client,
					   client_addr_len, EACCESS,
					   "Access denied");
				break;
			case DATA:
			case ACK:
			case ERROR:
				break;
			}
		}
	}
}

main(int argc, char **argv)
{
	int s;
	pid_t pid;
	FILE *fp;
	in_port_t port;
	struct sockaddr_in sin = {0,};

	s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s < 0) {
		perror("socket");
		exit(1);
	}

	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = htonl(INADDR_LOOPBACK);

	for (port = 0;; port++) {
		sin.sin_port = htons(IPPORT_USERRESERVED + port);
		if (!bind(s, (struct sockaddr *)&sin, sizeof(sin)))
			break;
		if (errno != EADDRINUSE || port >= 1023) {
			perror("bind");
			exit(1);
		}
	}

	if (argc > 2) {
		fp = fopen(argv[2], "w");
		if (fp) {
			fprintf(fp, "%ld\n", (long)ntohs(sin.sin_port));
			fclose(fp);
		}

		pid = fork();
		if (pid) {
			fp = fopen(argv[1], "w");
			if (fp) {
				fprintf(fp, "%ld\n", (long)pid);
				fclose(fp);
			}
			exit(0);
		}
	} else
		printf("%d\n", ntohs(sin.sin_port));

	serve(s);

	exit(0);
}