#define BPF_NO_GLOBAL_DATA
#define __TARGET_ARCH_x86
#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_core_read.h>
#include <linux/if_ether.h>
#include <linux/udp.h>
#include <linux/net.h>
#include <linux/socket.h>
#include <sys/socket.h>
#include <linux/net.h>
#include <arpa/inet.h>
#include <netinet/ip.h>
#include "common.h"


/*
 * Helper:
 * Issue: invalid indirect read from stack R2 off
 *   Fix: check if all variables is initialised
 * Issue: R1 invalid mem access 'inv'
 *   Fix: the value can be NULL
 */

struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 256 * 1024 /* 256kb */);
} m_data SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 32768); 
    __type(key, uint16_t);
    __type(value, struct dns_answer);
} m_tid SEC(".maps");

static size_t get_labels2(struct __sk_buff *skb, size_t offset, struct event *s_event){
    char c;
    int qname_len = 0;
    //bpf_printk("labels off: %d", offset);
    bpf_skb_load_bytes(skb, offset, &c, 1); // Get the first byte, which is the length
    int pos = 1;
    while (c != '\0') {
        bpf_skb_load_bytes(skb, offset + pos++, &c, 1);

        if(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z')
            s_event->qname[qname_len] = c;
        else
            s_event->qname[qname_len] = '.';
        qname_len++;
        if (pos == 128 || c == '\0')
            break;
    }
    s_event->qname[qname_len - 1] = '\0';
    qname_len++;
    // dquery->qname_len = qname_len;
    // bpf_printk("qname: %s", s_event->qname);
    // bpf_printk("qname len: %d", qname_len);
    return qname_len;
}
static size_t get_labels(struct __sk_buff *skb, size_t offset, size_t end, struct event *s_event, struct query_section *s_query){
    //size_t len;
    char buf[256] = {0};
    char *c;
    int index = 0;
    size_t qname_len = 0; // Full length of the qname field

    bpf_skb_load_bytes(skb, offset, &buf, 41);
    c = buf;

    /*
     * The qname is composed by a the number of bytes then follow by the label
     * For instance, for the qname www.bucchino.org,
     * the first byte is the number of byte, here, it's 3, then we have www (in hex)
     * Then, we have the byte of 8 and follow by the label bucchino (size 8)
     * And to finish, we have 3 follow by org and we finish with the \0 character
     * For instance, the result is:
     * 03 77 77 77 08 62 75 63 63 68 69 6e 6f 03 6f 72 67 00
     */
    while (*(c++) != '\0') {
        if(*c >= 'a' && *c <= 'z')
            s_event->qname[index] = *c;
        else if(*c >= 'A' && *c <= 'Z')
            s_event->qname[index] = *c;
        else
            s_event->qname[index] = '.';
        index++;
        qname_len++;
    }
    s_event->qname[--index] = '\0';
    qname_len++; // For the null character
    bpf_printk("qname: %s", s_event->qname);

    return qname_len;
}

/*
 * This function get the query field and the return the length of it
 */
static size_t get_query_section(struct __sk_buff *skb, struct event *s_event, uint8_t offset){
    size_t len;
    size_t qname_len = 0; // Full length of the qname field
    uint16_t class, type;

    offset += sizeof(struct dnshdr);
    qname_len = get_labels2(skb, offset, s_event);

    // Get class and type
    len = qname_len;
    bpf_skb_load_bytes(skb, offset + qname_len, &type, sizeof(uint16_t));
    s_event->type = ntohs(type);
    len += 2;
    bpf_skb_load_bytes(skb, offset + qname_len + 2, &class, sizeof(uint16_t));
    s_event->class = ntohs(class);
    len += 2;

    return len;
}
static unsigned int get_answer(struct __sk_buff *skb, struct event *s_event, size_t tlen, unsigned int index){
    unsigned char buf[2] = {0}; // Need to be unsigned, otherwise, the result is fffff
    unsigned int offset = index;

    // Get the 2 first bytes to identify if it's a message compression or not
    if(bpf_skb_load_bytes(skb, tlen, &buf, 2) < 0)
        return 0;
    tlen += 2; // For the message compression

    /*
     * According to the RFC 1035 (https://datatracker.ietf.org/doc/html/rfc1035#section-4.1.4)
     * In the section 4.1.4, message compression, the first two bits are set at 11 (0xc),
     * that's means, it's a pointer.
     * For instance, the two bytes 0xc00c, 0xc (11) it's the pointer and 0x00c is the position in the DNS header
     */
    if (buf[0] == 0xc0){
        /*
         * I cannot read the labels, because the eBPF verifier considerate as a infinity loop
         */

        /* 
         * According to the RFC 1035, the structure of answer is like that:
         * https://datatracker.ietf.org/doc/html/rfc1035#section-4.1.3
         */

        // Get the class and type
        if ((void*)(offset) >= MAX_UDP_PAYLOAD - sizeof(uint16_t))
            return 0;

        bpf_skb_load_bytes(skb, tlen, s_event->buf + offset, sizeof(uint16_t));
        uint16_t type = s_event->buf[0] + (s_event->buf[1] << 8);
        tlen += 2;
        if ((void*)(offset += 2) >= MAX_UDP_PAYLOAD - sizeof(uint16_t))
            return 0;
        //offset += 2;

	// For class
        if(bpf_skb_load_bytes(skb, tlen, s_event->buf + offset, sizeof(uint16_t)) < 0)
            return 0;
        tlen += 2;
        if ((void*)(offset += 2) >= MAX_UDP_PAYLOAD - sizeof(uint16_t))
            return 0;

        // Get ttl
        if(bpf_skb_load_bytes(skb, tlen, s_event->buf + offset, sizeof(uint32_t)) < 0)
            return 0;
        if ((void*)(offset += 4) >= MAX_UDP_PAYLOAD - sizeof(uint32_t))
            return 0;
        tlen += 4;

        // Get data size
        uint16_t size;
        bpf_skb_load_bytes(skb, tlen, &size, sizeof(uint16_t));
        bpf_skb_load_bytes(skb, tlen, s_event->buf + offset, sizeof(uint16_t));
        if ((void*)(offset += 2) >= MAX_UDP_PAYLOAD - sizeof(uint16_t))
            return 0;
        tlen += 2;

        uint32_t data;
        if (s_event->type == 1) { // -> A
            bpf_skb_load_bytes(skb, tlen, s_event->buf + offset, sizeof(uint32_t)); 
        }
        if ((void*)(offset += ntohs(size)) >= MAX_UDP_PAYLOAD - sizeof(uint16_t))
            return 0;
        tlen += ntohs(size);
    }
    else {
        // get_labels2(skb, sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr) + sizeof(struct dnshdr), s_event);
    }
    bpf_printk("End offset: %d", offset);
    return offset;
}
/*
 * https://datatracker.ietf.org/doc/html/rfc1035
 */
static int dnsquery(struct __sk_buff *skb, struct ethhdr eth, struct iphdr ip, struct udphdr udp, int dport, int sport){
    struct event *s_event;
    struct dnshdr dns = {0};
    struct query_section dquery = {0};

    /* Get DNS header */
    bpf_skb_load_bytes(skb, sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr), &dns, sizeof(struct dnshdr));

    // Check OpCode
    uint16_t qr = ntohs(dns.flags) & 0xF000; // Get the QR code: 0 -> query, 1 -> response
    /* If it's not a query, we do not continue */
    if(qr != 0x0)
        return 0;

    if (ntohs(dns.nbQuestions) == 0)
        return 0;

    s_event = bpf_ringbuf_reserve(&m_data, sizeof(*s_event), 0);
    if (!s_event)
       return 0;

    s_event->req_type = REQ_QUERY;

    /* Get IP header */
    s_event->client = ip.saddr;

    s_event->tid = ntohs(dns.transactionID);

    /* Get the query section */
    uint8_t tlen = sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr);
    size_t query_len = get_query_section(skb, s_event, tlen);

    // https://docs.cilium.io/en/stable/reference-guides/bpf/progtypes/
    s_event->dport = dport;
    s_event->sport = sport;

    bpf_ringbuf_submit(s_event, 0);

    return 0;
}

