#include <winsock2.h>
#include <divert.h>
#include <stdio.h>
#include <time.h>

typedef struct _TCP
{
	WORD SrcPort;
	WORD DstPort;
	DWORD SeqNum;
	DWORD AckNum;
	BYTE DataOff;
	BYTE Flags;
	WORD Window;
	WORD Chksum;
	WORD UrgPtr;
} 	TCP;

#define MAXBUF  0xFFFF

typedef struct ip_hdr
{
	unsigned char  ip_header_len:4;  // 4-bit header length (in 32-bit words) normally=5 (Means 20 Bytes may be 24 also)
	unsigned char  ip_version   :4;  // 4-bit IPv4 version
	unsigned char  ip_tos;           // IP type of service
	unsigned short ip_total_length;  // Total length
	unsigned short ip_id;            // Unique identifier 
	
	unsigned char  ip_frag_offset   :5;        // Fragment offset field
	
	unsigned char  ip_more_fragment :1;
	unsigned char  ip_dont_fragment :1;
	unsigned char  ip_reserved_zero :1;
	
	unsigned char  ip_frag_offset1;    //fragment offset
	
	unsigned char  ip_ttl;           // Time to live
	unsigned char  ip_protocol;      // Protocol(TCP,UDP etc)
	unsigned short ip_checksum;      // IP checksum
	unsigned int   ip_srcaddr;       // Source address
	unsigned int   ip_destaddr;      // Source address
}   IPV4_HDR, *PIPV4_HDR, FAR * LPIPV4_HDR;


TCP header;
typedef struct tcp_header 
{ 
	unsigned short source_port;  // source port 
	unsigned short dest_port;    // destination port 
	unsigned int   sequence;     // sequence number - 32 bits 
	unsigned int   acknowledge;  // acknowledgement number - 32 bits 
		
	unsigned char  ns   :1;          //Nonce Sum Flag Added in RFC 3540.
	unsigned char  reserved_part1:3; //according to rfc
	unsigned char  data_offset:4;    /*The number of 32-bit words in the TCP header. 
									   This indicates where the data begins. 
									   The length of the TCP header is always a multiple 
									   of 32 bits.*/
			
	unsigned char  fin  :1;      //Finish Flag
	unsigned char  syn  :1;      //Synchronise Flag
	unsigned char  rst  :1;      //Reset Flag
	unsigned char  psh  :1;      //Push Flag 
	unsigned char  ack  :1;      //Acknowledgement Flag 
	unsigned char  urg  :1;      //Urgent Flag
	
	unsigned char  ecn  :1;      //ECN-Echo Flag
	unsigned char  cwr  :1;      //Congestion Window Reduced Flag
	
	////////////////////////////////
	
	unsigned short window;  // window 
	unsigned short checksum;  // checksum 
	unsigned short urgent_pointer;  // urgent pointer 
}   TCP_HDR , *PTCP_HDR , FAR * LPTCP_HDR , TCPHeader , TCP_HEADER;

int generateICMPPacket(char *pPacket, char *strSrcIP, char *strDstIP)
{
	UINT packet_len;
	UINT local_ip;
	UINT remote_ip;
	
	DivertHelperParseIPv4Address(strSrcIP, &local_ip);
	DivertHelperParseIPv4Address(strDstIP, &remote_ip);

	PDIVERT_IPHDR ip_header = (PDIVERT_IPHDR) pPacket;
	ip_header->HdrLength = 5;
	ip_header->Version = 4;
	ip_header->TOS = 0;
	ip_header->Length = htons(sizeof(DIVERT_IPHDR) + sizeof(DIVERT_ICMPHDR) + 32);
	ip_header->Id = 0x1234;
	DIVERT_IPHDR_SET_FRAGOFF(ip_header, 0);
	DIVERT_IPHDR_SET_MF(ip_header, 0);
	DIVERT_IPHDR_SET_DF(ip_header, 0);
	DIVERT_IPHDR_SET_RESERVED(ip_header, 0);
	ip_header->TTL = 64;
	ip_header->Protocol = 1; //ICMP
	ip_header->Checksum = 0;
	ip_header->SrcAddr = htonl(local_ip);
	ip_header->DstAddr = htonl(remote_ip);

	PDIVERT_ICMPHDR icmp_header = (PDIVERT_ICMPHDR) ((PBYTE) ip_header + sizeof(DIVERT_IPHDR));
	icmp_header->Type = 8;
	icmp_header->Code = 0;
	icmp_header->Checksum = 0;
	icmp_header->Body = htonl(0x00010012);

	PBYTE icmp_data = (PBYTE) icmp_header + sizeof(DIVERT_ICMPHDR);
	for (int i = 0; i < 32; i ++)
	{
		*icmp_data = 'a' + i % 23;
		icmp_data ++;
	}

	packet_len = sizeof(DIVERT_IPHDR) + sizeof(DIVERT_ICMPHDR) + 32;
	DivertHelperCalcChecksums((PVOID) pPacket, packet_len, 0);

	return packet_len;
}

