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

#include <netinet/in.h>
#include <arpa/inet.h>

#include <netinet/tcp.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/socket.h>
#include <resolv.h>
#include <fcntl.h>
#include <sys/time.h>
#include <assert.h>
#define MAXBUFF 4096 


uint32_t * handle_key_tranfer(int sock, struct sockaddr_in * server_addr) 
{
	uint64 b,B,g,p,A,K;
	uint32_t v[2];
	char addr[INET_ADDRSTRLEN];
	uint32_t * k = (uint32_t *)malloc(sizeof(uint32_t) * 4);

	srand(getpid()^time(NULL));

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

	g = ((uint64)ntohl(v[1])) << 32 | ntohl(v[0]);

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

	p = ((uint64)ntohl(v[1])) << 32 | ntohl(v[0]);

	inet_ntop(AF_INET,&(server_addr->sin_addr.s_addr),addr,INET_ADDRSTRLEN);
	// printf("Received g %llu and p %llu from %s\n", g, p, addr);

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

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

		A = ((uint64)ntohl(v[1])) << 32 | ntohl(v[0]);

		b = rand()%(p-1);
		B = expm(g,b,p);

		memcpy(v,(void *)&B,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));
		
		K = expm(A, b, p);

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


int main(int argc, char * argv[]) {
	int sock,  z=0, fd;
	char buff[MAXBUFF];
	uint32_t * key;
	struct sockaddr_in server_addr;
	int addrlen; 

	if (argc != 4) {
		printf("Usage: %s <IP-address> <port number> <file to send>\n",argv[0]);
		exit(1);
	}

	sock = socket(AF_INET, SOCK_STREAM, 0);
	assert(sock != -1);

	bzero(&server_addr, sizeof(server_addr));
	server_addr.sin_family = AF_INET;
	server_addr.sin_port = htons(atoi(argv[2]));

	assert( inet_aton(argv[1], (struct in_addr*)&server_addr.sin_addr.s_addr) != 0 );

	addrlen = sizeof(server_addr);

	fd = open(argv[3],O_RDONLY,0666); 
	assert(fd != -1);

	assert(connect(sock,(struct sockaddr *)&server_addr, addrlen) != -1);
	key = handle_key_tranfer(sock, &server_addr);
	// printf("Client key for tea = (%x,%x,%x,%x)\n",key[0],key[1],key[2],key[3]);

	while(1) {
		
		ssize_t mbytes, nbytes;
		int blocks;
		bzero(buff,MAXBUFF);
		nbytes=read(fd,buff,MAXBUFF);

		if (nbytes==-1) {
			perror("read");
			exit(1);
		}
		if (nbytes==0) {
			break;
		}

		blocks = nbytes / sizeof(uint64);
		for (int i = 0; i < blocks; i++) {
			encrypt_tea((uint32_t *)(buff + i * sizeof(uint64)), key);
		}

		mbytes = write(sock,buff, blocks * sizeof(uint64));
		if (mbytes == -1) {
			perror("write");
			exit(1);
		}
		z=z+mbytes;
	}

	printf("file %s sent (%d bytes)\n",argv[3],z);

	close(sock);
	exit(0);
}