#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"


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

static size_t get_labels(struct __sk_buff *skb, size_t offset, struct event *s_event){
    char c;
    int qname_len = 0;

    bpf_skb_load_bytes(skb, offset, &c, 1); // Get the first byte, which is the length
    int pos = 1;
    /*
     * 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') {
        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 if(c >= '0' && c <= '9')
            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++;
    // bpf_printk("qname len: %d", qname_len);
    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_labels(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;
}
/*
 * 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};

    /* 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);
    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;
}

/*
 * This function read the dns answer section
 * All the answer is store in the buf variable and handle in the user space
 */
static void dnsanswer(struct __sk_buff *skb, struct iphdr ip, struct udphdr udp, int dport, int sport){
    struct event *s_event;
    struct dnshdr dns;
    uint32_t offset = sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr);
    size_t tlen = ntohs(udp.len);
    unsigned int index = 0U;

    if (tlen < 0 || tlen >= 256)
        return;
    
    // Load dns header
    if (bpf_skb_load_bytes(skb, offset, &dns, sizeof(struct dnshdr)) < 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;
    }

    s_event->numAns = ans;

    /*
     * Load query and answers
     * It's a little dirty to do that, to load byte by byte,
     * otherwise, I have an issue with the eBPF verifier
     */
    offset += sizeof(struct dnshdr) + query_len;
    for (index = 0; index < tlen || index == MAX_UDP_PAYLOAD; index++)
        bpf_skb_load_bytes(skb, offset + index, s_event->buf + index, 1);

    bpf_ringbuf_submit(s_event, 0);
}

/*
 * skb -> http://oldvger.kernel.org/~davem/skb_data.html
 */
SEC("socket")
int detect_dns(struct __sk_buff *skb) {
    struct ethhdr eth = {0};
    struct iphdr ip = {0};
    struct udphdr udp = {0};
    __u32 h_proto, p;
    __u32 dport;
    __u32 sport;


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

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

    bpf_skb_load_bytes(skb, sizeof(struct ethhdr), &ip, sizeof(struct iphdr));

    h_proto = ip.protocol;
    // If not UDP packet
    if (h_proto != 17) 
        return 0;

    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;
}
char LICENSE[] SEC("license") = "GPL";