/*
 * Copyright (c) 2023 - 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/rsa-aes.h"

#define MAX_PUB_KEY_SIZE 1000000

#define UDIV_UP(a, b) (((a) + (b) - 1) / (b))

int rsa_aes_send_public_key(struct nvnc_client* client)
{
	struct nvnc* server = client->server;

	if (!server->rsa_priv) {
		assert(!server->rsa_pub);

		nvnc_log(NVNC_LOG_WARNING, "An RSA key has not been set. A new key will be generated.");

		server->rsa_priv = crypto_rsa_priv_key_new();
		server->rsa_pub = crypto_rsa_pub_key_new();

		crypto_rsa_keygen(server->rsa_pub, server->rsa_priv);
	}
	assert(server->rsa_pub && server->rsa_priv);

	size_t key_len = crypto_rsa_pub_key_length(server->rsa_pub);
	size_t buf_len = sizeof(struct rfb_rsa_aes_pub_key_msg) + key_len * 2;

	char* buffer = calloc(1, buf_len);
	assert(buffer);
	struct rfb_rsa_aes_pub_key_msg* msg =
		(struct rfb_rsa_aes_pub_key_msg*)buffer;

	uint8_t* modulus = msg->modulus_and_exponent;
	uint8_t* exponent = msg->modulus_and_exponent + key_len;

	msg->length = htonl(key_len * 8);
	crypto_rsa_pub_key_modulus(server->rsa_pub, modulus, key_len);
	crypto_rsa_pub_key_exponent(server->rsa_pub, exponent, key_len);

	stream_send(client->net_stream, rcbuf_new(buffer, buf_len), NULL, NULL);
	return 0;
}

static int rsa_aes_send_challenge(struct nvnc_client* client,
		struct crypto_rsa_pub_key* pub)
{
	crypto_random(client->rsa.challenge, client->rsa.challenge_len);

	struct rfb_rsa_aes_challenge_msg* msg;
	size_t key_len = crypto_rsa_pub_key_length(client->rsa.pub);
	size_t msg_size = sizeof(*msg) + key_len;
	msg = calloc(1, msg_size);
	nvnc_assert(msg, "OOM");

	crypto_rsa_encrypt(pub, msg->challenge, key_len,
			client->rsa.challenge, client->rsa.challenge_len);
	msg->length = htons(key_len);

	stream_send(client->net_stream, rcbuf_new(msg, msg_size), NULL, NULL);
	return 0;
}

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

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

	uint32_t bit_length = ntohl(msg->length);
	size_t byte_length = UDIV_UP(bit_length, 8);

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

	if (byte_length > MAX_PUB_KEY_SIZE) {
		nvnc_log(NVNC_LOG_ERROR, "Client sent a ridiculously large public key. This can't be right.");
		nvnc_client_close(client);
		return -1;
	}

	const uint8_t* modulus = msg->modulus_and_exponent;
	const uint8_t* exponent = msg->modulus_and_exponent + byte_length;

	client->rsa.pub =
		crypto_rsa_pub_key_import(modulus, exponent, byte_length);
	if (!client->rsa.pub) {
		nvnc_log(NVNC_LOG_ERROR, "Failed to import public key from client");
		nvnc_client_close(client);
		return -1;
	}

	update_min_rtt(client);

	client->state = VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_CHALLENGE;
	rsa_aes_send_challenge(client, client->rsa.pub);

	return sizeof(*msg) + byte_length * 2;
}

static size_t client_rsa_aes_hash_len(const struct nvnc_client* client)
{
	switch (client->rsa.hash_type) {
	case CRYPTO_HASH_SHA1: return 20;
	case CRYPTO_HASH_SHA256: return 32;
	default:;
	}
	abort();
	return 0;
}

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

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

	uint16_t length = ntohs(msg->length);
	if (client->buffer_len - client->buffer_index < sizeof(*msg) + length)
		return 0;

	struct nvnc* server = client->server;

	uint8_t client_random[32] = {};
	ssize_t len = crypto_rsa_decrypt(server->rsa_priv, client_random,
			client->rsa.challenge_len, msg->challenge, length);
	if (len < 0) {
		nvnc_log(NVNC_LOG_ERROR, "Failed to decrypt client's challenge");
		nvnc_client_close(client);
		return -1;
	}

	// ClientSessionKey = the first 16 bytes of SHA1(ServerRandom || ClientRandom)
	uint8_t client_session_key[32];
	crypto_hash_many(client_session_key, client_rsa_aes_hash_len(client),
			client->rsa.hash_type, (const struct crypto_data_entry[]) {
		{ client->rsa.challenge, client->rsa.challenge_len },
		{ client_random, client->rsa.challenge_len },
		{}
	});

	// ServerSessionKey = the first 16 bytes of SHA1(ClientRandom || ServerRandom)
	uint8_t server_session_key[32];
	crypto_hash_many(server_session_key, client_rsa_aes_hash_len(client),
			client->rsa.hash_type, (const struct crypto_data_entry[]) {
		{ client_random, client->rsa.challenge_len },
		{ client->rsa.challenge, client->rsa.challenge_len },
		{}
	});

	stream_upgrade_to_rsa_eas(client->net_stream, client->rsa.cipher_type,
			server_session_key, client_session_key);

	size_t server_key_len = crypto_rsa_pub_key_length(server->rsa_pub);
	uint8_t* server_modulus = malloc(server_key_len * 2);
	uint8_t* server_exponent = server_modulus + server_key_len;

	crypto_rsa_pub_key_modulus(server->rsa_pub, server_modulus,
			server_key_len);
	crypto_rsa_pub_key_exponent(server->rsa_pub, server_exponent,
			server_key_len);

	size_t client_key_len = crypto_rsa_pub_key_length(client->rsa.pub);
	uint8_t* client_modulus = malloc(client_key_len * 2);
	uint8_t* client_exponent = client_modulus + client_key_len;

	crypto_rsa_pub_key_modulus(client->rsa.pub, client_modulus,
			client_key_len);
	crypto_rsa_pub_key_exponent(client->rsa.pub, client_exponent,
			client_key_len);

	uint32_t server_key_len_be = htonl(server_key_len * 8);
	uint32_t client_key_len_be = htonl(client_key_len * 8);

	uint8_t server_hash[32] = {};
	crypto_hash_many(server_hash, client_rsa_aes_hash_len(client),
			client->rsa.hash_type, (const struct crypto_data_entry[]) {
		{ (uint8_t*)&server_key_len_be, 4 },
		{ server_modulus, server_key_len },
		{ server_exponent, server_key_len },
		{ (uint8_t*)&client_key_len_be, 4 },
		{ client_modulus, client_key_len },
		{ client_exponent, client_key_len },
		{}
	});

	free(server_modulus);
	free(client_modulus);

	update_min_rtt(client);

	stream_write(client->net_stream, server_hash,
			client_rsa_aes_hash_len(client));

	client->state = VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_CLIENT_HASH;

	return sizeof(*msg) + length;
}

static int on_rsa_aes_client_hash(struct nvnc_client* client)
{
	const char* msg = (void*)(client->msg_buffer + client->buffer_index);

	if (client->buffer_len - client->buffer_index < client_rsa_aes_hash_len(client))
		return 0;

	struct nvnc* server = client->server;

	size_t server_key_len = crypto_rsa_pub_key_length(server->rsa_pub);
	uint8_t* server_modulus = malloc(server_key_len * 2);
	uint8_t* server_exponent = server_modulus + server_key_len;
	crypto_rsa_pub_key_modulus(server->rsa_pub, server_modulus,
			server_key_len);
	crypto_rsa_pub_key_exponent(server->rsa_pub, server_exponent,
			server_key_len);

	size_t client_key_len = crypto_rsa_pub_key_length(client->rsa.pub);
	uint8_t* client_modulus = malloc(client_key_len * 2);
	uint8_t* client_exponent = client_modulus + client_key_len;

	crypto_rsa_pub_key_modulus(client->rsa.pub, client_modulus,
			client_key_len);
	crypto_rsa_pub_key_exponent(client->rsa.pub, client_exponent,
			client_key_len);

	uint32_t server_key_len_be = htonl(server_key_len * 8);
	uint32_t client_key_len_be = htonl(client_key_len * 8);

	uint8_t client_hash[32] = {};
	crypto_hash_many(client_hash, client_rsa_aes_hash_len(client),
			client->rsa.hash_type, (const struct crypto_data_entry[]) {
		{ (uint8_t*)&client_key_len_be, 4 },
		{ client_modulus, client_key_len },
		{ client_exponent, client_key_len },
		{ (uint8_t*)&server_key_len_be, 4 },
		{ server_modulus, server_key_len },
		{ server_exponent, server_key_len },
		{}
	});

	free(client_modulus);
	free(server_modulus);

	if (memcmp(msg, client_hash, client_rsa_aes_hash_len(client)) != 0) {
		nvnc_log(NVNC_LOG_INFO, "Client hash mismatch");
		nvnc_client_close(client);
		return -1;
	}

	update_min_rtt(client);

	// TODO: Read this from config
	uint8_t subtype = RFB_RSA_AES_CRED_SUBTYPE_USER_AND_PASS;
	stream_write(client->net_stream, &subtype, 1);

	client->state = VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_CREDENTIALS;
	return client_rsa_aes_hash_len(client);
}

static int on_rsa_aes_credentials(struct nvnc_client* client)
{
	const uint8_t* msg = (void*)(client->msg_buffer + client->buffer_index);

	if (client->buffer_len - client->buffer_index < 2)
		return 0;

	size_t username_len = msg[0];
	if (client->buffer_len - client->buffer_index < 2 + username_len)
		return 0;

	size_t password_len = msg[1 + username_len];
	if (client->buffer_len - client->buffer_index < 2 + username_len +
			password_len)
		return 0;

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

	memcpy(creds->username, (const char*)(msg + 1), username_len);
	creds->username[username_len] = '\0';
	memcpy(creds->password, (const char*)(msg + 2 + username_len),
			password_len);
	creds->password[password_len] = '\0';

	update_min_rtt(client);

	security_handshake_authenticate(client, creds);

	return 2 + username_len + password_len;
}

int rsa_aes_handle_message(struct nvnc_client* client)
{
	switch (client->state) {
	case VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_PUBLIC_KEY:
		return on_rsa_aes_public_key(client);
	case VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_CHALLENGE:
		return on_rsa_aes_challenge(client);
	case VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_CLIENT_HASH:
		return on_rsa_aes_client_hash(client);
	case VNC_CLIENT_STATE_WAITING_FOR_RSA_AES_CREDENTIALS:
		return on_rsa_aes_credentials(client);
	default:;
	}
	nvnc_log(NVNC_LOG_PANIC, "Unhandled client state: %d", client->state);
	return 0;
}
