/*
 * Copyright (c) 2019 - 2024 Andri Yngvason
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
 * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
 * LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
 * OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 * PERFORMANCE OF THIS SOFTWARE.
 */

#include "client.h"
#include "server.h"
#include "stream/stream.h"
#include "auth/auth.h"
#include "auth/vencrypt.h"

#include <sys/param.h>

static int send_byte(struct nvnc_client* client, uint8_t value)
{
	return stream_write(client->net_stream, &value, 1);
}

int vencrypt_send_version(struct nvnc_client* client)
{
	struct rfb_vencrypt_version_msg msg = {
		.major = 0,
		.minor = 2,
	};

	return stream_write(client->net_stream, &msg, sizeof(msg));
}

static int on_vencrypt_version_message(struct nvnc_client* client)
{
	struct rfb_vencrypt_version_msg* msg =
		(struct rfb_vencrypt_version_msg*)&client->msg_buffer[client->buffer_index];

	if (client->buffer_len - client->buffer_index < sizeof(*msg))
		return 0;

	if (msg->major != 0 || msg->minor != 2) {
		security_handshake_failed(client, "Unsupported VeNCrypt version");
		return -1;
	}

	send_byte(client, 0);

	struct rfb_vencrypt_subtypes_msg result = { .n = 1, };
	result.types[0] = htonl(RFB_VENCRYPT_X509_PLAIN);

	update_min_rtt(client);

	stream_write(client->net_stream, &result, sizeof(result));

	client->state = VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_SUBTYPE;

	return sizeof(*msg);
}

static int on_vencrypt_subtype_message(struct nvnc_client* client)
{
	uint32_t* msg = (uint32_t*)&client->msg_buffer[client->buffer_index];

	if (client->buffer_len - client->buffer_index < sizeof(*msg))
		return 0;

	if (client->buffer_len - client->buffer_index != sizeof(*msg)) {
		nvnc_log(NVNC_LOG_WARNING, "Client tried to send additional data prior to TLS upgrade");
		send_byte(client, 0);
		nvnc_client_close(client);
		return -1;
	}

	enum rfb_vencrypt_subtype subtype = ntohl(*msg);

	if (subtype != RFB_VENCRYPT_X509_PLAIN) {
		send_byte(client, 0);
		nvnc_client_close(client);
		return -1;
	}

	update_min_rtt(client);

	send_byte(client, 1);

	if (stream_upgrade_to_tls(client->net_stream, client->server->tls_creds) < 0) {
		nvnc_client_close(client);
		return -1;
	}

	client->state = VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_PLAIN_AUTH;

	return sizeof(*msg);
}

static int on_vencrypt_plain_auth_message(struct nvnc_client* client)
{
	struct rfb_vencrypt_plain_auth_msg* msg =
	        (void*)(client->msg_buffer + client->buffer_index);

	if (client->buffer_len - client->buffer_index < sizeof(*msg))
		return 0;

	uint32_t ulen = ntohl(msg->username_len);
	uint32_t plen = ntohl(msg->password_len);

	if (ulen > NVNC_AUTH_USERNAME_MAX || plen > NVNC_AUTH_PASSWORD_MAX) {
		nvnc_log(NVNC_LOG_ERROR, "Client sent too long username/password");
		nvnc_client_close(client);
		return -1;
	}

	if (client->buffer_len - client->buffer_index < sizeof(*msg) + ulen + plen)
		return 0;

	struct nvnc_auth_creds* creds =
		nvnc_auth_creds_new(NVNC_AUTH_CREDS_PLAIN);
	nvnc_assert(creds, "Out of memory");

	memcpy(creds->username, msg->text,
			MIN(ulen, sizeof(creds->username) - 1));
	memcpy(creds->password, msg->text + ulen,
			MIN(plen, sizeof(creds->password) - 1));

	creds->username[MIN(ulen, sizeof(creds->username) - 1)] = '\0';
	creds->password[MIN(plen, sizeof(creds->password) - 1)] = '\0';

	update_min_rtt(client);

	security_handshake_authenticate(client, creds);

	return sizeof(*msg) + ulen + plen;
}

int vencrypt_handle_message(struct nvnc_client* client)
{
	switch (client->state) {
	case VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_VERSION:
		return on_vencrypt_version_message(client);
	case VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_SUBTYPE:
		return on_vencrypt_subtype_message(client);
	case VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_PLAIN_AUTH:
		return on_vencrypt_plain_auth_message(client);
	default:;
	}
	nvnc_log(NVNC_LOG_PANIC, "Unhandled client state: %d", client->state);
	return 0;
}
