#include <pcap.h>
#include <string.h>
#include <stdlib.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <net/ethernet.h>
#include "Diffie_hellman.c"

#define MAXBYTES2CAPTURE 2048

typedef struct{
	unsigned char client_ip[16];
	unsigned char server_ip[16];
	pcap_t *p;
}ip_rule;

struct psd_header{
	unsigned int saddr;
	unsigned int daddr;
	char mbz;
	char ptcl;
	unsigned short tcpl;
};
Mim_key dh_key;
unsigned char server_mac[]={0x00,0x0c,0x29,0x83,0xfc,0xdf};
unsigned char client_mac[]={0x00,0x0c,0x29,0xb5,0x23,0x98};
unsigned char self_mac[]={0x00, 0x0c, 0x29, 0x07, 0x31, 0x61};
//计算tcp校验和
uint16_t calc_checksum(void *pkt, int len) {
	uint16_t *buf = (uint16_t*)pkt;
	uint32_t checksum = 0;
	while (len > 1)
    {
        checksum += *buf++;
        checksum = (checksum >> 16) + (checksum & 0xffff);
		len -= 2;
	}
    if (len)
    {
        checksum += *((uint8_t*)buf);
        checksum = (checksum >> 16) + (checksum & 0xffff);
	}
	return (uint16_t)((~checksum) & 0xffff);
}
//伪造
void set_psdheader(struct psd_header* ph,struct iphdr *ip, uint16_t tl){
	ph->saddr=ip->saddr;
	ph->daddr=ip->daddr;
	ph->mbz=0;
	ph->ptcl=6;
	ph->tcpl=htons(tl);
}