static void dnsanswer_old(struct __sk_buff *skb, struct iphdr ip, struct udphdr udp, int dport, int sport){
     char buf[256] = {0}; // Max dns domain name length
    //__u16 udplen = 0U;

    // Check with ip.len
    //const __u32 tot_len = ntohs(ip.tot_len);
    // __u32 payload_len = (ntohs(ip.tot_len) - sizeof(struct iphdr) - sizeof(struct udphdr));
    __u32 payload_len = (ntohs(ip.tot_len) - sizeof(struct iphdr) - sizeof(struct udphdr)) & 0xff;

    if (payload_len > skb->len)
        return;

    bpf_printk("payload len %d", payload_len);
    if (payload_len <= -1) {
        bpf_printk("payload len %d", payload_len);
    }
    
    // Get udp len
    //udplen = ntohs(udp.len);

    //bpf_printk("udp len: %d", udplen);

    //if (udplen <= 0)
    //    return 0;

    /*if (offset + udplen > skb->len) {
        bpf_printk("outbound");
        plen = sizeof(struct dnshdr);
    }*/
    /*long err = bpf_skb_load_bytes(skb, offset, &buf, payload_len);
    if(err < 0){
        bpf_printk("failed");
        // bpf_ringbuf_discard(s_event, 0);
        return;
    }*/
    // Cast to dnshdr
    // dns = (struct dnshdr*)buf;

}


