// dns_ask: arma a mano una pregunta de DNS, la manda por UDP y muestra el encabezado y las
// respuestas de tipo A. Sirve para ver el formato del RFC 1035 sin ninguna biblioteca.
//   uso: dns_ask <IPv4 del servidor> <nombre> [--norecurse]
#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/random.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>

enum { HEADER_SIZE = 12, MAX_MESSAGE = 512, MAX_JUMPS = 32, TYPE_A = 1, CLASS_IN = 1 };

static uint16_t read_u16(const uint8_t *from) {
    return (uint16_t)((from[0] << 8) | from[1]);
}

static void write_u16(uint8_t *to, uint16_t value) {
    to[0] = (uint8_t)(value >> 8);
    to[1] = (uint8_t)(value & 0xFF);
}

// Escribe "www.uni.test" como 3www3uni4test0. Devuelve los bytes usados, o -1 si no cabe.
static int encode_name(const char *name, uint8_t *out, size_t room) {
    size_t used = 0;
    while (*name != '\0') {
        const char *dot = strchr(name, '.');
        size_t size = dot != NULL ? (size_t)(dot - name) : strlen(name);
        if (size == 0 || size > 63 || used + size + 2 > room) {
            return -1;
        }
        out[used++] = (uint8_t)size;
        memcpy(out + used, name, size);
        used += size;
        name += size;
        if (*name == '.') {
            name++;
        }
    }
    if (used + 1 > room) {
        return -1;
    }
    out[used++] = 0;
    return (int)used;
}

// Se salta un nombre y devuelve dónde sigue. Un puntero ocupa 2 bytes y termina el nombre.
static int skip_name(const uint8_t *message, int length, int offset) {
    while (offset < length) {
        uint8_t size = message[offset];
        if ((size & 0xC0) == 0xC0) {
            return offset + 2 <= length ? offset + 2 : -1;
        }
        if ((size & 0xC0) != 0) {
            return -1;
        }
        if (size == 0) {
            return offset + 1;
        }
        offset += 1 + size;
    }
    return -1;
}

// Imprime un nombre, siguiendo punteros. Solo se aceptan punteros hacia atrás.
static int print_name(const uint8_t *message, int length, int offset) {
    int jumps = 0;
    while (offset < length) {
        uint8_t size = message[offset];
        if ((size & 0xC0) == 0xC0) {
            if (offset + 2 > length || ++jumps > MAX_JUMPS) {
                return -1;
            }
            int target = ((size & 0x3F) << 8) | message[offset + 1];
            if (target >= offset) {
                return -1;
            }
            offset = target;
        } else if ((size & 0xC0) != 0) {
            return -1;
        } else if (size == 0) {
            return 0;
        } else {
            if (offset + 1 + size > length) {
                return -1;
            }
            printf("%.*s.", size, (const char *)message + offset + 1);
            offset += 1 + size;
        }
    }
    return -1;
}

