#include <memory.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include "Diffie_hellman.c"
#include "aes.h"
#include <unistd.h>


#define AES_DEBUG
#define MAX_size 256
#define PORT 8080
#define Sockaddr struct sockaddr

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;
    //encrypt
    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, "daaa", 4);
    int f;
    f = write(sockfd, cryptbuf, con_len + 44 + 4);
}


void DH_key_exchange_for_client(int sockfd, unsigned char * key_aes){
	Dh_key client_key;
    init_numbers(&client_key);
	client_key.urand = fopen("/dev/urandom","r");
    mpz_set_ui(client_key.base, (unsigned long)2);
	char buff[MAX_size];
	memset(buff, 0, MAX_size);

	//receive p
	memset(buff, 0, MAX_size);
	read(sockfd,buff,sizeof(buff));
    mpz_set_str(client_key.prime,buff+3,16);
	gmp_printf("p:%Zd\n",client_key.prime);

    //generate private key a
	generate_key(32, &client_key);
	gmp_printf("a:%Zd\n",client_key.private_key);

	//calculate public key A=g^a(mod p)
	mpz_powm(client_key.public_key, client_key.base, client_key.private_key, client_key.prime);
	gmp_printf("A:%Zd\n",client_key.public_key);

	//send A
	bzero(buff, MAX_size);
	memcpy(buff,"pud",3);
	mpz_get_str(buff+3,16,client_key.public_key);
	write(sockfd,buff,sizeof(buff));

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

	//calc s=g^(ab)
    mpz_powm(client_key.key, key_from_s, client_key.private_key, client_key.prime);

	//key -> char*
	bzero(buff, MAX_size);
	mpz_get_str(buff,16,client_key.key);
	memcpy(key_aes,str2hex(buff),32);
	fclose(client_key.urand);
	clear_numbers(&client_key);
}

void func(int sockfd)
{
	char buff[MAX_size];
	int n=0;
	unsigned char aes_key[32];

    DH_key_exchange_for_client(sockfd, aes_key);
	print_hex(aes_key, 32);
	for (;;) {
		bzero(buff, sizeof(buff));
		printf("Enter str : ");
		n = 0;
		while ((buff[n++] = getchar()) != '\n');
		send_msg(sockfd, buff, strlen(buff),aes_key);

		bzero(buff, sizeof(buff));

		recv_msg(sockfd, buff, aes_key);
		printf("From Server : %s", buff);

		if ((strncmp(buff, "exit", 4)) == 0) {
			printf("Client Exit\n");
			break;
		}
	}
}

int main(int argc,char **argv)
{
	if(argc!=3){
		puts("./client ip port");
		return 0;
	}
	int sockfd, connfd;
	struct sockaddr_in servaddr, cli;

	//创建套接字
	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));

	//绑定地址
	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = inet_addr(argv[1]);
	servaddr.sin_port = htons(atoi(argv[2]));
	if (connect(sockfd, (Sockaddr*)&servaddr, sizeof(servaddr)) != 0) {
		printf("connection with the server failed\n");
		exit(0);
	}
	else
		printf("connected to the server\n");

	//接收
	func(sockfd);
	close(sockfd);
}