//捕包
void tcp_callback(u_char *arg, const struct pcap_pkthdr* pkthdr, const u_char * packet){
	struct ether_header *ethernet;
	struct iphdr *ip;
	struct tcphdr *tcp;
	char srcip_str[16];
    int hdr_len = ETHER_HDR_LEN + sizeof(struct iphdr) + sizeof(struct tcphdr) + 12;
    int dl=pkthdr->len-hdr_len;

	memset(srcip_str,0,16);
	ip_rule * ip_r=(ip_rule*)arg;
	//数据帧头部
	ethernet = (struct ether_header*)(packet);
	//ip头部
	ip = (struct iphdr*)(packet + ETHER_HDR_LEN);
	//tcp头部
	tcp = (struct tcphdr*)(packet + ETHER_HDR_LEN+ sizeof(struct iphdr));
	inet_ntop(AF_INET,&(ip->saddr),srcip_str,16);
	memcpy(ethernet->ether_shost, self_mac, 6);

	//处理服务器端发送的信息
    if(strncmp(srcip_str, ip_r->server_ip, strlen(srcip_str))==0)
    {
		//模
		if(strncmp(packet+hdr_len,"pri",3)==0)
        {
			mpz_set_str(dh_key.prime, packet+hdr_len+3, 16);
		}
        else if(strncmp(packet+hdr_len,"data",4)==0)
        {
			char *buff=packet+hdr_len;
			char *plain=(char*)malloc(dl-48);
			bzero(plain, sizeof(buff));
			char key_string[128],key_value[32];
			mpz_get_str(key_string, 16, dh_key.key_for_server);
			memcpy(key_value, str2hex(key_string), sizeof(key_value));

			//解密
			unsigned char iv[12],add[16],tag[16];
			memcpy(add,buff+4,16);
			memcpy(iv,buff+16+4,12);
			memset(tag,0,16);
			int tmp=aes_gcm_ad(key_value, 32,iv, 12,buff+44+4, dl-44-4,add, 16,tag, plain);
			printf("[server]%s",plain);

			//加密
			bzero(key_string, sizeof(key_string));
			mpz_get_str(key_string, 16, dh_key.key_for_client);
			memcpy(key_value, str2hex(key_string), sizeof(key_value));
			tmp=aes_gcm_ae(key_value, 32,iv,12,plain,dl-48,add, 16,buff+44+4, tag);

			//伪造tcp头
			uint16_t tcp_len =pkthdr->len-ETHER_HDR_LEN-sizeof(struct iphdr);
			unsigned char * data_for_checksum=(unsigned char*)malloc(tcp_len+sizeof(struct psd_header));
			struct psd_header ph;
			bzero(data_for_checksum, tcp_len+sizeof(ph));
			set_psdheader(&ph, ip, tcp_len);
			memcpy(data_for_checksum, (void*)(&ph), sizeof(ph));
			tcp->check=0;
			memcpy(data_for_checksum+sizeof(ph), tcp, tcp_len);
			uint16_t cs= calc_checksum(data_for_checksum, tcp_len + sizeof(ph));
			tcp->check=cs;
		}
		else if(strncmp(packet+hdr_len,"pub",3)==0)
        {
			mpz_t key_from_s;
			mpz_init(key_from_s);
			mpz_set_str(key_from_s,packet+hdr_len+3,16);
			mpz_powm(dh_key.key_for_server, key_from_s, dh_key.private_key, dh_key.prime);
			unsigned char pub[64];
			mpz_get_str(pub,16,dh_key.public_key);
			memcpy(packet+hdr_len+3,pub,sizeof(pub));

			//伪造tcp头
			uint16_t tcp_len =pkthdr->len-ETHER_HDR_LEN-sizeof(struct iphdr);
			unsigned char * data_for_checksum=(unsigned char*)malloc(tcp_len+sizeof(struct psd_header));
			struct psd_header ph;
			bzero(data_for_checksum, tcp_len+sizeof(ph));
			set_psdheader(&ph, ip, tcp_len);
			memcpy(data_for_checksum, (void*)(&ph), sizeof(ph));
			tcp->check=0;
			memcpy(data_for_checksum+sizeof(ph), tcp, tcp_len);
			unsigned short cs= calc_checksum(data_for_checksum, tcp_len + sizeof(ph));
			tcp->check=cs;
		}
		memcpy(ethernet->ether_dhost,client_mac,6);
	}
	//处理客户端
	else
    {
        if(strncmp(packet+hdr_len,"data",4)==0)
        {
			//decrypt
			char *buff=packet+hdr_len;
			char *plain=(char*)malloc(dl-48);
			bzero(plain, sizeof(buff));
			char key_string[128],key_value[32];
			mpz_get_str(key_string, 16, dh_key.key_for_client);
			memcpy(key_value, str2hex(key_string), sizeof(key_value));

			//解密
			unsigned char iv[12],add[16],tag[16];
			memcpy(add,buff+4,16);
			memcpy(iv,buff+16+4,12);
			memset(tag,0,16);
			int tmp=aes_gcm_ad(key_value, 32,iv, 12,buff+44+4, dl-44-4,add, 16,tag, plain);
			printf("[client]%s",plain);

			//加密
			bzero(key_string, sizeof(key_string));
			mpz_get_str(key_string, 16, dh_key.key_for_server);
			memcpy(key_value, str2hex(key_string), sizeof(key_value));
			tmp=aes_gcm_ae(key_value, 32,iv,12,plain,dl-48,add, 16,buff+44+4, tag);

			//伪造
			uint16_t tcp_len =pkthdr->len-ETHER_HDR_LEN-sizeof(struct iphdr);
			unsigned char * data_for_checksum=(unsigned char*)malloc(tcp_len+sizeof(struct psd_header));
			struct psd_header ph;
			bzero(data_for_checksum, tcp_len+sizeof(ph));
			set_psdheader(&ph, ip, tcp_len);
			memcpy(data_for_checksum, (void*)(&ph), sizeof(ph));
			tcp->check=0;
			memcpy(data_for_checksum+sizeof(ph), tcp, tcp_len);
			uint16_t cs= calc_checksum(data_for_checksum, tcp_len + sizeof(ph));
			tcp->check=cs;
		}
		else if(strncmp(packet+hdr_len,"pub",3)==0)
        {
            unsigned char pub[64];
			generate_key(32,&dh_key);
			mpz_t key_from_c;
			mpz_init(key_from_c);
			mpz_set_str(key_from_c,packet+hdr_len+3,16);
			mpz_powm(dh_key.key_for_client, key_from_c, dh_key.private_key, dh_key.prime);
			mpz_powm(dh_key.public_key, dh_key.base, dh_key.private_key, dh_key.prime);
			memcpy(ethernet->ether_dhost,server_mac,6);

			mpz_get_str(pub,16,dh_key.public_key);
			memcpy(packet+hdr_len+3,pub,sizeof(pub));

			uint16_t tcp_len =pkthdr->len-ETHER_HDR_LEN-sizeof(struct iphdr);
			unsigned char * data_for_checksum=(unsigned char*)malloc(tcp_len+sizeof(struct psd_header));
			struct psd_header ph;
			bzero(data_for_checksum, tcp_len+sizeof(ph));
			set_psdheader(&ph, ip, tcp_len);
			memcpy(data_for_checksum, (void*)(&ph), sizeof(ph));
			tcp->check=0;
			memcpy(data_for_checksum+sizeof(ph), tcp, tcp_len);
			uint16_t cs= calc_checksum(data_for_checksum, tcp_len + sizeof(ph));
			tcp->check=cs;
		}
		memcpy(ethernet->ether_dhost,server_mac,6);
	}
	pcap_sendpacket(ip_r->p,packet,pkthdr->len);
	return;
}

