#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 *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 '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[] = {
        {"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.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("bind to %s: %s\n", name, strerror(errno));
        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;

    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;
        default:
           strncpy(tmp, "Unknown", 8);
           break;
    }
    return tmp;
}

int handle_event(void *ctx, void *data, size_t data_sz){
    struct event *s_event = (struct event*)data;
    char *req_type, *class, *type;
    printf("%s:%-10d", inet_ntoa(*(struct in_addr*)&s_event->client), s_event->dport);
    printf("%-5x", s_event->tid);

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

    printf("%-30s", s_event->qname);

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

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

    if (s_event->req_type == REQ_ANSWER){
        printf("%-15s", inet_ntoa(*(struct in_addr*)&s_event->ans));
        printf("%-5d", s_event->ttl);
    }


    /*printf("IP: %s\n", inet_ntoa(*(struct in_addr*)&s_event->client));
    printf("dport: %d\n", s_event->dport);
    printf("sport: %d\n", s_event->sport);
    printf("Transaction ID: %x\n", s_event->tid);
    printf("Request type: ");
    mapReqType(s_event->req_type);
    printf("qname: %s\n", s_event->qname);
    printf("Class: ");
    mapClass(s_event->class);
    printf("Type: ");
    mapType(s_event->type);

    if (s_event->req_type == REQ_ANSWER)
        printf("Data: %s\n", inet_ntoa(*(struct in_addr*)&s_event->ans));
    */
    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;

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

    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 sock = open_raw_sock("wlp0s20f3");
    int sock = open_raw_sock("enx98e743c667fc");
    printf("Socket: %d\n", sock);
    int prog_fd = bpf_program__fd(programSkb);
    printf("Program fd: %d\n", prog_fd);
    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;
}