#include <netdb.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include "Diffie_hellman.c"
#include "aes.h"
#define MAX_size 256
#define Sockaddr struct sockaddr

#define DEBUF
// Function designed for chat between client and server.

void recv_msg(int sockfd, char * buf, unsigned char *key){
	char buff[MAX_size];
	//add(16)+iv(12)+tag(16)+data
	unsigned char iv[12],add[16],tag[16];
	int l=read(sockfd, buff, sizeof(buff));
	memcpy(add,buff+4,16);
	memcpy(iv,buff+16+4,12);
	memset(tag,0,16);
	int tmp=aes_gcm_ad(key, 32, iv, 12,buff+44+4, l-44-4, add, 16, tag, buf);

}

void send_msg(int sockfd, char * buff, int con_len, unsigned char *key){
    //add(16)+iv(12)+tag(16)+data
    unsigned char * cryptbuf=(unsigned char*)malloc(con_len + 44 + 4);
    memset(cryptbuf, 0, con_len + 44 + 4);
    //add+iv
    unsigned char iv[12],add[16],tag[16];
    memset(tag,0,16);
    for(int j=0;j<16;j++)
        add[j]=rand()%256;
    for(int j=0;j<12;j++)
        iv[j]=rand()%256;
    //加密
    int tmp;
    tmp = aes_gcm_ae(key, 32, iv, 12, buff, con_len, add, 16, cryptbuf + 44 + 4, tag);
    memcpy(cryptbuf + 4, add, 16);
    memcpy(cryptbuf + 20, iv, 12);
    memcpy(cryptbuf + 32, tag, 16);
    memcpy(cryptbuf, "data", 4);
    int f;
    f = write(sockfd, cryptbuf, con_len + 44 + 4);
}

void DH_key_exchange_for_server(int sockfd, unsigned char* aes_key){
	Dh_key dh_key;
	char buff[MAX_size];
    init_numbers(&dh_key);

    dh_key.urand = fopen("/dev/urandom","r");
    if(NULL == dh_key.urand)
    {
        fprintf(stderr, "Failed to init randomization\n");
        exit(1);
    }
    init_prime(32,&dh_key);
	//send p
	bzero(buff, MAX_size);
	memcpy(buff,"pri",3);
	mpz_get_str(buff+3,16,dh_key.prime);
	write(sockfd,buff,sizeof(buff));
	gmp_printf("p:%Zd\n",dh_key.prime);

	//receive A
	bzero(buff, MAX_size);
	read(sockfd,buff,sizeof(buff));
	mpz_t key_from_c;
	mpz_init(key_from_c);
	mpz_set_str(key_from_c,buff+3,16);
	gmp_printf("A:%Zd\n",key_from_c);

	//calculate private key b
    generate_key(32, &dh_key);
	gmp_printf("b:%Zd\n",dh_key.private_key);

	//B=g^b(mod p)
    mpz_powm(dh_key.public_key, dh_key.base, dh_key.private_key, dh_key.prime);
	gmp_printf("B:%Zd\n",dh_key.public_key);

	//send B
	bzero(buff, MAX_size);
	memcpy(buff,"pub",3);
	mpz_get_str(buff+3,16,dh_key.public_key);
	write(sockfd, buff, sizeof(buff));

	//calc s=g^(ab)
	bzero(buff,sizeof(buff));
    mpz_powm(dh_key.key, key_from_c, dh_key.private_key, dh_key.prime);
    mpz_get_str(buff,16,dh_key.key);
    memcpy(aes_key,str2hex(buff),32);
	fclose(dh_key.urand);
    clear_numbers(&dh_key);
}

void echo(int sockfd)
{
	char buff[MAX_size];
	int n;
	char aes_key[32];
	char echobuf[MAX_size+5];
	bzero(echobuf, sizeof(echobuf));
	// negotiate key for aes 256
    DH_key_exchange_for_server(sockfd, aes_key);
    printf("\n");
    print_hex(aes_key, 32);
	for (;;) {
		bzero(buff, sizeof(buff));
		// read message from client
		recv_msg(sockfd, buff, aes_key);
		// print buffer
		printf("Client: %s", buff);

		// generate echo
		memcpy(echobuf,"ECHO:",5);
		memcpy(echobuf+5,buff,strlen(buff));
		// send echo
		send_msg(sockfd, echobuf, strlen(echobuf),aes_key);
		bzero(echobuf, sizeof(echobuf));
		// Exit to stop service
		if (strncmp("exit", buff, 4) == 0) {
			printf("Server Exit...\n");
			break;
		}
	}
}

// main echo
int main(int argc, char **argv)
{
    int sockfd, connfd, len;
    struct sockaddr_in servaddr, cli;

    // socket create and verification
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd == -1) {
		printf("socket creation failed...\n");
		exit(0);
	}
	else
		printf("Socket successfully created..\n");
	bzero(&servaddr, sizeof(servaddr));

	// assign IP, PORT
	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
	servaddr.sin_port = htons(atoi(argv[1]));

	// Binding newly created socket to given IP and verification
	if ((bind(sockfd, (Sockaddr*)&servaddr, sizeof(servaddr))) != 0) {
		printf("socket bind failed\n");
		exit(0);
	}
	else
		printf("Socket successfully binded..\n");

	// start listen
	if ((listen(sockfd, 5)) != 0) {
		printf("listen failed\n");
		exit(0);
	}
	else
		printf("start listening\n");
	len = sizeof(cli);

	// Accept the data packet from client and verification
	connfd = accept(sockfd, (Sockaddr*)&cli, &len);
	if (connfd < 0) {
		printf("server acccept failed\n");
		exit(0);
	}
	else
		printf("connected\n");

	// callback echo
    echo(connfd);

	// close socket
	close(sockfd);
}