#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/netpoll.h>

#include <linux/ip.h>
#include <linux/udp.h>
#include <linux/if_ether.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>

#include <net/udp.h>


/* Used to describe our Netfilter hooks */
struct nf_hook_ops pre_hook;  /* Incoming */
struct nf_hook_ops post_hook;  /* Outgoing */


unsigned int add_to_queue(struct sk_buff* skb);
void calc_transport_csum(struct sk_buff *skb);
int is_a2s_info_request(uint8_t* data, size_t len);
int is_inet_search_request(uint8_t* data, size_t len);


/* This is the hook function itself */
unsigned int hook_func(const struct nf_hook_ops *ops,
                        struct sk_buff *skb,
                        const struct net_device *in,
                        const struct net_device *out,
                        int (*okfn)(struct sk_buff *))
{
    struct iphdr* ip_header = ip_hdr(skb);

    // Handle only UDP IPv4 packets
    if (
        skb->protocol != htons(ETH_P_IP)
        || ip_header->version != 4
        || ip_header->protocol != IPPROTO_UDP
    ) {
        return NF_ACCEPT;
    }

    struct udphdr* udp_header = udp_hdr(skb);

    if (out != NULL && strcmp(out->name, "lo") != 0)
    {
        // Redirect all proxy outgoing traffic

        uint16_t src_port = ntohs(udp_header->source);
        uint16_t new_src_port = src_port;

        switch(src_port)
        {
            case 27921:
                new_src_port = 27021;
                break;
            case 27920:
                new_src_port = 27020;
                break;
            case 27919:
                new_src_port = 27019;
                break;
            case 27918:
                new_src_port = 27018;
                break;
            case 27917:
                new_src_port = 27017;
                break;
            case 27916:
                new_src_port = 27016;
                break;
            case 27915:
                new_src_port = 27015;
                break;
        }

        if (new_src_port != src_port)
        {
            udp_header->source = htons(new_src_port);
            calc_transport_csum(skb);
        }
    }
    else if (in != NULL && strcmp(in->name, "lo") != 0)
    {
        // Redirect incoming requests to proxy

        const size_t IP_UDP_HDR_SIZE = sizeof(struct iphdr) + sizeof(struct udphdr);

        size_t data_len = skb->len - IP_UDP_HDR_SIZE;
        uint8_t* data = (uint8_t *)skb_header_pointer(skb, IP_UDP_HDR_SIZE, 0, NULL);

        if (is_a2s_info_request(data, data_len) || is_inet_search_request(data, data_len))
        {
            //printk(KERN_INFO "A2S: %d", is_a2s_info_request(data, data_len));
            //printk(KERN_INFO "InetSearch: %d\n", is_inet_search_request(data, data_len));

            uint16_t dst_port = ntohs(udp_header->dest);
            uint16_t new_dst_port = dst_port;

            switch(dst_port)
            {
                case 27021:
                    new_dst_port = 27921;
                    break;
                case 27020:
                    new_dst_port = 27920;
                    break;
                case 27019:
                    new_dst_port = 27919;
                    break;
                case 27018:
                    new_dst_port = 27918;
                    break;
                case 27017:
                    new_dst_port = 27917;
                    break;
                case 27016:
                    new_dst_port = 27916;
                    break;
                case 27015:
                    new_dst_port = 27915;
                    break;
            }

            if (new_dst_port != dst_port)
            {
                udp_header->dest = htons(new_dst_port);
                udp_header->check = 0;
                //calc_transport_csum(skb);
            }
        }
    }

    return NF_ACCEPT;
}

unsigned int add_to_queue(struct sk_buff* skb)
{
    // Create eth header to requeue packet again

    struct ethhdr *eth = (struct ethhdr *) skb_push(skb, ETH_HLEN);
    skb_reset_mac_header(skb);
    skb->protocol = eth->h_proto = htons(ETH_P_IP);
    memcpy(eth->h_dest, skb->dev->dev_addr, ETH_ALEN);
    memset(eth->h_source, 0xFF, ETH_ALEN);

    dev_queue_xmit(skb);

    return NF_STOLEN;
}


void calc_transport_csum(struct sk_buff *skb)
{
    struct iphdr *ip_header = ip_hdr(skb);
    struct udphdr *udp_header = udp_hdr(skb);

    int transport_len = skb->len - skb_transport_offset(skb);

    ip_header->check = 0;
    ip_send_check(ip_header);
    udp_header->check = 0;
    udp_header->check = ~(csum_tcpudp_magic(ip_header->saddr, ip_header->daddr, transport_len, IPPROTO_UDP, 0));
}


int is_a2s_info_request(uint8_t* data, size_t len)
{
    static uint8_t payload[] = "\xFF\xFF\xFF\xFFTSource Engine Query";
    static size_t pl_len = sizeof(payload);

    if (len < pl_len) {
        return 0;
    }

    return memcmp(data, payload, pl_len) == 0;
}


#define INET_SEARCH_HDR_LEN 17
int is_inet_search_request(uint8_t* data, size_t len)
{
    static uint8_t payload[] = "InetSearchServerDetails";
    static size_t pl_len = sizeof(payload);

    if (len < INET_SEARCH_HDR_LEN + pl_len) {
        return 0;
    }

    // skip header and compare body only
    return memcmp(&data[INET_SEARCH_HDR_LEN], payload, pl_len) == 0;
}


/* Initialisation routine */
int init_module()
{
    /* Fill in our hook structure */
    pre_hook.pf = PF_INET;
    pre_hook.hook = hook_func;  /* Handler function */
    pre_hook.priority = NF_IP_PRI_FIRST;  /* Make our function first */
    pre_hook.hooknum  = NF_INET_PRE_ROUTING;  /* First hook for IPv4 */

    nf_register_hook(&pre_hook);

    post_hook.pf = PF_INET;
    post_hook.hook = hook_func;  /* Handler function */
    post_hook.priority = NF_IP_PRI_FIRST;  /* Make our function first */
    post_hook.hooknum  = NF_INET_POST_ROUTING;  /* First hook for IPv4 */

    nf_register_hook(&post_hook);

    return 0;
}

/* Cleanup routine */
void cleanup_module()
{
    nf_unregister_hook(&pre_hook);
    nf_unregister_hook(&post_hook);
}