/*
 * Storage implementation for antispam plugin
 * Copyright 2007-2008, 2025	Johannes Berg <johannes@sipsolutions.net>
 * Copyright 2009	Jonas Maurus <jonas@maurus.net>
 *
 * Derived from Quota plugin:
 * Copyright (C) 2005 Timo Sirainen
 */

#include <sys/stat.h>

#include "lib.h"
#include "array.h"
#include "istream.h"
#include "mail-search.h"
#include "mail-index.h"
#include "mailbox-list-private.h"
#include "mail-storage-private.h"
#include "settings.h"
#include "imap-common.h"

#include "antispam-plugin.h"

#define ANTISPAM_CONTEXT(obj) \
	MODULE_CONTEXT(obj, antispam_storage_module)
#define ANTISPAM_USER_CONTEXT(obj) \
	MODULE_CONTEXT(obj, antispam_user_module)
#define ANTISPAM_MAIL_CONTEXT(obj) \
	MODULE_CONTEXT(obj, antispam_mail_module)

static MODULE_CONTEXT_DEFINE_INIT(antispam_storage_module,
                                  &mail_storage_module_register);
static MODULE_CONTEXT_DEFINE_INIT(antispam_user_module,
                                  &mail_user_module_register);
static MODULE_CONTEXT_DEFINE_INIT(antispam_mail_module,
                                  &mail_module_register);

                                 
enum mailbox_move_type {
	MMT_APPEND,
	MMT_UNINTERESTING,
	MMT_TO_CLEAN,
	MMT_TO_SPAM,
};

struct antispam_internal_context {
	union mailbox_transaction_module_context module_ctx;
	struct antispam_transaction_context *backendctx;
	struct mail *mail;
};

static enum classification move_to_class(enum mailbox_move_type tp)
{
	switch (tp) {
	case MMT_TO_CLEAN:
		return CLASS_NOTSPAM;
	case MMT_TO_SPAM:
		return CLASS_SPAM;
	default:
		i_assert(0);
	}
}

struct antispam_mailbox {
	union mailbox_module_context module_ctx;

	const struct antispam_settings *cfg;

	enum mailbox_move_type movetype;

	/* used to check if copy was implemented with save */
	unsigned int save_hack:1;
};

struct antispam_mail_user {
	union mail_user_module_context module_ctx;

	struct antispam_settings *cfg;
};

struct antispam_mail {
	union mail_module_context module_ctx;

	const struct antispam_settings *cfg;
};

static int
antispam_copy(struct mail_save_context *ctx, struct mail *mail)
{
	struct mailbox_transaction_context *t = ctx->transaction;
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(t->box);
	struct antispam_internal_context *ast = ANTISPAM_CONTEXT(t);
	int ret;
	bool src_trash, dst_trash;

	if (!ctx->dest_mail) {
		/* always need mail */
		if (!ast->mail)
			ast->mail = mail_alloc(t, MAIL_FETCH_STREAM_HEADER |
						  MAIL_FETCH_STREAM_BODY,
					       NULL);
		ctx->dest_mail = ast->mail;
	}

	i_assert(mail->box);

	asbox->save_hack = FALSE;
	asbox->movetype = MMT_UNINTERESTING;

	if (mailbox_is_unsure(asbox->cfg, t->box)) {
		mail_storage_set_error(t->box->storage, MAIL_ERROR_NOTPOSSIBLE,
				       "Cannot copy to unsure folder");
		mailbox_save_cancel(&ctx);
		return -1;
	}

	src_trash = mailbox_is_trash(asbox->cfg, mail->box);
	dst_trash = mailbox_is_trash(asbox->cfg, t->box);

	debug_verbose(asbox->cfg,
		      "mail copy: from trash: %d, to trash: %d\n",
	              src_trash, dst_trash);

	if (!src_trash && !dst_trash) {
		bool src_spam = mailbox_is_spam(asbox->cfg, mail->box);
		bool dst_spam = mailbox_is_spam(asbox->cfg, t->box);
		bool src_unsu = mailbox_is_unsure(asbox->cfg, mail->box);

		debug_verbose(asbox->cfg,
			      "mail copy: src spam: %d, dst spam: %d,"
		              " src unsure: %d\n",
		              src_spam, dst_spam, src_unsu);

		if ((src_spam || src_unsu) && !dst_spam)
			asbox->movetype = MMT_TO_CLEAN;
		else if ((!src_spam || src_unsu) && dst_spam)
			asbox->movetype = MMT_TO_SPAM;
	}

	if (asbox->module_ctx.super.copy(ctx, mail) < 0)
		return -1;

	/*
	 * If copying used saving internally, we already have treated the mail
	 */
	if (asbox->save_hack || asbox->movetype == MMT_UNINTERESTING)
		ret = 0;
	else
		ret = asbox->cfg->backend->handle_mail(
				asbox->cfg, t, ast->backendctx,
				ctx->dest_mail,
				move_to_class(asbox->movetype));

	/*
	 * Both save_hack and movetype are only valid within a copy operation,
	 * i.e. they are now invalid. Because, in theory, another operation
	 * could be done after mailbox_open(), we need to reset the movetype
	 * variable here. save_hack doesn't need to be reset because it is
	 * only ever set within the save function and tested within this copy
	 * function after being reset at the beginning of the copy, movetype
	 * however is tested within the save_finish() function and a subsequent
	 * save to the mailbox should not invoke the backend.
	 */
	asbox->movetype = MMT_APPEND;
	return ret;
}

