/*
 * Ibis - IRCv3 Library
 * Copyright (C) 2022-2024 Ibis Developers
 *
 * Ibis is the legal property of its developers, whose names are too numerous
 * to list here.  Please refer to the COPYRIGHT file distributed with this
 * source distribution.
 *
 * This library 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 2 of the License, or
 * (at your option) any later version.
 *
 * This library 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 library; if not, see <https://www.gnu.org/licenses/>.
 */

#include <glib/gi18n-lib.h>

#include <hasl.h>

#include "ibissasl.h"

#include "ibisconstants.h"
#include "ibiscore.h"
#include "ibisstring.h"

#define IBIS_SASL_DATA_KEY ("sasl-data")

typedef struct {
	GString *server_in_buffer;
} IbisSASLData;

/******************************************************************************
 * SASL Helpers
 *****************************************************************************/
static void
ibis_sasl_data_free(IbisSASLData *data) {
	g_string_free(data->server_in_buffer, TRUE);

	g_free(data);
}

static void
ibis_sasl_data_add(IbisClient *client) {
	IbisSASLData *data = NULL;

	data = g_new0(IbisSASLData, 1);
	g_object_set_data_full(G_OBJECT(client), IBIS_SASL_DATA_KEY,
	                       data, (GDestroyNotify)ibis_sasl_data_free);

	/* We truncate the server_in_buffer when we need to so that we can minimize
	 * allocations and simplify the logic involved with it.
	 */
	data->server_in_buffer = g_string_new("");
}

static void
ibis_sasl_attempt(IbisClient *client) {
	IbisMessage *message = NULL;
	HaslContext *hasl_context = NULL;
	const char *next_mechanism = NULL;

	hasl_context = ibis_client_get_hasl_context(client);

	next_mechanism = hasl_context_next(hasl_context);
	if(next_mechanism == NULL) {
		GError *error = g_error_new(IBIS_DOMAIN, 0,
		                            _("No valid SASL mechanisms found"));

		hasl_context_done(hasl_context);

		ibis_client_disconnect(client, error, NULL);

		return;
	}

	g_message("trying SASL '%s' mechanism", next_mechanism);

	message = ibis_message_new(IBIS_MSG_AUTHENTICATE);
	ibis_message_set_params(message, next_mechanism, NULL);
	ibis_client_write(client, message);
}

/******************************************************************************
 * Callbacks
 *****************************************************************************/
static gboolean
ibis_sasl_logged_in(G_GNUC_UNUSED IbisClient *client,
                    G_GNUC_UNUSED const char *command,
                    G_GNUC_UNUSED IbisMessage *message,
                    G_GNUC_UNUSED gpointer data)
{
	return TRUE;
}

static gboolean
ibis_sasl_logged_out(G_GNUC_UNUSED IbisClient *client,
                     G_GNUC_UNUSED const char *command,
                     G_GNUC_UNUSED IbisMessage *message,
                     G_GNUC_UNUSED gpointer data)
{
	/* Not sure how to trigger this or what we should do in this case to be
	 * honest, so just note it for now.
	 * -- GK 2023-01-12
	 */
	g_warning("Server sent SASL logged out");

	return TRUE;
}

static gboolean
ibis_sasl_nick_locked(IbisClient *client,
                      G_GNUC_UNUSED const char *command,
                      IbisMessage *ibis_message,
                      G_GNUC_UNUSED gpointer data)
{
	GError *error = NULL;
	GStrv params = NULL;
	char *message = NULL;

	params = ibis_message_get_params(ibis_message);
	message = g_strjoinv(" ", params);

	error = g_error_new(IBIS_DOMAIN, 0, _("Nick name is locked: %s"), message);
	g_free(message);

	ibis_client_disconnect(client, error, NULL);

	return TRUE;
}


static gboolean
ibis_sasl_success(IbisClient *client,
                  G_GNUC_UNUSED const char *command,
                  G_GNUC_UNUSED IbisMessage *message,
                  G_GNUC_UNUSED gpointer user_data)
{
	IbisCapabilities *capabilities = NULL;
	HaslContext *hasl_context = NULL;

	capabilities = ibis_client_get_capabilities(client);
	ibis_capabilities_remove_wait(capabilities);

	hasl_context = ibis_client_get_hasl_context(client);
	g_message("successfully authenticated with SASL '%s' mechanism.",
	          hasl_context_get_current_mechanism(hasl_context));
	hasl_context_done(hasl_context);

	return TRUE;
}

static gboolean
ibis_sasl_failed(IbisClient *client,
                 G_GNUC_UNUSED const char *command,
                 G_GNUC_UNUSED IbisMessage *message,
                 G_GNUC_UNUSED gpointer data)
{
	HaslContext *hasl_context = NULL;
	const char *current = NULL;

	hasl_context = ibis_client_get_hasl_context(client);
	current = hasl_context_get_current_mechanism(hasl_context);
	if(current != NULL) {
		g_message("SASL '%s' mechanism failed", current);
	}

	ibis_sasl_attempt(client);

	return TRUE;
}

