#include <stdio.h>
#include <stdlib.h>
#include <linux/bpf.h>
#include <bpf/bpf.h>
#include <bpf/libbpf.h>
#include <linux/perf_event.h>
#include <argp.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <fcntl.h>
#include <net/if.h> /* if_nametoindex */
#include <linux/if_ether.h>
#include <linux/if_packet.h>
#include <arpa/inet.h>

#include "common.h"

static struct arguments arguments;
static int running = 1;

struct arguments {
    char *interface;
    char *filename;
    int to_output;
};

/*
 * Functions for arguments
 * https://www.gnu.org/software/libc/manual/html_node/Argp-Example-3.html
 */

static char doc[] = "DNS Trace usage:";
static char args_doc[] = "ARG1 ARG2";

static error_t parse_opts(int key, char *arg, struct argp_state *state){
    struct arguments *arguments = state->input;
    switch(key){
        case 'f':
            arguments->filename = arg;
            arguments->to_output = 0;
            break;
        case 'i':
            arguments->interface = arg;
            break;
        case 'o':
            arguments->to_output = 1;
            break;
        case ARGP_KEY_ARG:
            break;
        case ARGP_KEY_END:
            break;
        default:
            return ARGP_ERR_UNKNOWN;

    }
    return 0;
}

struct arguments parse_args(int argc, char *argv[]){
    static const struct argp_option opts[] = {
        {"interface", 'i', "INTERFACE", 0, "Ifname for listening"},
        {"filename", 'f', "FILENAME", 0, "Save result to logs"},
        {"to-output", 'o', NULL, 0, "Print to output"},
        {NULL, 'h', NULL, OPTION_HIDDEN, "help"},
        {},
    };
    struct arguments arguments;
    arguments.interface = NULL;
    arguments.filename = NULL;
    arguments.to_output = 1;
    static struct argp argp = {opts, parse_opts, args_doc, doc};

    argp_parse(&argp, argc, argv, 0, 0, &arguments);
    return arguments;
}

/* End functions arguments */

static void signalHandler(int signum){
    running = 0;
}

static int open_raw_sock(const char *name) {
	struct sockaddr_ll sll;
	int sock;

	sock = socket(PF_PACKET, SOCK_RAW | SOCK_NONBLOCK | SOCK_CLOEXEC, htons(ETH_P_ALL));
	if (sock < 0) {
		printf("cannot create raw socket\n");
		return -1;
	}

	memset(&sll, 0, sizeof(sll));
	sll.sll_family = AF_PACKET;
	sll.sll_ifindex = if_nametoindex(name);
	sll.sll_protocol = htons(ETH_P_ALL);
	if (bind(sock, (struct sockaddr *)&sll, sizeof(sll)) < 0) {
        printf("Failed to bind the interface %s\n", name);
		close(sock);
		return -1;
	}

	return sock;
}