static int antispam_save_begin(struct mail_save_context *ctx, 
                               struct istream *input)
{
	struct mailbox_transaction_context *t = ctx->transaction;
	struct antispam_internal_context *ast = ANTISPAM_CONTEXT(t);
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(t->box);
	int ret;

	if (!ctx->dest_mail) {
		if (!ast->mail)
			ast->mail = mail_alloc(t, MAIL_FETCH_STREAM_HEADER |
						  MAIL_FETCH_STREAM_BODY,
					       NULL);
		ctx->dest_mail = ast->mail;
	}
	ret = asbox->module_ctx.super.save_begin(ctx, input);

	return ret;
}

static int antispam_save_finish(struct mail_save_context *ctx)
{
	struct antispam_mailbox *asbox =
		ANTISPAM_CONTEXT(ctx->transaction->box);
	struct antispam_internal_context *ast =
		ANTISPAM_CONTEXT(ctx->transaction);
	struct mail *dest_mail;
	int ret;

	if (asbox->module_ctx.super.save_finish(ctx) < 0)
		return -1;

	dest_mail = ctx->dest_mail ? : ast->mail;

	asbox->save_hack = TRUE;

	ret = 0;

	switch (asbox->movetype) {
	case MMT_UNINTERESTING:
		break;
	case MMT_APPEND:
		/* Disallow APPENDs to UNSURE folders. */
		if (mailbox_is_unsure(asbox->cfg, dest_mail->box)) {
			ret = -1;
			mail_storage_set_error(dest_mail->box->storage,
					MAIL_ERROR_NOTPOSSIBLE,
					"Cannot APPEND to an UNSURE folder.");
			break;
		} else if (mailbox_is_spam(asbox->cfg, dest_mail->box)) {
			/*
			 * The client is APPENDing a message to a SPAM folder
			 * so we try to train the backend on it. For most of
			 * the backends, that can only succeed if the message
			 * contains appropriate information.
			 *
			 * This happens especially when offlineimap is used and
			 * the user moved a message to the SPAM folder while
			 * offline---offlineimap cannot reproduce the COPY but
			 * rather APPENDs the moved message on the next sync.
			 *
			 * This could be a bad if the spam headers were not
			 * generated on our server, but since the user can
			 * always APPEND to another folder and then COPY to a
			 * SPAM folder backends need to be prepared for cases
			 * like this anyway. With dspam, for example, the worst
			 * that can happen is that the APPEND fails with a
			 * training error from dspam.
			 *
			 * Unfortunately, we cannot handle the cases where
			 *  (1) the user moved a message from one folder that
			 *      contains SPAM to another folder containing SPAM
			 *  (2) the user moved a message out of the SPAM folder
			 *  (3) the user recovered a message from trash
			 *
			 * Because of these limitations, this behaviour needs
			 * to be enabled with an option.
			 */
			if (!asbox->cfg->allow_append_to_spam) {
				ret = -1;
				mail_storage_set_error(
					dest_mail->box->storage,
					MAIL_ERROR_NOTPOSSIBLE,
					"Cannot APPEND to a SPAM folder.");
				break;
			}
			asbox->movetype = MMT_TO_SPAM;
			/* fall through to default case to invoke backend */
		} else {
			/* neither UNSURE nor SPAM, regular folder */
			break;
		}
		/* fall through */
	default:
		ret = asbox->cfg->backend->handle_mail(
				asbox->cfg, ctx->transaction, ast->backendctx,
				dest_mail, move_to_class(asbox->movetype));
	}

	return ret;
}

static struct antispam_transaction_context *
antispam_transaction_begin(struct mailbox *box)
{
	struct antispam_transaction_context *ast;
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(box);

	ast = asbox->cfg->backend->start(asbox->cfg, box);
	i_assert(ast != NULL);

	return ast;
}

