/* BSD-2-Clause license
 * 
 * Copyright (c) 2018-2023 NST <www.newinfosec.ru>, sss <sss at dark-alexandr dot net>. 
 * 
 */

#include <openssl/err.h>
#include <ev.h>
#include <stdbool.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <json.h>

#include "globals.h"
#include "socket_helpers.h"
#include "wrdp_thpool.h"
#include "wrdp_thpool_internals.h"
#include "remote_control.h"
#include "json_helpers.h"

#include "webrdp_core_api.h"
#include "log.h"

static int ctl_server_tcp_sock = -1, ctl_server_unix_sock = -1;
static SSL_CTX *ctx;

/* TODO: find out why this code affecting global certificate store */

static void
init_openssl()
{
	SSL_load_error_strings();
	OpenSSL_add_ssl_algorithms();
}

static SSL_CTX *
create_context()
{
	const SSL_METHOD *method = TLS_server_method();

	ctx = SSL_CTX_new(method);
	if (!ctx)
	{
		perror("Unable to create SSL context");
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}

	return ctx;
}

/* this is default behaviour, no need to override */
/*static int
verify_callback(int verify_ok, X509_STORE_CTX *s)
{
	return verify_ok;
}*/

static void
configure_context(SSL_CTX *ctx)
{
	if (g_globals.settings.ctl_ssl_cafile
	    || g_globals.settings.ctl_ssl_capath)
	{
		if (!SSL_CTX_load_verify_locations(ctx,
			g_globals.settings.ctl_ssl_cafile,
			g_globals.settings.ctl_ssl_capath))
		{
			ERR_print_errors_fp(stderr);
			exit(EXIT_FAILURE);
		}
	}
	else
	{
		if (!SSL_CTX_set_default_verify_paths(ctx))
		{
			ERR_print_errors_fp(stderr);
			exit(EXIT_FAILURE);
		}
	}

	SSL_CTX_set_ecdh_auto(ctx, 1);

	/*TODO: configurable file path*/
	/* Set the key and cert */
	if (SSL_CTX_use_certificate_file(
		ctx, g_globals.settings.ctl_ssl_cert, SSL_FILETYPE_PEM)
	    <= 0)
	{
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}

	if (SSL_CTX_use_PrivateKey_file(
		ctx, g_globals.settings.ctl_ssl_key, SSL_FILETYPE_PEM)
	    <= 0)
	{
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}

	SSL_CTX_set_verify(
	    ctx, SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, 0);
}

static void
on_ctl_task_destroy(wrdp_thpool_task *task)
{
	ctl_task_info *info = task->userdata;
	ctl_session *session = info->session;
	{
		const char *msg = "cleaning task slot";
		log_msg(
		    (const uint8_t *)msg, strlen(msg), wrdp_log_level_trace, 0);
	}
	ev_io_stop(task->thread->ev_th_loop, &(info->session->ev_con_fd_r));
	if (ev_is_active(&(info->session->ev_con_fd_w)))
	{
		ev_io_stop(
		    task->thread->ev_th_loop, &(info->session->ev_con_fd_w));
	}
	if (session->ssl)
	{
		SSL_free(session->ssl);
	}
	close(info->session->connection_fd);
	free(info->session);
	free(info);
}

void
ctl_destroy_task(wrdp_thpool_task *task)
{
	wrdp_thread_pool_destroy_task(task, on_ctl_task_destroy);
}

int
ctl_server_init_tcp()
{
	if (g_globals.settings.ctl_port <= 0)
	{
		ctl_server_tcp_sock = -1;
		return ctl_server_tcp_sock;
	}
	ctl_server_tcp_sock
	    = create_listen_socket_tcp(g_globals.settings.ctl_port);
	if (ctl_server_tcp_sock != -1)
	{
		socket_make_non_block(ctl_server_tcp_sock);
	}
	return ctl_server_tcp_sock;
}

int
ctl_server_init_unix()
{
	if (!g_globals.settings.ctl_socket_path
	    || !g_globals.settings.ctl_socket_path[0])
	{
		ctl_server_unix_sock = -1;
		return ctl_server_unix_sock;
	}
	ctl_server_unix_sock
	    = create_listen_socket_unix(g_globals.settings.ctl_socket_path);
	if (ctl_server_unix_sock != -1)
	{
		socket_make_non_block(ctl_server_unix_sock);
	}
	return ctl_server_unix_sock;
}

ctl_task_info *
ctl_create_task(int connection_fd)
{
	/* TODO: handle errors */
	ctl_task_info *info = 0;
	ctl_session *session = 0;
	int ssl_accept_ret = 0;
	info = calloc(1, sizeof(ctl_task_info));
	if (!info)
	{
		perror("calloc");
		goto error;
	}
	session = calloc(1, sizeof(ctl_session));
	if (!session)
	{
		perror("calloc");
		goto error;
	}
	session->connection_fd = connection_fd;
	info->session = session;
	session->ssl = SSL_new(ctx);
	SSL_set_fd(session->ssl, connection_fd);
	ssl_accept_ret = SSL_accept(session->ssl);
	if (ssl_accept_ret <= 0)
	{
		switch (SSL_get_error(session->ssl, ssl_accept_ret))
		{
			case SSL_ERROR_WANT_READ:
			{
				/* do nothing here, will be handled later */
			}
			break;
			case SSL_ERROR_WANT_WRITE:
			{
				/* TODO:
				 * 1. set write watcher
				 * 2. do SSL_write ?
				 */
			}
			break;
			default:
				goto error;
		}
	}
	else
	{
		session->state = ctl_session_connected;
	}
	return info;
error:
	if (session)
	{
		if (session->ssl)
		{
			SSL_free(session->ssl);
		}
		free(session);
	}
	if (info)
	{
		free(info);
	}
	return 0;
}

