#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <fcntl.h>

#include "nicStress.h"

static char *srvName = NULL;

static void usage(void)
{
	fprintf(stderr, "\n");
	fprintf(stderr, "%s \n", srvName);
	fprintf(stderr, "Usage: %s -i IPADDR | [-h] | [-c] | [-v] | [-s SIZE]\n", srvName);
	fprintf(stderr, "\t-h\tprint this message\n");
	fprintf(stderr, "\t-i\tIPv4 address\n");
    fprintf(stderr, "\t-s\tsize of buffer in bytes\n");
	fprintf(stderr, "\t-c\tdo not run as daemon\n");
	fprintf(stderr, "\t-v\tverbose\n");
	fprintf(stderr, "\n");
}

static unsigned char isIPv4Format(const char *str)
{
	int i, j, pos[3], len;
	char buf[IP_STR_BUF];
	unsigned int addr[4];
	unsigned char valid = TRUE;

	// Length check
	len = strlen(str);
	if (len > IP_STR_LEN || len < IP_STR_LEN_MIN)
		return FALSE;
	
	strncpy(buf, str, IP_STR_BUF);

	// Look for positions of delimiters
	for (i = 0, j = 0 ; i < len ; i++) {
		
		if ( buf[i] == '.') {
			// Exceed the limit
			if (j >= 3)
				return FALSE;
			// Record & Terminate string
			pos[j] = i;
			buf[i] = 0;
			j++;
		}
	}

	// Must have 3 dots at least or at most
	if (j != 3)
		return FALSE;

	// Convert strings
	addr[0] = strtol((const char *)&buf[0], NULL, 10);
	addr[1] = strtol((const char *)&buf[pos[0] + 1], NULL, 10);
	addr[2] = strtol((const char *)&buf[pos[1] + 1], NULL, 10);
	addr[3] = strtol((const char *)&buf[pos[2] + 1], NULL, 10);

	// Check range
	for (i = 0 ; i < 4 ; i++)
    if (addr[ i ] > 255) {
		valid = FALSE;
		break;
	}

	// Not ZERO
	snprintf(buf, IP_STR_BUF, "%d.%d.%d.%d", addr[0], addr[1], addr[2], addr[3]);
	if (strcmp(buf, str))
		valid = FALSE;

	return valid;
}


int main(int argc, char **argv)
{
	int ret = 0;
	int sfd, rbytes, wbytes;
	int addr_len = sizeof(struct sockaddr_in);
	struct sockaddr_in servaddr;
	struct sockaddr_in clntaddr;
	pid_t pid, sid;
	int verbose = 0;
	char daemon = 1;
	int c;
	char *rbuf = NULL;
	char *ipaddr = NULL, ipbuf[ IP_STR_BUF + 1 ];
	unsigned int size = MAX_UDP_BUFSIZE;


	srvName = strdup(argv[0]);

	// Handle parameters
	while ((c = getopt(argc, argv, "i:s:vch")) != EOF) {
		switch (c) {
			case 'i':
				strncpy(ipbuf, optarg, sizeof(ipbuf));
				if (isIPv4Format(ipbuf) == TRUE)
					ipaddr = ipbuf;
				break;
			case 'c':
				daemon = 0;
				break;
            case 's':
                size = strtol(optarg, NULL, 10);
                break;
			case 'v':
				verbose = 1;
				break;
			case 'h':
			default:
				usage();
				free(srvName);
				return -1;
		}
	}

	// Sanity check
	if (size <= 0 || size > MAX_UDP_BUFSIZE) {
		fprintf(stderr, "%s : Invalid size(%d)\n", srvName, size);
		free(srvName);
		return -1;
	}
	if (!ipaddr) {
		fprintf(stderr, "%s : Invalid ip addr\n", srvName);
		free(srvName);
		return -1;
	}

	// Become Daemon
	if (daemon) {
    
		// Fork a child process
		pid = fork();
		if (pid < 0) {
			fprintf(stderr, "%s : Cannot fork child process.\n", srvName);
			exit(1);
		}
		// Terminate parent process
		if (pid > 0)
			exit( 0 );
		
		umask(0);
		// Set SID
		sid = setsid();
		if (sid < 0) {
			exit( 1 );
		} 
		// Change location to root
		if (chdir("/") < 0) {
			exit(1);
		}

		// Close standard input, output, and error
		close(0);
		close(1);
		close(2);
	}

	// Allocate buffer
	rbuf = malloc(size);
	if (!rbuf) {
		fprintf(stderr, "%s : Out of memory\n", srvName);
		return -1;
	}

	sfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sfd < 0) {
		fprintf(stderr, "%s : Cannot open socket\n", srvName);
		ret = 1;
		goto errexit;
	}

	// Set address
	memset(&servaddr, 0, sizeof(struct sockaddr_in));
	servaddr.sin_family = AF_INET;
	servaddr.sin_port   = htons(UDP_PORT_NUM);
	ret = inet_aton(ipaddr, &servaddr.sin_addr);
	if (ret < 0) {
		fprintf( stderr, "%s : Cannot set IP address %s\n", srvName, ipaddr );
		goto errexit;
	}

	// Bind network port
	ret = bind(sfd, (struct sockaddr *)&servaddr, sizeof(struct sockaddr_in));
	if (ret < 0) {
		fprintf(stderr, "%s : Cannot bind network port %d\n", srvName, UDP_PORT_NUM);
		goto errexit;
	}

	while (1) {
		
		rbytes = recvfrom(sfd, rbuf, size, 0, (struct sockaddr*) &clntaddr, (socklen_t*) &addr_len);
		if (verbose)
			printf("recv %d from %s\n", rbytes, inet_ntoa(clntaddr.sin_addr));

		wbytes = sendto(sfd, rbuf, rbytes, 0, (struct sockaddr*) &clntaddr, addr_len);
		if (verbose)
			printf("send %d to %s\n", wbytes, inet_ntoa(clntaddr.sin_addr));
	}

errexit:
	// Close socket
	close(sfd);    

	// Free buffer
	if (rbuf)
		free(rbuf);

	if (srvName)
		free(srvName);

	return ret;
}