#include "../lib_df/df.h"
#include "../lib_tea/tea.h"


#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/socket.h>
#include <resolv.h>
#include <arpa/inet.h>
#include <time.h>
#include <string.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <assert.h>
#define MAXBUFF 4096
#define PRIME_LENGTH 14


const uint64 prime_list[] = {
	4611686018427394499,
	4611686018427395927,
	4611686018427402023,
	4611686018427402599,
	4611686018427407519,
	4611686018427407939,
	4611686018427408827,
	4611686018427412487,
	4611686018427412619,
	4611686018427413183,
	4611686018427415163,
	4611686018427419483,
	4611686018427419627,
	4611686018427427067
};

uint32_t * handle_key_tranfer(int sock, struct sockaddr_in * client_addr) 
{
	uint64 a,A,g,p,B,K;
	uint32_t v[2];
	uint32_t * k = (uint32_t *)malloc(sizeof(uint32_t) * 4);

	srand(getpid()^time(NULL));
	p = prime_list[rand()%PRIME_LENGTH];
	g = generateur(p);
	// printf("g = %llu p = %llu\n", g, p);

	memcpy((void *)v,(void *)&g,2*sizeof(uint32_t));

	v[0] = htonl(v[0]);
	v[1] = htonl(v[1]);

	write(sock, (void *)v, sizeof(uint32_t));
	write(sock, (void *)(v+1), sizeof(uint32_t));
	
	memcpy(v,(void *)&p,2*sizeof(uint32_t));

	v[0] = htonl(v[0]);
	v[1] = htonl(v[1]);

	write(sock, (void *)v, sizeof(uint32_t));
	write(sock, (void *)(v+1), sizeof(uint32_t));

	for (int i = 0;i < 2; i++){

		a = rand()%(p-1);
		A = expm(g,a,p);

		memcpy(v,(void *)&A,2*sizeof(uint32_t));

		v[0] = htonl(v[0]);
		v[1] = htonl(v[1]);

		write(sock, (void *)v, sizeof(uint32_t));
		write(sock, (void *)(v+1), sizeof(uint32_t));

		read(sock, (void *)v, sizeof(uint32_t));
		read(sock, (void *)(v+1), sizeof(uint32_t));

		B = ((uint64)ntohl(v[1])) << 32 | ntohl(v[0]);
		K = expm(B, a, p);

		memcpy((void*)(k+2*i),(void*)&K,sizeof(uint64));
	}
	// printf("Server key = %llx%llx\n", k[1], k[0]);

	return k;
}


int handle_transferred_file(int sock,struct sockaddr_in * client_addr)
{
	char bufname[256];
	uint32_t * key;
	char buff[MAXBUFF];
	int fd, blocks;
	ssize_t nbytes,z;

	char addr[INET_ADDRSTRLEN];
	snprintf(bufname,256,"%s.%d",inet_ntoa(client_addr->sin_addr),ntohs(client_addr->sin_port));
	fd = open(bufname,O_WRONLY|O_CREAT|O_TRUNC,0666); 
	if (fd == -1) {
		perror("open");
		exit(1);
	}

	key = handle_key_tranfer(sock, client_addr);
	// printf("Server key for tea = (%x,%x,%x,%x)\n",key[0],key[1],key[2],key[3]);

	while (1) {
		//bzero(buff,MAXBUFF);
		nbytes = read(sock,buff,MAXBUFF);
		blocks = nbytes / sizeof(uint64);
		for (int i = 0; i < blocks; i++) {
			decrypt_tea((uint32_t *)(buff + i * sizeof(uint64)), key);
		}
		if (nbytes == -1) {
			perror("read");
			exit(1);
		}
		if (nbytes==0) break;

		write(fd,buff,blocks * sizeof(uint64));
		z=z+nbytes;
	}

	inet_ntop(AF_INET,&(client_addr->sin_addr.s_addr),addr,INET_ADDRSTRLEN);
	printf("%ld bytes recieved from %s\n", z, addr);
	shutdown(sock,SHUT_RDWR);
	close(sock);

	exit(0);

}


int main(int argc, char * argv[]) 
{
	int 				/* file descriptor for the tranferred file */ 
		listen_sock,    /* socket for listenning connection */
		service_sock,   /* socket for a connected client */
		nbytes, z=0;
	int yes=1;

	struct sockaddr_in server_addr, 
					   client_addr;

	int addr_len; 

	struct sigaction sa;

	if (argc != 2) {
		printf("Usage: %s <port number>\n",argv[0]);
		exit(1);
	}

	sigemptyset(&sa.sa_mask);
	sa.sa_flags = SA_NOCLDWAIT;
	sa.sa_handler = SIG_DFL;
	assert (sigaction(SIGCHLD, &sa, NULL) != -1);



	listen_sock = socket(AF_INET, SOCK_STREAM, 0);
	assert( listen_sock > 0);

	assert (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int)) != -1);


	bzero(&server_addr, sizeof(server_addr));
	server_addr.sin_family = AF_INET;
	server_addr.sin_port = htons(atoi(argv[1]));
	server_addr.sin_addr.s_addr = INADDR_ANY;
	addr_len = sizeof(server_addr);

	assert(bind(listen_sock, (struct sockaddr *)&server_addr, sizeof(struct sockaddr)) != -1);

	assert(listen(listen_sock,10) == 0);



	while(1){

		service_sock = accept(listen_sock,(struct sockaddr *)&client_addr , &addr_len);

		assert(service_sock >= 0);

		pid_t p = fork();
		if (p==0){
			close(listen_sock);
			handle_transferred_file(service_sock,&client_addr);
		}
	}
}