/*
TODO: je recupere tout le skb->data, grace au skb->len
grace a ca, j'aurai tout le payload udp et toute les donnees dns
je pourrai facilement parcourir les data
*/
static void dnsanswer(struct __sk_buff *skb, struct iphdr ip, struct udphdr udp, int dport, int sport){
    struct event *s_event;
    struct dnshdr dns;
    uint16_t tid = 0U;
    uint32_t offset = sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr);

    // Load dns header
    if (bpf_skb_load_bytes(skb, offset, &dns, 12) < 0)
        return;

    // Check OpCode
    uint16_t qr = ntohs(dns.flags) & 0xF000; // Get the QR code: 0 -> query, 1 -> response
    if(qr != 0x8000) // Not a response, we do not continue
        return;

    if (ntohs(dns.nbQuestions) == 0)
        return;

    s_event = bpf_ringbuf_reserve(&m_data, sizeof(*s_event), 0);
    if (!s_event)
       return;

    s_event->req_type = REQ_ANSWER;

    /* Get IP header */
    s_event->client = ip.daddr;

    s_event->dport = dport;
    s_event->sport = sport;

    /* Get the Transaction ID */
    s_event->tid = ntohs(dns.transactionID);
   
    /* Get the query section */
    size_t query_len = get_query_section(skb, s_event, offset);

    /* Get answer section */
    uint16_t ans = ntohs(dns.nbAnswerRRs);
    if (ans <= 0){
        bpf_ringbuf_discard(s_event, 0);
        return;
    }

    if (ans > 0){
        /*
         * We get a least the 5 last answer
         * In the RFC 1035 (https://datatracker.ietf.org/doc/html/rfc1035#section-2.3.4) the max udp payload is 512 bytes
         * The program limit size of the answer
         */
        offset += sizeof(struct dnshdr) + query_len; // For the pos in the answer section in the skb
        unsigned int offset_ans = 0;
        for (uint16_t i = 0; i < ans; i++){
            offset_ans += get_answer(skb, s_event, offset, offset_ans);
	    offset += offset_ans + 2; // +2 for the message compression
            //offset_ans += offset_ans;
            // For eBPF verifier, to be sure we leave the loop
            if (i == ans || i == 5 || offset_ans >= 512)
                break;
        }
        s_event->numAns = ans;
    }
    if (ntohs(dns.nbAuthorityRRs) > 0){

    }
    /*
     * In the user space, if the haven't have the answer, we can have an error
     * The solution is to push to the ring buffer and the answer is store in
     * the struct event
     * Or, we push to the ring buffer and the query only with the map
     * but, if we haven't have the answer, we need print the query
     */

    /*
        Pour recuperer les infos:
         1 - dans le getquery, on push dans le ringbuffer et dans le userspace, on recupere aussi la reponse
         mais si la reponse, nous l'avons pas encore, ca fail et dans le get answer on push dans une map
         2 - on push dans le ring buffer quand on a la reponse avec la requette car c'est dans le field query
         cependant, si on a pas la reponse, on n'aura jamais la query
         3 - dans le get query et get answer, on push dans le ring buffer et tout est store dans le struct event
     */

    /* Get the answer */
    bpf_ringbuf_submit(s_event, 0);
}

