/*

 * 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);

}