// Lee el primer mensaje de un cliente TLS sin saber nada de TLS más allá del formato, y lo
// muestra. No tiene clave ni certificado: ve lo mismo que cualquier middlebox en el camino.
//
//   sni_peek <puerto> [conexiones]
#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/socket.h>
#include <sys/time.h>
#include <unistd.h>

#define MAX_HELLO 16384
#define RECORD_HEADER 5
#define HANDSHAKE_HEADER 4
#define TYPE_HANDSHAKE 22
#define HANDSHAKE_CLIENT_HELLO 1
#define EXT_SERVER_NAME 0
#define EXT_ALPN 16
#define EXT_SUPPORTED_VERSIONS 43
#define EXT_KEY_SHARE 51

// Un cursor sobre bytes recibidos: toda lectura revisa primero que quede espacio.
struct cursor {
    const uint8_t *data;
    size_t left;
};

static int read_u8(struct cursor *c, unsigned *value) {
    if (c->left < 1) return -1;
    *value = c->data[0];
    c->data += 1;
    c->left -= 1;
    return 0;
}

static int read_u16(struct cursor *c, unsigned *value) {
    if (c->left < 2) return -1;
    *value = ((unsigned)c->data[0] << 8) | c->data[1];
    c->data += 2;
    c->left -= 2;
    return 0;
}

static int skip(struct cursor *c, size_t n) {
    if (c->left < n) return -1;
    c->data += n;
    c->left -= n;
    return 0;
}

// Saca un pedazo de n bytes, con su propio límite.
static int slice(struct cursor *c, size_t n, struct cursor *part) {
    if (c->left < n) return -1;
    part->data = c->data;
    part->left = n;
    return skip(c, n);
}

// Lo que manda el cliente no es de fiar: nada de bytes de control hacia la terminal.
static void print_text(const uint8_t *data, size_t len) {
    for (size_t i = 0; i < len; i++) {
        if (data[i] >= 0x20 && data[i] < 0x7f) putchar(data[i]);
        else printf("\\x%02x", data[i]);
    }
}

static int show_server_name(struct cursor ext) {
    unsigned list_len, name_type, name_len;
    if (read_u16(&ext, &list_len) < 0 || read_u8(&ext, &name_type) < 0 ||
        read_u16(&ext, &name_len) < 0 || ext.left < name_len)
        return -1;
    printf("  SNI: ");
    print_text(ext.data, name_len);
    printf("\n");
    return 0;
}

static int show_alpn(struct cursor ext) {
    unsigned list_len, len;
    if (read_u16(&ext, &list_len) < 0) return -1;
    printf("  ALPN:");
    while (ext.left > 0) {
        if (read_u8(&ext, &len) < 0 || ext.left < len) return -1;
        printf(" ");
        print_text(ext.data, len);
        if (skip(&ext, len) < 0) return -1;
    }
    printf("\n");
    return 0;
}

static int show_versions(struct cursor ext) {
    unsigned list_len, version;
    if (read_u8(&ext, &list_len) < 0) return -1;
    printf("  versiones:");
    while (ext.left > 0) {
        if (read_u16(&ext, &version) < 0) return -1;
        if (version == 0x0304) printf(" TLS1.3");
        else if (version == 0x0303) printf(" TLS1.2");
        else printf(" 0x%04x", version);
    }
    printf("\n");
    return 0;
}

static int show_key_shares(struct cursor ext) {
    unsigned list_len, group, len;
    if (read_u16(&ext, &list_len) < 0) return -1;
    printf("  key_share:");
    while (ext.left > 0) {
        if (read_u16(&ext, &group) < 0 || read_u16(&ext, &len) < 0 || skip(&ext, len) < 0)
            return -1;
        printf(" grupo 0x%04x (%u bytes)", group, len);
    }
    printf("\n");
    return 0;
}

