/*
 * WebSocket SSH Library - WebSocket functions implementation
 *
 * Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include "websocket.h"
#include "wssshlib.h"
#include <openssl/err.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

int websocket_handshake(SSL *ssl, const char *host, int port, const char *path) {
    char request[1024];
    char response[BUFFER_SIZE];
    int bytes_read;

    // Send WebSocket handshake
    snprintf(request, sizeof(request),
             "GET %s HTTP/1.1\r\n"
             "Host: %s:%d\r\n"
             "Upgrade: websocket\r\n"
             "Connection: upgrade\r\n"
             "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
             "Sec-WebSocket-Version: 13\r\n"
             "\r\n",
             path, host, port);

    if (SSL_write(ssl, request, strlen(request)) <= 0) {
        ERR_print_errors_fp(stderr);
        fprintf(stderr, "WebSocket handshake send failed\n");
        return 0;
    }

    // Read response
    bytes_read = SSL_read(ssl, response, sizeof(response) - 1);
    if (bytes_read <= 0) {
        ERR_print_errors_fp(stderr);
        fprintf(stderr, "WebSocket handshake recv failed\n");
        return 0;
    }

    response[bytes_read] = '\0';

    // Check for successful handshake
    if (strstr(response, "101 Switching Protocols") == NULL) {
        fprintf(stderr, "WebSocket handshake failed\n");
        return 0;
    }

    return 1;
}

int send_json_message(SSL *ssl, const char *type, const char *client_id, const char *request_id) {
    char message[1024];

    if (request_id) {
        snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"client_id\":\"%s\",\"request_id\":\"%s\"}",
                 type, client_id, request_id);
    } else {
        snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"client_id\":\"%s\"}",
                 type, client_id);
    }

    // Send as WebSocket frame
    return send_websocket_frame(ssl, message);
}

int send_websocket_frame(SSL *ssl, const char *data) {
    char frame[BUFFER_SIZE];
    frame[0] = 0x81; // FIN + text opcode
    int msg_len = strlen(data);
    int header_len = 2;

    if (msg_len <= 125) {
        frame[1] = 0x80 | msg_len; // MASK + length
    } else if (msg_len <= 65535) {
        frame[1] = 0x80 | 126; // MASK + extended length
        frame[2] = (msg_len >> 8) & 0xFF;
        frame[3] = msg_len & 0xFF;
        header_len = 4;
    } else {
        frame[1] = 0x80 | 127; // MASK + extended length
        frame[2] = 0;
        frame[3] = 0;
        frame[4] = 0;
        frame[5] = 0;
        frame[6] = (msg_len >> 24) & 0xFF;
        frame[7] = (msg_len >> 16) & 0xFF;
        frame[8] = (msg_len >> 8) & 0xFF;
        frame[9] = msg_len & 0xFF;
        header_len = 10;
    }

    // Add mask key
    char mask_key[4];
    for (int i = 0; i < 4; i++) {
        mask_key[i] = rand() % 256;
        frame[header_len + i] = mask_key[i];
    }
    header_len += 4;

    // Mask payload
    for (int i = 0; i < msg_len; i++) {
        frame[header_len + i] = data[i] ^ mask_key[i % 4];
    }

    int frame_len = header_len + msg_len;
    return SSL_write(ssl, frame, frame_len) > 0;
}

int send_pong_frame(SSL *ssl, const char *ping_payload, int payload_len) {
    char frame[BUFFER_SIZE];
    frame[0] = 0x8A; // FIN + pong opcode
    int header_len = 2;

    if (payload_len <= 125) {
        frame[1] = 0x80 | payload_len; // MASK + length
    } else if (payload_len <= 65535) {
        frame[1] = 0x80 | 126; // MASK + extended length
        frame[2] = (payload_len >> 8) & 0xFF;
        frame[3] = payload_len & 0xFF;
        header_len = 4;
    } else {
        frame[1] = 0x80 | 127; // MASK + extended length
        frame[2] = 0;
        frame[3] = 0;
        frame[4] = 0;
        frame[5] = 0;
        frame[6] = (payload_len >> 24) & 0xFF;
        frame[7] = (payload_len >> 16) & 0xFF;
        frame[8] = (payload_len >> 8) & 0xFF;
        frame[9] = payload_len & 0xFF;
        header_len = 10;
    }

    // Add mask key
    char mask_key[4];
    for (int i = 0; i < 4; i++) {
        mask_key[i] = rand() % 256;
        frame[header_len + i] = mask_key[i];
    }
    header_len += 4;

    // Mask payload
    for (int i = 0; i < payload_len; i++) {
        frame[header_len + i] = ping_payload[i] ^ mask_key[i % 4];
    }

    int frame_len = header_len + payload_len;
    return SSL_write(ssl, frame, frame_len) > 0;
}

int parse_websocket_frame(const char *buffer, int bytes_read, char **payload, int *payload_len) {
    if (bytes_read < 2) {
        return 0; // Not enough data for a frame
    }

    unsigned char frame_type = buffer[0] & 0x8F;
    if (frame_type != 0x81 && frame_type != 0x82 && frame_type != 0x88 && frame_type != 0x89 && frame_type != 0x8A) {
        return 0; // Not a supported frame type (text, binary, close, ping, pong)
    }

    int masked = buffer[1] & 0x80;
    int len_indicator = buffer[1] & 0x7F;
    int header_len = 2;


    if (len_indicator <= 125) {
        *payload_len = len_indicator;
    } else if (len_indicator == 126) {
        if (bytes_read < 4) return 0;
        *payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
        header_len = 4;
    } else if (len_indicator == 127) {
        if (bytes_read < 10) return 0;
        unsigned long long full_len = ((unsigned long long)(unsigned char)buffer[2] << 56) |
                                      ((unsigned long long)(unsigned char)buffer[3] << 48) |
                                      ((unsigned long long)(unsigned char)buffer[4] << 40) |
                                      ((unsigned long long)(unsigned char)buffer[5] << 32) |
                                      ((unsigned long long)(unsigned char)buffer[6] << 24) |
                                      ((unsigned long long)(unsigned char)buffer[7] << 16) |
                                      ((unsigned long long)(unsigned char)buffer[8] << 8) |
                                      (unsigned char)buffer[9];
        if (full_len > INT_MAX) {
            return 0; // Payload too large
        }
        *payload_len = (int)full_len;
        header_len = 10;
    } else {
        // Invalid length indicator
        return 0;
    }

    if (masked) {
        header_len += 4;
    }

    if (bytes_read < header_len + *payload_len) {
        return 0; // Incomplete frame
    }

    *payload = (char *)buffer + header_len;
    if (masked) {
        char *mask_key = (char *)buffer + header_len - 4;
        for (int i = 0; i < *payload_len; i++) {
            (*payload)[i] ^= mask_key[i % 4];
        }
    }

    return 1;
}