#define __KERNEL__
#define MODULE


#include	<linux/module.h>
#include <linux/kernel.h>
#include <linux/netfilter.h>
#include <linux/skbuff.h>
#include <linux/ip.h>
#include <linux/netdevice.h> 
#include <linux/if_ether.h> 
#include <linux/if_packet.h>
#include <net/tcp.h>
#include <linux/netfilter_ipv4.h>


static struct nf_hook_ops nfho;
static struct nf_hook_ops nfhoRecv;
static struct nf_hook_ops nfhoSend;

short checksum(unsigned short* buffer, int size)
{
	unsigned long cksum = 0;
	 
	while(size>1)
	{
	    cksum += *buffer++;
	    size -= sizeof(unsigned short);
	}
	if(size)
	{
	    cksum += *(unsigned char*)buffer;
	}
	cksum = (cksum>>16) + (cksum&0xffff);  //將高16bit與低16bit相加
	cksum += (cksum>>16);             //將進(jìn)位到高位的16bit與低16bit 再相加
	 
	return (unsigned short)(~cksum);
}

unsigned int hook_func(unsigned int hooknum,
                       struct sk_buff **skb,
                       const struct net_device *in,
                       const struct net_device *out,
                       int (*okfn)(struct sk_buff *))
{
     struct sk_buff *sb = *skb;
    unsigned char src_ip[4];
    *(unsigned int *)src_ip = sb->nh.iph->saddr;
    printk("hook_func A packet from:%d.%d.%d.%d Detected!",
                 src_ip[0],src_ip[1],src_ip[2],src_ip[3]);
   switch(sb->nh.iph->protocol)
    {
       case IPPROTO_TCP:
           printk("It's a TCP PACKET\n");break;
       case IPPROTO_ICMP:
          printk("It's a ICMP PACKET\n");break;
      case IPPROTO_UDP:
         printk("It's a UDP PACKET\n");break;
    }
    
    return NF_ACCEPT;          
}

unsigned int hook_func_recv(unsigned int hooknum,
                       struct sk_buff **skb,
                       const struct net_device *in,
                       const struct net_device *out,
                       int (*okfn)(struct sk_buff *))
{
     struct sk_buff *sb = *skb;
    unsigned char src_ip[4];
    *(unsigned int *)src_ip = sb->nh.iph->saddr;
    printk("hook_func_recv A packet from:%d.%d.%d.%d Detected!",
                 src_ip[0],src_ip[1],src_ip[2],src_ip[3]);
   switch(sb->nh.iph->protocol)
    {
       case IPPROTO_TCP:
           printk("It's a TCP PACKET\n");break;
       case IPPROTO_ICMP:
          printk("It's a ICMP PACKET\n");break;
      case IPPROTO_UDP:
         printk("It's a UDP PACKET\n");break;
    }
    
    return NF_ACCEPT;          
}

unsigned int hook_func_send(unsigned int hooknum,
                       struct sk_buff **skb,
                       const struct net_device *in,
                       const struct net_device *out,
                       int (*okfn)(struct sk_buff *))
{
     struct sk_buff *sb = *skb;
    unsigned char src_ip[4];
    *(unsigned int *)src_ip = sb->nh.iph->saddr;
   //printk("hook_func_send A packet from:%d.%d.%d.%d Detected!",
     //           src_ip[0],src_ip[1],src_ip[2],src_ip[3]);
   switch(sb->nh.iph->protocol)
    {
       case IPPROTO_TCP:
           //printk("It's a TCP PACKET\n");break;
       case IPPROTO_ICMP:
       		//sb->nh.iph->tos=0xa0;
         // printk("It's a ICMP PACKET totlen=%d.\n",sb->nh.iph->tot_len);break;
      case IPPROTO_UDP:
      	if((sb->nh.iph->tot_len==200)||(sb->nh.iph->tot_len==280)
      	||(sb->nh.iph->tot_len==60)||(sb->nh.iph->tot_len==64))
      	{
      		sb->nh.iph->tos=0xa0;
      		sb->nh.iph->check=0;
      		
      		//sb->nh.iph->check=checksum((unsigned short *)sb->nh.iph,((sb->nh.iph->ihl)|0x0F)*4);  
      		sb->nh.iph->check=checksum((unsigned short *)sb->nh.iph,sb->nh.iph->ihl*4);
      		 //printk("###modify udp tos.\n");
      	}
         //printk("It's a UDP PACKET totlen=%d.\n",sb->nh.iph->tot_len);
         break;
         
    }
    
    return NF_ACCEPT;          
}

int NET_init_module()
{
  
    nfho.hook = hook_func;         
    nfho.hooknum  = NF_IP_PRE_ROUTING; 
    nfho.pf       = PF_INET;
    nfho.priority = NF_IP_PRI_FIRST; 
   // nf_register_hook(&nfho);
    
    nfhoRecv.hook = hook_func_recv;         
    nfhoRecv.hooknum  = NF_IP_LOCAL_IN; 
    nfhoRecv.pf       = PF_INET;
    nfhoRecv.priority = NF_IP_PRI_FIRST; 
 //  nf_register_hook(&nfhoRecv);
    
    
    nfhoSend.hook = hook_func_send;         
    nfhoSend.hooknum  = NF_IP_LOCAL_OUT; 
    nfhoSend.pf       = PF_INET;
    nfhoSend.priority = NF_IP_PRI_FIRST; 
   nf_register_hook(&nfhoSend);
    
    return 0;
}

void NET_cleanup_module()
{
   // nf_unregister_hook(&nfho);
   // nf_unregister_hook(&nfhoRecv);
    nf_unregister_hook(&nfhoSend);
}

module_init(NET_init_module);
module_exit(NET_cleanup_module);