// Devuelve 0 si pudo leer el ClientHello completo, -1 si está mal formado o incompleto.
static int parse_client_hello(const uint8_t *buf, size_t len) {
    struct cursor msg = {buf, len};
    struct cursor extensions;
    unsigned type, major, minor, length, hs_high, hs_low, id_len, suites_len, comp_len, ext_len;

    if (read_u8(&msg, &type) < 0 || type != TYPE_HANDSHAKE) return -1;
    if (read_u8(&msg, &major) < 0 || read_u8(&msg, &minor) < 0) return -1;
    if (read_u16(&msg, &length) < 0 || msg.left < length) return -1;
    if (read_u8(&msg, &type) < 0 || type != HANDSHAKE_CLIENT_HELLO) return -1;
    if (read_u8(&msg, &hs_high) < 0 || read_u16(&msg, &hs_low) < 0) return -1;
    printf("  registro: tipo 22 (handshake), versión %u.%u, %u bytes\n", major, minor, length);
    if (skip(&msg, 2 + 32) < 0) return -1;  // legacy_version y random
    if (read_u8(&msg, &id_len) < 0 || skip(&msg, id_len) < 0) return -1;
    if (read_u16(&msg, &suites_len) < 0 || skip(&msg, suites_len) < 0) return -1;
    printf("  cipher suites ofrecidas: %u\n", suites_len / 2);
    if (read_u8(&msg, &comp_len) < 0 || skip(&msg, comp_len) < 0) return -1;
    if (read_u16(&msg, &ext_len) < 0 || slice(&msg, ext_len, &extensions) < 0) return -1;

    while (extensions.left > 0) {
        unsigned ext_type, body_len;
        struct cursor body;
        int status = 0;
        if (read_u16(&extensions, &ext_type) < 0 || read_u16(&extensions, &body_len) < 0 ||
            slice(&extensions, body_len, &body) < 0)
            return -1;
        if (ext_type == EXT_SERVER_NAME) status = show_server_name(body);
        else if (ext_type == EXT_ALPN) status = show_alpn(body);
        else if (ext_type == EXT_SUPPORTED_VERSIONS) status = show_versions(body);
        else if (ext_type == EXT_KEY_SHARE) status = show_key_shares(body);
        if (status < 0) return -1;
    }
    return 0;
}

// Lee un registro TLS entero: primero 5 bytes, y ahí dice cuánto falta.
static ssize_t read_record(int fd, uint8_t *buf, size_t size) {
    size_t have = 0;
    size_t want = RECORD_HEADER;
    while (have < want) {
        ssize_t got = recv(fd, buf + have, want - have, 0);
        if (got < 0) {
            if (errno == EINTR) continue;
            return -1;
        }
        if (got == 0) return (ssize_t)have;
        have += (size_t)got;
        if (have >= RECORD_HEADER && want == RECORD_HEADER) {
            size_t total = RECORD_HEADER + (((size_t)buf[3] << 8) | buf[4]);
            if (total > size) return (ssize_t)have;
            want = total;
        }
    }
    return (ssize_t)have;
}

int main(int argc, char **argv) {
    if (argc < 2) {
        fprintf(stderr, "uso: %s <puerto> [conexiones]\n", argv[0]);
        return 2;
    }
    int port = atoi(argv[1]);
    int count = argc > 2 ? atoi(argv[2]) : -1;
    if (port <= 0 || port > 65535) {
        fprintf(stderr, "puerto inválido: %s\n", argv[1]);
        return 2;
    }

    int listener = socket(AF_INET, SOCK_STREAM, 0);
    if (listener < 0) {
        perror("socket");
        return 1;
    }
    int one = 1;
    if (setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &one, sizeof one) < 0) {
        perror("setsockopt");
        return 1;
    }
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof addr);
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons((uint16_t)port);
    if (bind(listener, (struct sockaddr *)&addr, sizeof addr) < 0) {
        perror("bind");
        return 1;
    }
    if (listen(listener, 8) < 0) {
        perror("listen");
        return 1;
    }
    printf("escuchando en el puerto %d\n", port);
    fflush(stdout);

    static uint8_t buf[RECORD_HEADER + MAX_HELLO];
    for (int served = 0; count < 0 || served < count; served++) {
        struct sockaddr_in peer;
        socklen_t peer_len = sizeof peer;
        int fd = accept(listener, (struct sockaddr *)&peer, &peer_len);
        if (fd < 0) {
            if (errno == EINTR) continue;
            perror("accept");
            return 1;
        }
        // Un cliente que no manda nada no debe trabar el programa.
        struct timeval wait = {.tv_sec = 2, .tv_usec = 0};
        if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &wait, sizeof wait) < 0) perror("setsockopt");
        char text[INET_ADDRSTRLEN];
        if (inet_ntop(AF_INET, &peer.sin_addr, text, sizeof text) == NULL) {
            perror("inet_ntop");
            text[0] = '\0';
        }
        printf("conexión de %s\n", text);
        ssize_t got = read_record(fd, buf, sizeof buf);
        if (got < 0) perror("recv");
        else if (parse_client_hello(buf, (size_t)got) < 0) printf("  no es un ClientHello completo\n");
        fflush(stdout);
        if (close(fd) < 0) perror("close");
    }
    if (close(listener) < 0) perror("close");
    return 0;
}