/*

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

 }