#include <stdio.h>
#include <unistd.h>
#include <sched.h>
#include <sys/mount.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <stdlib.h> 
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <net/if.h>
#include <netinet/in.h>
#include <string.h>
#include <dirent.h>
#include <sys/stat.h>
#include <pthread.h>
#include <resolv.h>
#include <netdb.h>
#include "test.h"

#define TEST(c, ...) \
	( (c) || (t_error(#c " failed: " __VA_ARGS__),0) )

// Example IP
#define EXPECTED_IP_ADDRESS "192.168.17.1"

// Simple answer resource record
#define EXAMPLE_ANSWER_RR "\xc0\x0c" \
                          "\x00\x01" \
                          "\x00\x01" \
                          "\x00\x00\x02\x58" \
                          "\x00\x04" \
                          "\xc0\xa8\x11\x01" // IPv4 192.168.17.1

// Wait until serving thread is ready to receive
pthread_barrier_t sync_barrier;

static size_t construct_response(uint16_t id, unsigned char *question, unsigned char *response)
{
    HEADER dns_header;
    const size_t dns_header_size = sizeof(dns_header);
    const size_t expected_question_size = 17U;
    memset(&dns_header, 0, dns_header_size);

    dns_header.id = id;
    dns_header.qr = 0x01U;
    dns_header.rd = 0x01U;
    dns_header.ra = 0x01U;
    dns_header.qdcount = 0x0100U; // 1 question
    dns_header.ancount = 0x0100U; // 1 answer

    memcpy(response, &dns_header, sizeof(dns_header));
    memcpy(&response[dns_header_size], &question[dns_header_size], expected_question_size);
    unsigned char answer_buffer[] = EXAMPLE_ANSWER_RR;

    memcpy(&response[dns_header_size + expected_question_size], &answer_buffer[0], sizeof(answer_buffer) - 1);

    return dns_header_size + expected_question_size + sizeof(answer_buffer) - 1; //ignore null terminator
}

static int bind_to_socket(int s)
{
    struct sockaddr_in dns_server;

    memset(&dns_server, 0, sizeof(dns_server));
    dns_server.sin_addr.s_addr = inet_addr("127.0.0.1");
    dns_server.sin_family = AF_INET;
    dns_server.sin_port = htons(53);

    return bind(s, (struct sockaddr*)&dns_server, sizeof(dns_server));
}

static void set_environment(void)
{
    FILE *ft = fopen("/etc/resolv.conf", "w");
    if (ft == NULL) t_error("unable to open namespaced resolv.conf\n");
    fprintf(ft, "nameserver 127.0.0.1");
    fclose(ft);
    
    ft = fopen("/etc/hosts", "w");
    if (ft == NULL) t_error("unable to open namespaced resolv.conf\n");
    fprintf(ft, "127.0.0.1 localhost");
    fclose(ft);
}

void *dns_server(void *arguments)
{
    int s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    int status = bind_to_socket(s);
    
    struct sockaddr_in from = {0};
    socklen_t from_length = sizeof(from);
    unsigned char packet_buffer[NS_PACKETSZ];
    pthread_barrier_wait(&sync_barrier);
    int packet_size = recvfrom(s, packet_buffer, NS_PACKETSZ, 0,
                              (struct sockaddr*)&from, &from_length);

    unsigned char response_buffer[NS_PACKETSZ] = {0};
    const uint16_t response_id = (packet_buffer[1] << 8) | packet_buffer[0];
    size_t response_size = construct_response(response_id, packet_buffer, response_buffer);
    status = sendto(s, response_buffer, response_size, 0, (struct sockaddr*)&from, from_length);
    
    return 0; 
}

int main(int argc, char **argv)
{
    pthread_barrier_init(&sync_barrier, NULL, 2);
    set_environment();

    pthread_t dns_thread;
    int status = pthread_create(&dns_thread, 0, dns_server, 0);
    pthread_barrier_wait(&sync_barrier);
    
    unsigned char res_buffer[NS_PACKETSZ] = {0};
    char *query = "example.com";
    int length = res_query(query, C_ANY, T_PTR, res_buffer, sizeof(res_buffer));
   
    // Simple test, expect the resulting IP to be in the last four bytes of 
    // the buffer
    struct in_addr returned_address = {
        .s_addr = res_buffer[length - 1] << 24
                | res_buffer[length - 2] << 16
                | res_buffer[length - 3] << 8
                | res_buffer[length - 4]
    };
    
    TEST(!strcmp(inet_ntoa(returned_address), EXPECTED_IP_ADDRESS), 
         "Expected ip address %s, got %s\n", EXPECTED_IP_ADDRESS, inet_ntoa(returned_address));
    
    void *thread_return;
    pthread_join(dns_thread, &thread_return);

    return t_status;
}