static void
antispam_transaction_rollback(const struct antispam_settings *cfg,
			      struct antispam_transaction_context **_ast)
{
	struct antispam_transaction_context *ast = *_ast;

	cfg->backend->rollback(cfg, ast);
	*_ast = NULL;
}

static int
antispam_transaction_commit(const struct antispam_settings *cfg,
			    struct mailbox_transaction_context *ctx,
			    struct antispam_transaction_context **_ast)
{
	struct antispam_transaction_context *ast = *_ast;
	int ret;

	ret = cfg->backend->commit(cfg, ctx, ast);
	*_ast = NULL;
	return ret;
}

static void
antispam_mail_update_keywords(struct mail *mail,
			      enum modify_type modify_type,
			      struct mail_keywords *keywords)
{
	struct mail_private *pmail = (struct mail_private *)mail;
	struct antispam_mail *amail = ANTISPAM_MAIL_CONTEXT(pmail);
	unsigned int i, numkwds;
	const ARRAY_TYPE(keywords) *idxkwd = mail_index_get_keywords(keywords->index);
	const char *const *keyword_names = array_get(idxkwd, &numkwds);
	const char *const *orig_keywords;
	bool previous_spam_keyword, now_spam_keyword;

	switch (modify_type) {
	case MODIFY_ADD:
		debug(amail->cfg, "adding keyword(s)\n");
		break;
	case MODIFY_REMOVE:
		debug(amail->cfg, "removing keyword(s)\n");
		break;
	case MODIFY_REPLACE:
		debug(amail->cfg, "replacing keyword(s)\n");
		break;
	default:
		i_assert(0);
	}

	orig_keywords = pmail->v.get_keywords(mail);
	if (orig_keywords) {
		debug(amail->cfg, "original keyword list:\n");
		while (*orig_keywords) {
			debug(amail->cfg, " * %s\n", *orig_keywords);
			if (keyword_is_spam(amail->cfg, *orig_keywords))
				previous_spam_keyword = TRUE;
			orig_keywords++;
		}
	}

	debug(amail->cfg, "keyword list:\n");

	for (i = 0; i < keywords->count; i++) {
		unsigned int idx = keywords->idx[i];

		i_assert(idx < numkwds);

		debug(amail->cfg, " * %s\n", keyword_names[idx]);

		switch (modify_type) {
		case MODIFY_ADD:
		case MODIFY_REPLACE:
			if (keyword_is_spam(amail->cfg, keyword_names[idx]))
				now_spam_keyword = TRUE;
			break;
		case MODIFY_REMOVE:
			if (keyword_is_spam(amail->cfg, keyword_names[idx]))
				now_spam_keyword = FALSE;
			break;
		default:
			i_assert(0);
		}
	}

	amail->module_ctx.super.update_keywords(mail, modify_type, keywords);

	debug(amail->cfg, "previous-spam, now-spam: %d, %d\n",
	      previous_spam_keyword, now_spam_keyword);

	if (previous_spam_keyword != now_spam_keyword) {
		/*
		 * Call backend here.
		 *
		 * TODO: It is not clear how to roll back the
		 *       keyword change if the backend fails.
		 */
	}
}

static struct mailbox_transaction_context *
antispam_mailbox_transaction_begin(struct mailbox *box,
				   enum mailbox_transaction_flags flags,
				   const char *reason)
{
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(box);
	struct mailbox_transaction_context *t;
	struct antispam_transaction_context *ast;
	struct antispam_internal_context *aic;

	t = asbox->module_ctx.super.transaction_begin(box, flags, reason);
	aic = i_new(struct antispam_internal_context, 1);
	ast = antispam_transaction_begin(box);
	aic->backendctx = ast;

	MODULE_CONTEXT_SET(t, antispam_storage_module, aic);
	return t;
}

static int
antispam_mailbox_transaction_commit(struct mailbox_transaction_context *ctx,
				    struct mail_transaction_commit_changes *changes_r)
{
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(ctx->box);
	struct antispam_internal_context *ast = ANTISPAM_CONTEXT(ctx);

	if (antispam_transaction_commit(asbox->cfg, ctx, &ast->backendctx) < 0) {
		if (ast->mail)
			mail_free(&ast->mail);

		asbox->module_ctx.super.transaction_rollback(ctx);
		return -1;
	}

	if (ast->mail)
		mail_free(&ast->mail);

	return asbox->module_ctx.super.transaction_commit(ctx, changes_r);
}