static int print_reply(const uint8_t *reply, int length, uint16_t id) {
    if (length < HEADER_SIZE) {
        fprintf(stderr, "dns_ask: respuesta más corta que el encabezado\n");
        return -1;
    }
    if (read_u16(reply) != id) {
        fprintf(stderr, "dns_ask: el ID no coincide; se descarta\n");
        return -1;
    }
    uint16_t flags = read_u16(reply + 2);
    uint16_t questions = read_u16(reply + 4);
    uint16_t answers = read_u16(reply + 6);
    printf("id=%u qr=%d aa=%d tc=%d rd=%d ra=%d rcode=%d\n", id, (flags >> 15) & 1, (flags >> 10) & 1,
           (flags >> 9) & 1, (flags >> 8) & 1, (flags >> 7) & 1, flags & 0xF);
    printf("preguntas=%u respuestas=%u autoridad=%u adicionales=%u\n", questions, answers,
           read_u16(reply + 8), read_u16(reply + 10));

    int offset = HEADER_SIZE;
    for (int i = 0; i < questions; i++) {
        offset = skip_name(reply, length, offset);
        if (offset < 0 || offset + 4 > length) {
            fprintf(stderr, "dns_ask: pregunta mal formada\n");
            return -1;
        }
        offset += 4;
    }
    for (int i = 0; i < answers; i++) {
        int name_at = offset;
        offset = skip_name(reply, length, offset);
        if (offset < 0 || offset + 10 > length) {
            fprintf(stderr, "dns_ask: respuesta mal formada\n");
            return -1;
        }
        uint16_t type = read_u16(reply + offset);
        uint32_t ttl = ((uint32_t)read_u16(reply + offset + 4) << 16) | read_u16(reply + offset + 6);
        uint16_t size = read_u16(reply + offset + 8);
        offset += 10;
        if (offset + size > length) {
            fprintf(stderr, "dns_ask: datos cortados\n");
            return -1;
        }
        if (type == TYPE_A && size == 4) {
            char text[INET_ADDRSTRLEN];
            if (inet_ntop(AF_INET, reply + offset, text, sizeof text) == NULL) {
                perror("inet_ntop");
                return -1;
            }
            if (print_name(reply, length, name_at) < 0) {
                fprintf(stderr, "dns_ask: nombre mal formado\n");
                return -1;
            }
            printf(" ttl=%u A %s\n", ttl, text);
        }
        offset += size;
    }
    return 0;
}

int main(int argc, char *argv[]) {
    if (argc < 3 || argc > 4 || (argc == 4 && strcmp(argv[3], "--norecurse") != 0)) {
        fprintf(stderr, "uso: %s <IPv4 del servidor> <nombre> [--norecurse]\n", argv[0]);
        return 2;
    }
    struct sockaddr_in server = {.sin_family = AF_INET, .sin_port = htons(53)};
    if (inet_pton(AF_INET, argv[1], &server.sin_addr) != 1) {
        fprintf(stderr, "dns_ask: dirección inválida: %s\n", argv[1]);
        return 2;
    }

    uint16_t id;
    if (getrandom(&id, sizeof id, 0) != (ssize_t)sizeof id) {
        perror("getrandom");
        return 1;
    }

    // Nunca se copia una struct a la red: cada campo se escribe en orden de red.
    uint8_t query[MAX_MESSAGE] = {0};
    write_u16(query, id);
    write_u16(query + 2, argc == 4 ? 0x0000 : 0x0100);  // RD pide recursión
    write_u16(query + 4, 1);                            // una pregunta
    int name_size = encode_name(argv[2], query + HEADER_SIZE, sizeof query - HEADER_SIZE - 4);
    if (name_size < 0) {
        fprintf(stderr, "dns_ask: nombre inválido: %s\n", argv[2]);
        return 2;
    }
    write_u16(query + HEADER_SIZE + name_size, TYPE_A);
    write_u16(query + HEADER_SIZE + name_size + 2, CLASS_IN);
    size_t query_size = (size_t)(HEADER_SIZE + name_size + 4);

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket");
        return 1;
    }
    struct timeval timeout = {.tv_sec = 2};
    if (setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof timeout) < 0) {
        perror("setsockopt");
        close(sock);
        return 1;
    }
    // Con connect, el kernel solo nos entrega datagramas que vengan de ese servidor.
    if (connect(sock, (struct sockaddr *)&server, sizeof server) < 0) {
        perror("connect");
        close(sock);
        return 1;
    }
    if (send(sock, query, query_size, 0) != (ssize_t)query_size) {
        perror("send");
        close(sock);
        return 1;
    }
    uint8_t reply[MAX_MESSAGE];
    ssize_t got = recv(sock, reply, sizeof reply, 0);
    if (got < 0) {
        fprintf(stderr, "dns_ask: sin respuesta: %s\n", strerror(errno));
        close(sock);
        return 1;
    }
    if (close(sock) < 0) {
        perror("close");
        return 1;
    }
    return print_reply(reply, (int)got, id) == 0 ? 0 : 1;
}