#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 void mapClass(const int class){
    switch(class){
        case 1:
           printf("IN\n");
           break;
        case 2:
           printf("CS\n");
           break;
        case 3:
           printf("CH\n");
           break;
        case 4:
           printf("HS\n");
           break;
        default:
           printf("Unknown\n");
           break;
    }
}
static void mapType(const int type){
    switch(type){
        case 1:
           printf("A");
           break;
        case 2:
           printf("NS");
           break;
        case 3:
           printf("MD");
           break;
        case 4:
           printf("MF");
           break;
        case 5:
           printf("CNAME");
           break;
        case 6:
           printf("SOA");
           break;
        case 7:
           printf("MB");
           break;
        case 8:
           printf("MG");
           break;
        case 9:
           printf("MR");
           break;
        case 10:
           printf("NULL");
           break;
        case 11:
           printf("WKS");
           break;
        case 12:
           printf("PTR");
           break;
        case 13:
           printf("HINFO");
           break;
        case 14:
           printf("MINFO");
           break;
        case 15:
           printf("MX");
           break;
        case 16:
           printf("TXT");
           break;
        default:
           printf("Unknown\n");
           break;
    }
    printf("\n");
}

int handle_event(void *ctx, void *data, size_t data_sz){
    struct event *s_event = (struct event*)data;
    printf("IP: %s\n", inet_ntoa(*(struct in_addr*)&s_event->saddr));
    printf("dport: %d\n", s_event->dport);
    printf("sport: %d\n", s_event->sport);
    printf("qname: %s\n", s_event->qname);
    printf("Class: ");
    mapClass(s_event->class);
    printf("Type: ");
    mapType(s_event->type);
    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;
}