static char *mapReqType(const int req){
    char *tmp = malloc(8);
    if (tmp == NULL)
        return NULL;

    switch(req){
        case 0x00:
            strncpy(tmp, "Query", 6);
            break;
        case 0x01:
            strncpy(tmp, "Answer", 7);
            break;
        default:
            strncpy(tmp, "Unknown", 8);
    };
    return tmp;
}
static char *mapClass(const int class){
    char *tmp = malloc(8);
    if (tmp == NULL)
        return NULL;

    memset(tmp, 0, 8);
    switch(class){
        case 1:
           strncpy(tmp, "IN", 3);
           break;
        case 2:
           strncpy(tmp, "CS", 3);
           break;
        case 3:
           strncpy(tmp, "CH", 3);
           break;
        case 4:
           strncpy(tmp, "HS", 3);
           break;
        default:
           strncpy(tmp, "Unknown", 8);
           break;
    }
    return tmp;
}
static char *mapType(const int type){
    char *tmp = malloc(8);
    if (tmp == NULL)
        return NULL;

    /*
     * type DNS defined in RFC:
     * https://datatracker.ietf.org/doc/html/rfc1035#section-3.2.2
     * https://datatracker.ietf.org/doc/html/rfc3596
     */
    switch(type){
        case 1:
           strncpy(tmp, "A", 2);
           break;
        case 2:
           strncpy(tmp, "NS", 3);
           break;
        case 3:
           strncpy(tmp, "MD", 3);
           break;
        case 4:
           strncpy(tmp, "MF", 3);
           break;
        case 5:
           strncpy(tmp, "CNAME", 6);
           break;
        case 6:
           strncpy(tmp, "SOA", 4);
           break;
        case 7:
           strncpy(tmp, "MB", 3);
           break;
        case 8:
           strncpy(tmp, "MG", 3);
           break;
        case 9:
           strncpy(tmp, "MR", 3);
           break;
        case 10:
           strncpy(tmp, "NULL", 5);
           break;
        case 11:
           strncpy(tmp, "WKS", 4);
           break;
        case 12:
           strncpy(tmp, "PTR", 4);
           break;
        case 13:
           strncpy(tmp, "HINFO", 6);
           break;
        case 14:
           strncpy(tmp, "MINFO", 6);
           break;
        case 15:
           strncpy(tmp, "MX", 3);
           break;
        case 16:
           strncpy(tmp, "TXT", 4);
           break;
        case 28:
            strncpy(tmp, "AAAA", 5);
            break;
        default:
           strncpy(tmp, "Unknown", 8);
           break;
    }
    return tmp;
}