/* main(): Main function. Opens network interface and calls pcap_loop() */
int main(int argc, char *argv[] ){
	int i=0, count=0;
	pcap_t *descr = NULL;
	char errbuf[PCAP_ERRBUF_SIZE], *device=NULL;
    struct bpf_program filter;

	memset(errbuf,0,PCAP_ERRBUF_SIZE);
	init_numbers(&dh_key);
    mpz_set_ui(dh_key.base, (unsigned long)2);
	dh_key.urand = fopen("/dev/urandom","r");

	//获得网卡信息
	if ( (device = pcap_lookupdev(errbuf)) == NULL)
    {
		fprintf(stderr, "ERROR: %s\n", errbuf);
		exit(1);
	}
	printf("Opening device %s\n", device);
	//打开网卡
	if ( (descr = pcap_open_live(device, MAXBYTES2CAPTURE, 1,  512, errbuf)) == NULL)
    {
		fprintf(stderr, "ERROR: %s\n", errbuf);
		exit(1);
	}

    //设置过滤信息
	char rule[128];
	memset(rule,0,128);
	strncat(rule,"(src host ",10);
	strncat(rule,argv[1],strlen(argv[1]));
	strncat(rule," and dst host ",14);
	strncat(rule,argv[2],strlen(argv[2]));
	strncat(rule,") or (src host ",15);
	strncat(rule,argv[2],strlen(argv[2]));
	strncat(rule," and dst host ",14);
	strncat(rule,argv[1],strlen(argv[1]));
	strncat(rule, ")", 1);
    if(pcap_compile(descr,&filter,rule,1,0)<0)
    {
        printf("error\n");
        return 0;
    }
    if(pcap_setfilter(descr,&filter)<0)
    {
        printf("error\n");
        return 0;
    }
	ip_rule ip_r;
	ip_r.p=descr;
	memset(ip_r.client_ip,0,15);
	memcpy(ip_r.client_ip, argv[1],strlen(argv[1]));
	memset(ip_r.server_ip,0,15);
	memcpy(ip_r.server_ip, argv[2],strlen(argv[2]));
	//循环捕包
	if (pcap_loop(descr, -1, tcp_callback, (u_char *) &ip_r) == -1)
    {
		fprintf(stderr, "ERROR: %s\n", pcap_geterr(descr) );
		exit(1);
	}

	return 0;
}