/*
 * 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/apple-dh.h"

#define APPLE_DH_SERVER_KEY_LENGTH 256

int apple_dh_send_public_key(struct nvnc_client* client)
{
	client->apple_dh_secret = crypto_keygen();
	assert(client->apple_dh_secret);

	struct crypto_key* pub =
		crypto_derive_public_key(client->apple_dh_secret);
	assert(pub);

	uint8_t mod[APPLE_DH_SERVER_KEY_LENGTH] = {};
	crypto_key_p(pub, mod, sizeof(mod));

	uint8_t q[APPLE_DH_SERVER_KEY_LENGTH] = {};
	crypto_key_q(pub, q, sizeof(q));

	struct rfb_apple_dh_server_msg msg = {
		.generator = htons(crypto_key_g(client->apple_dh_secret)),
		.key_size = htons(sizeof(q)),
	};

	stream_write(client->net_stream, &msg, sizeof(msg));
	stream_write(client->net_stream, mod, sizeof(mod));
	stream_write(client->net_stream, q, sizeof(q));

	crypto_key_del(pub);
	return 0;
}

int apple_dh_handle_response(struct nvnc_client* client)
{
	struct rfb_apple_dh_client_msg* msg =
	        (void*)(client->msg_buffer + client->buffer_index);

	uint8_t p[APPLE_DH_SERVER_KEY_LENGTH];
	int key_len = crypto_key_p(client->apple_dh_secret, p, sizeof(p));
	assert(key_len == sizeof(p));

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

	int g = crypto_key_g(client->apple_dh_secret);

	struct crypto_key* remote_key = crypto_key_new(g, p, key_len,
			msg->public_key, key_len);
	assert(remote_key);

	struct crypto_key* shared_secret =
		crypto_derive_shared_secret(client->apple_dh_secret, remote_key);
	assert(shared_secret);
	crypto_key_del(remote_key);

	uint8_t shared_buf[APPLE_DH_SERVER_KEY_LENGTH];
	crypto_key_q(shared_secret, shared_buf, sizeof(shared_buf));
	crypto_key_del(shared_secret);

	uint8_t hash[16] = {};
	crypto_hash_one(hash, sizeof(hash), CRYPTO_HASH_MD5, shared_buf,
			sizeof(shared_buf));

	struct crypto_cipher* cipher;
	cipher = crypto_cipher_new(NULL, hash, CRYPTO_CIPHER_AES128_ECB);
	assert(cipher);

	char username[128] = {};
	char* password = username + 64;

	crypto_cipher_decrypt(cipher, (uint8_t*)username, NULL,
			msg->encrypted_credentials, sizeof(username), NULL, 0);
	username[63] = '\0';
	username[127] = '\0';
	crypto_cipher_del(cipher);

	update_min_rtt(client);

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

	strncpy(creds->username, username, sizeof(creds->username) - 1);
	strncpy(creds->password, password, sizeof(creds->password) - 1);

	static_assert(sizeof(creds->username) >= sizeof(username) / 2,
			"creds.username is too small");
	static_assert(sizeof(creds->password) >= sizeof(username) / 2,
			"creds.password is too small");

	security_handshake_authenticate(client, creds);

	return sizeof(*msg) + key_len;
}