static void print_query(struct event *s_event){
    char *req_type, *class, *type;
    req_type = mapReqType(s_event->req_type);
    printf("%s ", req_type);
    free(req_type);

    printf("%5s:%d\t\t", inet_ntoa(*(struct in_addr*)&s_event->client), s_event->dport);
    printf("%x\t", s_event->tid);

    class = mapClass(s_event->class);
    printf("%-5s", class);
    free(class);

    type = mapType(s_event->type);
    printf("%-5s", type);
    free(type);

    printf("%s", s_event->qname);
}
static void get_labels(unsigned char *buf, char *qname){
    int pos = 0;
    while (*buf++ != '\0') {
        if((*buf >= 'a' && *buf <= 'z') || (*buf >= 'A' && *buf <= 'Z'))
            *(qname + pos) = *buf;
        else if (*buf >= '0' && *buf <= '9')
            *(qname + pos) = *buf;
        else
            *(qname + pos) = '.';
        pos++;
    }
    qname[pos - 1] = '\0';
}
static void print_answer_hdr(struct event *s_event, int *pos, uint16_t *type, uint16_t *size, uint16_t *class, uint32_t *ttl){
    int p = *pos;
    uint16_t msg = s_event->buf[p++];
    msg |= s_event->buf[p++] << 8;
    // printf("%x\n", ntohs(msg));

    char *req_type;
    req_type = mapReqType(s_event->req_type);
    printf("%s ", req_type);
    free(req_type);

    printf("%5s:%5d\t", inet_ntoa(*(struct in_addr*)&s_event->client), s_event->dport);
    printf("%x\t", s_event->tid);

    *type = s_event->buf[p++];
    *type |= s_event->buf[p++] << 8;

    *class = s_event->buf[p++];
    *class |= s_event->buf[p++] << 8;

    *ttl = s_event->buf[p++];
    *ttl |= s_event->buf[p++] << 8;
    *ttl |= s_event->buf[p++] << 16;
    *ttl |= s_event->buf[p++] << 24;

    *size = s_event->buf[p++];
    *size |= s_event->buf[p++] << 8;

    *type = ntohs(*type);
    *class = ntohs(*class);
    *ttl = ntohl(*ttl);
    *size = ntohs(*size);

    *pos = p;
}
int handle_event(void *ctx, void *data, size_t data_sz){
    struct event *s_event = (struct event*)data;
    if (s_event->req_type == REQ_QUERY){
        print_query(s_event);
    }
    if (s_event->req_type == REQ_ANSWER){
        int pos = 0;
        for (int i = 0; i < s_event->numAns; i++){
            // print_query(s_event);
            uint16_t type, size, class;
            uint32_t ttl;
            char *s_type, *s_class;
            print_answer_hdr(s_event, &pos, &type, &size, &class, &ttl);

            s_class = mapClass(class);
            printf("%-5s", s_class);
            free(s_class);

            s_type = mapType(type);
            printf("%s\t", s_type);
            free(s_type);


            if (type == 1) { // -> A
                uint32_t ip = s_event->buf[pos] + (s_event->buf[pos+1] << 8) + (s_event->buf[pos+2] << 16) + (s_event->buf[pos+3] << 24);
                printf("%s %5d", inet_ntoa(*(struct in_addr*)&ip), ttl);
            }
            if (type == 5) { // -> CNAME
                char cname[size];
                get_labels(s_event->buf + pos, cname);
                printf("%s ", cname);
            }
            if (type == 28){ // -> AAAA
                int p = 0;
                for (int i = 0; i < size; i++){
                    if (i % 2 == 0)
                        printf("%x", s_event->buf[pos + p++]);
                    else{
                        if (i<size - 1)
                            printf("%x:", s_event->buf[pos + p++]);
                        else
                            printf("%x", s_event->buf[pos + p++]);
                    }
                }
            }
            pos += size;
            printf("\n");
        }
    }

    printf("\n");
    
    return 0;
}
int main(int argc, char *argv[]){
    const char *fileObj = "src/dns-trace.ebpf.o";
    struct bpf_object *obj;
    struct bpf_program *programSkb;
    struct ring_buffer *rb;
    int err;
    int fd_map_data;
    int sock;

    arguments = parse_args(argc, argv); // Parsing arguments

    if (arguments.interface == NULL){
        printf("You must specified the interface name");
        exit(-1);
    }
    printf("Listen to %s\n", arguments.interface);

    sock = open_raw_sock(arguments.interface);
    if (sock == -1){
        printf("Failed to listen to the interface %s\n", arguments.interface);
        exit(-1);
    }

    signal(SIGINT, signalHandler);

    /* Open and load our eBPF object */
    obj = bpf_object__open_file(fileObj, NULL);
    if (!obj){
        printf("Failed to open the file\n");
        return -1;
    }

    err = bpf_object__load(obj);
    if (err){
        printf("Failed to load object\n");
        return -1;
    }

    /* Retrieving fd of maps */
    fd_map_data = bpf_object__find_map_fd_by_name(obj, "m_data");
    if (!fd_map_data){
        printf("Failed to find the fd map data\n");
        bpf_object__close(obj);
        return -1;
    }

    /* Retrieving our programs */
    programSkb = bpf_object__find_program_by_name(obj, "detect_dns");

    if (!programSkb){
        printf("Failed to find program\n");
        bpf_object__close(obj);
        return -1;
    }

    bpf_program__attach(programSkb);
    int prog_fd = bpf_program__fd(programSkb);

    setsockopt(sock, SOL_SOCKET, SO_ATTACH_BPF, &prog_fd, sizeof(int));

    /* Start the ringbuffer */
    rb = ring_buffer__new(fd_map_data, handle_event, NULL, NULL);
    if (!rb){
        printf("Failed to create the ringbuf\n");
        bpf_object__close(obj);
        return -1;
    }

    while(running){
        err = ring_buffer__poll(rb, 100 /* timeout, ms */);
        if (err == -EINTR){
            printf("Failed to get the ringbuf\n");
            running = 0;
            break;
         }
    }

    ring_buffer__free(rb);
    bpf_object__close(obj);

    return 0;
}