int generateTCPPacket(char *pPacket, char *strSrcIP, char *strDstIP, short iDstPort)
{
	UINT packet_len;
	IPV4_HDR *v4hdr = NULL;
	TCP_HDR  *tcphdr = NULL;
	char *data = NULL;
	int payload = 0;

	srand((unsigned) time(0));

	v4hdr = (IPV4_HDR *) pPacket; //lets point to the ip header portion
	v4hdr->ip_version = 4;
	v4hdr->ip_header_len = 5;
	v4hdr->ip_tos = 0;
	v4hdr->ip_total_length = htons (sizeof(IPV4_HDR) + sizeof(TCP_HDR) + payload);
	v4hdr->ip_id = 0; //htons(2);
	v4hdr->ip_frag_offset = 0;
	v4hdr->ip_frag_offset1 = 0;
	v4hdr->ip_reserved_zero = 0;
	v4hdr->ip_dont_fragment = 0;
	v4hdr->ip_more_fragment = 0;
	v4hdr->ip_ttl = 123;
	v4hdr->ip_protocol = IPPROTO_TCP;
	v4hdr->ip_srcaddr  = inet_addr(strSrcIP);
	v4hdr->ip_destaddr = inet_addr(strDstIP);
	v4hdr->ip_checksum = 0;

	tcphdr = (TCP_HDR *) &pPacket[sizeof(IPV4_HDR)]; //get the pointer to the tcp header in the packet

	tcphdr->source_port = htons(rand() + 32767);
	tcphdr->dest_port = htons(iDstPort);
	tcphdr->sequence = htonl(rand() + 0x12345678);
	tcphdr->acknowledge = 0;

	tcphdr->ns = 0;
	tcphdr->reserved_part1 = 0;
	tcphdr->data_offset = 5;

	tcphdr->cwr = 0;
	tcphdr->ecn = 0;
	tcphdr->urg = 0;
	tcphdr->ack = 0;
	tcphdr->psh = 0;
	tcphdr->rst = 0;
	tcphdr->syn = 1;
	tcphdr->fin = 0;
	tcphdr->ns = 0;

	tcphdr->window = htons(1480);
	tcphdr->urgent_pointer = 0;
	tcphdr->checksum = 0;


	// Initialize the TCP payload to some rubbish
	data = &pPacket[sizeof(IPV4_HDR) + sizeof(TCP_HDR)];
	memset(data, '^', payload);

	packet_len = sizeof(IPV4_HDR) + sizeof(TCP_HDR) + payload;
	DivertHelperCalcChecksums((PVOID) pPacket, packet_len, 0);

	return packet_len;
}

int main(int argc,char **argv)
{
	

	HANDLE handle;  		// Divert handle
	DIVERT_ADDRESS addr;	// Packet address
	char packet[MAXBUF];	// Packet buffer
	UINT packet_len;

	addr.IfIdx = 0;
	addr.SubIfIdx = 0;
	addr.Direction = DIVERT_DIRECTION_OUTBOUND;

	char strSrcIP[30] = "127.0.0.1";
	char strDstIP[30] = "127.0.0.1";
	char strDstPort[10] = "80";

	if (argc == 1)
	{
		printf("Usage:\nPacketInjecter.exe [Source IP] [Destination IP] [Destination Port]");
		return 0;
	}
	else if (argc == 4)
	{
		strcpy_s(strSrcIP, argv[1]);
		strcpy_s(strDstIP, argv[2]);
		strcpy_s(strDstPort, argv[3]);
	}

	short iDstPort = (short) atoi(strDstPort);

	packet_len = generateTCPPacket(packet, strSrcIP, strDstIP, iDstPort);

	handle = DivertOpen("true", (DIVERT_LAYER)0, 1, 0);   // Open some filter
	if (handle == INVALID_HANDLE_VALUE)
	{
		// Handle error
		exit(1);
	}

// 	// Read packet.
// 	if (!DivertRecv(handle, packet, sizeof(packet), &addr, &packet_len))
// 	{
// 		fprintf(stderr, "warning: failed to read packet (%d)\n",
// 			GetLastError());
// 		DivertClose(handle);
// 		return 0;
// 	}

	// Send packet.
	if (!DivertSend(handle, packet, packet_len, &addr, NULL))
	{
		// Handle send error
		//BOOL a = DivertSend(handle, (PVOID) packet, packet_len, &addr, NULL);
		DivertClose(handle);
		DWORD dwError = GetLastError();
		printf("DivertSend Error.\n");
	}
	else
	{
		DivertClose(handle);
		printf("DivertSend Success.\n");
	}

	return 0;
}