/*
 * skb -> http://oldvger.kernel.org/~davem/skb_data.html
 */

SEC("socket")
int detect_dns(struct __sk_buff *skb) {
    //void *data = (void *)(long)skb->data;
    //void *data_end = (void *)(long)skb->data_end;
    //struct ethhdr *eth2 = data;
    struct ethhdr eth = {0};
    struct iphdr ip = {0};
    struct udphdr udp = {0};
    __u32 h_proto, p;
    __u32 dport;
    __u32 sport;

    //if (data + sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr) > data_end)
    //    return 0;

    if (skb->len < sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr))
        return 0;

    //bpf_skb_load_bytes(skb, 12, &p, 2);
    bpf_skb_load_bytes(skb, 0, &eth, sizeof(struct ethhdr));
    p = eth.h_proto;
    if (ntohs(p) != ETH_P_IP)
        return 0;

    // bpf_printk("ip: %d",ntohs(p));

    //ip = (struct iphdr*)(data + sizeof(struct ethhdr));
    bpf_skb_load_bytes(skb, sizeof(struct ethhdr), &ip, sizeof(struct iphdr));

    //bpf_skb_load_bytes(data, sizeof(struct ethhdr), ip, sizeof(struct ip));
    h_proto = ip.protocol;
    // If not UDP packet
    if (h_proto != 17) 
        return 0;

    // bpf_printk("proto: %d", h_proto);
    //udp = (struct udphdr*)(data + sizeof(struct ethhdr) + sizeof(struct iphdr));

    bpf_skb_load_bytes(skb, sizeof(struct ethhdr) + sizeof(struct iphdr), &udp, sizeof(struct udphdr));

    if (udp.len == 0)
        return 0;

    // Check if DNS port
    dport = ntohs(udp.dest);
    sport = ntohs(udp.source);
    
    if (dport == 53)
        dnsquery(skb, eth, ip, udp, dport, sport);
    else if(sport == 53)
        dnsanswer(skb, ip, udp, dport, sport);

    return 0;
}
/*SEC("xdp")
int detect_dns(struct xdp_md *ctx){
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;
    struct ethhdr *eth = data;
    struct iphdr *ip;
    struct udphdr *udp;
    __u16 h_proto;
    __u16 dport;

    if (data + sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr) > data_end)
        return XDP_DROP;

    ip = (struct iphdr*)(data + sizeof(struct ethhdr));
    udp = (struct udphdr*)(data + sizeof(struct ethhdr) + sizeof(struct iphdr));

    h_proto = ip->protocol;
    // If not UDP packet
    if (h_proto != 17) 
        return XDP_PASS;

    // Check if DNS port
    dport = udp->dest;
    bpf_printk("Dport: %d", (dport));

    return XDP_PASS;
}*/
char LICENSE[] SEC("license") = "GPL";