static gboolean
ibis_sasl_message_too_long(G_GNUC_UNUSED IbisClient *client,
                           G_GNUC_UNUSED const char *command,
                           G_GNUC_UNUSED IbisMessage *message,
                           G_GNUC_UNUSED gpointer data)
{
	g_warning("received SASL_MESSAGE_TOO_LONG");

	return TRUE;
}

static gboolean
ibis_sasl_aborted(IbisClient *client,
                  G_GNUC_UNUSED const char *command,
                  G_GNUC_UNUSED IbisMessage *message,
                  G_GNUC_UNUSED gpointer data)
{
	HaslContext *hasl_context = NULL;

	hasl_context = ibis_client_get_hasl_context(client);
	hasl_context_done(hasl_context);

	/* Return FALSE so the clients can handle this how they like. */

	return FALSE;
}

static gboolean
ibis_sasl_already_authed(IbisClient *client,
                         G_GNUC_UNUSED const char *command,
                         G_GNUC_UNUSED IbisMessage *message,
                         G_GNUC_UNUSED gpointer data)
{
	HaslContext *hasl_context = NULL;

	hasl_context = ibis_client_get_hasl_context(client);
	hasl_context_done(hasl_context);

	/* Return FALSE so clients can handle this how they like. */

	return FALSE;
}

static gboolean
ibis_sasl_mechanisms(G_GNUC_UNUSED IbisClient *client,
                     G_GNUC_UNUSED const char *command,
                     IbisMessage *message,
                     G_GNUC_UNUSED gpointer data)
{
	GStrv params = NULL;
	guint n_params = 0;

	params = ibis_message_get_params(message);
	if(params != NULL) {
		n_params = g_strv_length(params);
	}

	/* We need to find a server that sends this message. The specification says
	 * it _may_ be sent when the client sends AUTHENTICATE with an unknown
	 * mechanism, but ergo doesn't.
	 *
	 * We can't just blindly accept the new list either, as depending on how
	 * the server implements it, we'll need to remove mechanisms we've already
	 * tried in the event the server just dumps the entire list. As we're not
	 * currently tracking which mechanisms we've tried, this will have to be
	 * addressed as well.
	 */
	if(n_params > 0) {
		char *message = g_strjoinv(" ", params);

		g_message("Server sent the following SASL mechanisms: %s", message);

		g_free(message);
	} else {
		g_message("Server sent an empty list of SASL mechanisms");
	}

	return TRUE;
}

static gboolean
ibis_sasl_authenticate(IbisClient *client,
                       G_GNUC_UNUSED const char *command,
                       IbisMessage *message,
                       G_GNUC_UNUSED gpointer user_data)
{
	IbisSASLData *data = NULL;
	GStrv params = NULL;
	char *payload = NULL;
	gboolean done = FALSE;
	guint n_params = 0;

	params = ibis_message_get_params(message);
	if(params != NULL) {
		n_params = g_strv_length(params);
	}

	if(n_params != 1) {
		g_warning("ignoring AUTHENTICATE with %d parameters", n_params);

		return TRUE;
	}

	payload = params[0];
	data = g_object_get_data(G_OBJECT(client), IBIS_SASL_DATA_KEY);
	if(data == NULL) {
		g_warning("AUTHENTICATE received with no SASL data present");

		return TRUE;
	}

	/* If the server sent us a payload, combine the chunks. */
	if(payload[0] != '+') {
		g_string_append(data->server_in_buffer, payload);

		if(strlen(payload) < 400) {
			done = TRUE;
		}
	} else {
		/* The server sent a + which is an empty message or the final message
		 * ended on a 400 byte barrier. */
		done = TRUE;
	}

	if(done) {
		IbisMessage *response = NULL;
		HaslContext *hasl_context = NULL;
		HaslMechanismResult res = 0;
		GError *error = NULL;
		guint8 *server_in = NULL;
		guint8 *client_out = NULL;
		gsize server_in_length = 0;
		size_t client_out_length = 0;

		/* If we have a buffer, base64 decode it, and then truncate it. */
		if(data->server_in_buffer->len > 0) {
			server_in = g_base64_decode(data->server_in_buffer->str,
			                            &server_in_length);
			g_string_truncate(data->server_in_buffer, 0);
		}

		/* Try to move to the next step of the sasl client. */
		hasl_context = ibis_client_get_hasl_context(client);
		res = hasl_context_step(hasl_context, server_in, server_in_length,
		                        &client_out, &client_out_length, &error);

		/* We should be done with server_in, so free it.*/
		g_clear_pointer(&server_in, g_free);

		if(res == HASL_MECHANISM_RESULT_ERROR) {
			const char *error_message = "Unknown error";

			if(error != NULL) {
				error_message = error->message;
			}

			g_warning("hasl mechanism failed: %s", error_message);

			g_clear_error(&error);

			return TRUE;
		}

		if(error != NULL) {
			g_warning("hasl_context_step returned an error without an error "
			          "status: %s", error->message);

			g_clear_error(&error);
		}

		/* If we got an output for the client, write it out. */
		response = ibis_message_new(IBIS_MSG_AUTHENTICATE);

		if(client_out_length > 0) {
			char *encoded = NULL;

			encoded = g_base64_encode(client_out, client_out_length);
			ibis_message_set_params(response, encoded, NULL);

			g_clear_pointer(&client_out, g_free);
			g_free(encoded);
		} else {
			ibis_message_set_params(response, "+", NULL);
		}

		ibis_client_write(client, response);
	}

	return TRUE;
}