typedef enum
{
	cmd_get_sessions,
	cmd_kill
} cmd_type;

typedef struct
{
	cmd_type type;
	char **sessions, *message;
} ctl_cmd;

static bool
ctl_handle_json_option(struct json_object_element_s *json_option,
    wrdp_thpool_task *task, ctl_cmd *cmd)
{
	switch (json_option->value->type)
	{
		case json_type_object:
		{
		}
		break;
		default:
		{
			const char *msg = "unsupported json setting type";
			log_msg((const uint8_t *)msg, strlen(msg),
			    wrdp_log_level_error, 0);
			return false;
		}
		break;
	}
	return true;
}

static bool
ctl_handle_json_array_cmd(wrdp_thpool_task *task)
{
	ctl_task_info *info = task->userdata;
	ctl_session *session = info->session;
	struct json_parse_result_s res = {0};
	struct json_value_s *root = 0;
	ctl_cmd cmd;
	root = json_parse_ex(session->read_buf, session->cmd_size,
	    json_parse_flags_allow_json5, 0, 0, &res);
	if (!root)
	{
		const char *msg
		    = "Failed to parse remote control cmd data json";
		log_msg(
		    (const uint8_t *)msg, strlen(msg), wrdp_log_level_error, 0);
		return false;
	}
	{
		bool fail = false;
		struct json_object_s *object
		    = (struct json_object_s *)root->payload;
		struct json_object_element_s *option = object->start;
		while (option)
		{
			fail = !ctl_handle_json_option(option, task, &cmd);
			if (fail)
				break;
			option = option->next;
		}
		if (fail)
		{
			free(root);
			return false;
		}
	}
	free(root);
	return true;
}

bool
ctl_server_handle_data(void *taskdata)
{
	wrdp_thpool_task *task = taskdata;
	ctl_task_info *info = task->userdata;
	ctl_session *session = info->session;
	switch (session->state)
	{
		case ctl_session_ssl_handshake:
		{
			int ssl_accept_ret = 0;
			ssl_accept_ret = SSL_accept(session->ssl);
			if (ssl_accept_ret <= 0)
			{
				/* ERR_print_errors_fp(stderr); */
				int err = SSL_get_error(
				    session->ssl, ssl_accept_ret);
				if (err != SSL_ERROR_WANT_READ
				    && err != SSL_ERROR_WANT_WRITE)
				{
					char buf[64];
					snprintf(
					    buf, 63, "SSL_accept: %d", err);
					log_msg((const uint8_t *)buf,
					    strlen(buf), wrdp_log_level_error,
					    0);
					return false;
				}
			}
			else
			{
				session->state = ctl_session_connected;
			}
		}
		break;
		case ctl_session_connected:
		{
			if (!(session->cmd_size_known))
			{
				int32_t ret = SSL_read(session->ssl,
				    &(session->cmd_size) + session->read_size,
				    4 - session->read_size);
				if (ret <= 0)
				{
					int err
					    = SSL_get_error(session->ssl, ret);
					if (err != SSL_ERROR_WANT_READ
					    && err != SSL_ERROR_WANT_WRITE)
					{
						char buf[64];
						snprintf(buf, 63,
						    "SSL_read: %d", err);
						log_msg((const uint8_t *)buf,
						    strlen(buf),
						    wrdp_log_level_error, 0);
						return false;
					}
				}
				session->read_size += ret;
				if (session->read_size == 4)
				{
					session->cmd_size_known = true;
					session->cmd_size
					    = ntohl(session->cmd_size);
					session->read_buf
					    = malloc(session->cmd_size);
					if (!session->read_buf)
					{
						perror("malloc");
						return false;
					}
					/* reset read_size */
					session->read_size = 0;
				}
			}
			else
			{
				int32_t ret = SSL_read(session->ssl,
				    session->read_buf + session->read_size,
				    session->cmd_size - session->read_size);
				if (ret <= 0)
				{
					int err
					    = SSL_get_error(session->ssl, ret);
					if (err != SSL_ERROR_WANT_READ
					    && err != SSL_ERROR_WANT_WRITE)
					{
						char buf[64];
						snprintf(buf, 63,
						    "SSL_read: %d", err);
						log_msg((const uint8_t *)buf,
						    strlen(buf),
						    wrdp_log_level_error, 0);
						return false;
					}
				}
				session->read_size += ret;
				if (session->read_size == session->cmd_size)
				{
					ctl_handle_json_array_cmd(task);
					/* reset cmd_size_known */
					session->cmd_size_known = false;
					/* reset read_size, cmd_size */
					session->cmd_size = session->read_size
					    = 0;
					/* free read_buf */
					free(session->read_buf);
				}
			}
		}
		break;
		default:
			break;
	}
	return true;
}

void
init_remote_control()
{
	init_openssl();
	create_context();
	configure_context(ctx);
}