static void
antispam_mailbox_transaction_rollback(struct mailbox_transaction_context *ctx)
{
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(ctx->box);
	struct antispam_internal_context *ast = ANTISPAM_CONTEXT(ctx);

	if (ast->mail)
		mail_free(&ast->mail);

	asbox->module_ctx.super.transaction_rollback(ctx);

	antispam_transaction_rollback(asbox->cfg, &ast->backendctx);
}

static void antispam_mail_allocated(struct mail *_mail)
{
	struct mail_private *mail = (struct mail_private *)_mail;
	struct mail_vfuncs *v = mail->vlast;
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(_mail->box);
	struct antispam_mail *amail;

	if (asbox == NULL)
		return;

	amail = p_new(mail->pool, struct antispam_mail, 1);
	amail->module_ctx.super = *v;
	mail->vlast = &amail->module_ctx.super;

	amail->cfg = asbox->cfg;
	if (!!asbox->cfg->spam_keywords)
		v->update_keywords = antispam_mail_update_keywords;

	MODULE_CONTEXT_SET(mail, antispam_mail_module, amail);
}

static void antispam_mailbox_free(struct mailbox *box)
{
	struct antispam_mailbox *asbox = ANTISPAM_CONTEXT(box);

	asbox->module_ctx.super.free(box);
}

static void antispam_mailbox_allocated(struct mailbox *box)
{
	struct mail_user *user = box->storage->user;
	struct mailbox_vfuncs *v = box->vlast;
	struct antispam_mailbox *asbox;
	struct antispam_mail_user *asuser = ANTISPAM_USER_CONTEXT(user);

	if (asuser == NULL)
		return;

	asbox = p_new(box->pool, struct antispam_mailbox, 1);
	asbox->module_ctx.super = *v;
	box->vlast = &asbox->module_ctx.super;

	asbox->save_hack = FALSE;
	asbox->movetype = MMT_APPEND;
	asbox->cfg = asuser->cfg;

	v->free = antispam_mailbox_free;

	if (asbox->cfg->need_folder_hook) {
		/* override save_init to override want_mail, we need that */
		v->save_begin = antispam_save_begin;
		v->save_finish = antispam_save_finish;
		v->transaction_begin = antispam_mailbox_transaction_begin;
		v->transaction_commit = antispam_mailbox_transaction_commit;
		v->transaction_rollback = antispam_mailbox_transaction_rollback;
		v->copy = antispam_copy;
	}

	MODULE_CONTEXT_SET(box, antispam_storage_module, asbox);
}

static void antispam_user_deinit(struct mail_user *user)
{
	struct antispam_mail_user *asuser = ANTISPAM_USER_CONTEXT(user);

	antispam_free_config(asuser->cfg);

	asuser->module_ctx.super.deinit(user);
}

void antispam_mail_user_created(struct mail_user *user)
{
	struct mail_user_vfuncs *v = user->vlast;
	struct antispam_mail_user *asuser;

	asuser = p_new(user->pool, struct antispam_mail_user, 1);
	asuser->module_ctx.super = *v;
	user->vlast = &asuser->module_ctx.super;
	v->deinit = antispam_user_deinit;

	MODULE_CONTEXT_SET(user, antispam_user_module, asuser);
}

static struct mail_storage_hooks antispam_mail_storage_hooks = {
	.mail_user_created = antispam_mail_user_created,
	.mailbox_allocated = antispam_mailbox_allocated,
	.mail_allocated = antispam_mail_allocated,
};

static imap_client_created_func_t *next_hook_client_created;
static struct module *antispam_module;

static void antispam_storage_client_created(struct client *client)
{
	struct mail_user *user = client->user;
	struct antispam_mail_user *asuser = ANTISPAM_USER_CONTEXT(user);
	struct antispam_settings *set;
	const char *error;

	if (settings_get(client->event, &antispam_setting_parser_info, 0,
			 &set, &error) < 0) {
		e_error(client->event, "%s", error);
		return;
	}

	asuser->cfg = set;
}

static void antispam_client_created(struct client **clientp)
{
	struct client *client = *clientp;
	struct mail_user *user = client->user;

	if (mail_user_is_plugin_loaded(user, antispam_module))
		antispam_storage_client_created(client);

	if (next_hook_client_created)
		next_hook_client_created(clientp);
}

void antispam_storage_init(struct module *module)
{
	antispam_module = module;
	mail_storage_hooks_add(module, &antispam_mail_storage_hooks);
	next_hook_client_created =
		imap_client_created_hook_set(antispam_client_created);
}

void antispam_storage_deinit(void)
{
	mail_storage_hooks_remove(&antispam_mail_storage_hooks);
	imap_client_created_hook_set(next_hook_client_created);
}