static void
ibis_sasl_setup(IbisCapabilities *capabilities, IbisClient *client) {
	HaslContext *hasl_context = NULL;
	const char *server_mechanisms = NULL;
	gboolean found = FALSE;

	hasl_context = ibis_client_get_hasl_context(client);

	/* Reset the context to its initial state. */
	hasl_context_done(hasl_context);

	/* Set our server mechanisms to whatever the server advertised. */
	server_mechanisms = ibis_capabilities_lookup(capabilities,
	                                             IBIS_CAPABILITY_SASL,
	                                             &found);
	if(found) {
		hasl_context_set_server_mechanisms(hasl_context, server_mechanisms);
	}

	/* There isn't a good way to know if our handlers are here or not, so we
	 * just clear them and reconnect them.
	 */
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_logged_in),
	                                     NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_logged_out),
	                                     NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_nick_locked),
	                                     NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_success), NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_failed), NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_message_too_long),
	                                     NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_aborted), NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_already_authed),
	                                     NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_mechanisms),
	                                     NULL);
	g_signal_handlers_disconnect_by_func(client,
	                                     G_CALLBACK(ibis_sasl_authenticate),
	                                     NULL);

	g_signal_connect(client, "message::" IBIS_RPL_LOGGEDIN,
	                 G_CALLBACK(ibis_sasl_logged_in), NULL);
	g_signal_connect(client, "message::" IBIS_RPL_LOGGEDOUT,
	                 G_CALLBACK(ibis_sasl_logged_out), NULL);
	g_signal_connect(client, "message::" IBIS_ERR_NICKLOCKED,
	                 G_CALLBACK(ibis_sasl_nick_locked), NULL);
	g_signal_connect(client, "message::" IBIS_RPL_SASLSUCCESS,
	                 G_CALLBACK(ibis_sasl_success), NULL);
	g_signal_connect(client, "message::" IBIS_ERR_SASLFAIL,
	                 G_CALLBACK(ibis_sasl_failed), NULL);
	g_signal_connect(client, "message::" IBIS_ERR_SASLTOOLONG,
	                 G_CALLBACK(ibis_sasl_message_too_long), NULL);
	g_signal_connect(client, "message::" IBIS_ERR_SASLABORTED,
	                 G_CALLBACK(ibis_sasl_aborted), NULL);
	g_signal_connect(client, "message::" IBIS_ERR_SASLALREADY,
	                 G_CALLBACK(ibis_sasl_already_authed), NULL);
	g_signal_connect(client, "message::" IBIS_RPL_SASLMECHS,
	                 G_CALLBACK(ibis_sasl_mechanisms), NULL);
	g_signal_connect(client, "message::" IBIS_MSG_AUTHENTICATE,
	                 G_CALLBACK(ibis_sasl_authenticate), NULL);

	ibis_sasl_data_add(client);

	ibis_capabilities_add_wait(capabilities);

	ibis_sasl_attempt(client);
}

static void
ibis_sasl_ack_cb(IbisCapabilities *capabilities,
                 IbisClient *client,
                 G_GNUC_UNUSED const char *capability,
                 G_GNUC_UNUSED gpointer data)
{
	ibis_sasl_setup(capabilities, client);
}

static void
ibis_sasl_new_cb(IbisCapabilities *capabilities,
                 G_GNUC_UNUSED IbisClient *client,
                 GStrv caps,
                 G_GNUC_UNUSED gpointer data)
{
	/* We have to run through the new caps as we only want to re-authenticate
	 * when SASL comes back and not every time a new capability is added.
	 */

	for(guint i = 0; caps[i] != NULL; i++) {
		if(g_str_has_prefix(caps[i], IBIS_CAPABILITY_SASL "=")) {
			ibis_capabilities_request(capabilities, IBIS_CAPABILITY_SASL);

			break;
		}
	}
}

/******************************************************************************
 * Internal API
 *****************************************************************************/
void
ibis_sasl_request(IbisCapabilities *capabilities) {
	ibis_capabilities_request(capabilities, IBIS_CAPABILITY_SASL);

	g_signal_connect(capabilities, "ack::" IBIS_CAPABILITY_SASL,
	                 G_CALLBACK(ibis_sasl_ack_cb), NULL);

	g_signal_connect(capabilities, "new", G_CALLBACK(ibis_sasl_new_cb), NULL);
}
