From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Tue, 3 May 2016 11:23:18 +0000 (+0100)
Subject: [Feature] Allow setting scores and actions from lua
X-Git-Tag: 1.3.0~571
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=cf13c6d210dfeab26634ab29bb2db2575882ebe0;p=rspamd.git

[Feature] Allow setting scores and actions from lua
---

diff --git a/src/libmime/filter.c b/src/libmime/filter.c
index b0f93ec4b..3a3203ed9 100644
--- a/src/libmime/filter.c
+++ b/src/libmime/filter.c
@@ -57,14 +57,14 @@ rspamd_create_metric_result (struct rspamd_task *task, const gchar *name)
 	metric_res->metric = metric;
 	metric_res->grow_factor = 0;
 	metric_res->score = 0;
-	g_hash_table_insert (task->results, (gpointer) metric->name,
-			metric_res);
 
 	for (i = 0; i < METRIC_ACTION_MAX; i++) {
 		metric_res->actions_limits[i] = metric->actions[i].score;
 	}
 
 	metric_res->action = METRIC_ACTION_MAX;
+	g_hash_table_insert (task->results, (gpointer) metric->name,
+			metric_res);
 
 	return metric_res;
 }
diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 3571e2e41..a394ddaff 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -835,7 +835,10 @@ rspamd_metric_result_ucl (struct rspamd_task *task,
 	const gchar *subject;
 
 	m = mres->metric;
-	mres->action = rspamd_check_action_metric (task, mres);
+
+	if (mres->action == METRIC_ACTION_MAX) {
+		mres->action = rspamd_check_action_metric (task, mres);
+	}
 
 	action = mres->action;
 	is_spam = (action < METRIC_ACTION_GREYLIST);
@@ -1064,7 +1067,14 @@ rspamd_protocol_http_reply (struct rspamd_http_message *msg,
 		/* Update stat for default metric */
 		metric_res = g_hash_table_lookup (task->results, DEFAULT_METRIC);
 		if (metric_res != NULL) {
-			action = rspamd_check_action_metric (task, metric_res);
+
+			if (metric_res->action != METRIC_ACTION_MAX) {
+				action = metric_res->action;
+			}
+			else {
+				action = rspamd_check_action_metric (task, metric_res);
+			}
+
 			if (action <= METRIC_ACTION_NOACTION) {
 #ifndef HAVE_ATOMIC_BUILTINS
 				task->worker->srv->stat->actions_stat[action]++;
diff --git a/src/libstat/stat_process.c b/src/libstat/stat_process.c
index 1c8fbf070..78166906f 100644
--- a/src/libstat/stat_process.c
+++ b/src/libstat/stat_process.c
@@ -793,7 +793,10 @@ rspamd_stat_check_autolearn (struct rspamd_task *task)
 					mres = g_hash_table_lookup (task->results, DEFAULT_METRIC);
 
 					if (mres) {
-						mres->action = rspamd_check_action_metric (task, mres);
+
+						if (mres->action == METRIC_ACTION_MAX) {
+							mres->action = rspamd_check_action_metric (task, mres);
+						}
 
 						if (mres->action == METRIC_ACTION_REJECT) {
 							task->flags |= RSPAMD_TASK_FLAG_LEARN_SPAM;
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 1f0383764..134876940 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -468,6 +468,21 @@ LUA_FUNCTION_DEF (task, get_metric_score);
  * @return {string} the current action of the metric as a string
  */
 LUA_FUNCTION_DEF (task, get_metric_action);
+/***
+ * @method task:set_metric_score(name, score)
+ * Set the current score of metric `name`. Should be used in post-filters only.
+ * @param {string} name name of a metric
+ * @param {number} score the current score of the metric
+ */
+LUA_FUNCTION_DEF (task, set_metric_score);
+/***
+ * @method task:set_metric_action(name, action)
+ * Set the current action of metric `name`. Should be used in post-filters only.
+ * @param {string} name name of a metric
+ * @param {string} action name to set
+ */
+LUA_FUNCTION_DEF (task, set_metric_action);
+
 /***
  * @method task:learn(is_spam[, classifier)
  * Learn classifier `classifier` with the task. If `is_spam` is true then message
@@ -645,6 +660,8 @@ static const struct luaL_reg tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, get_timeval),
 	LUA_INTERFACE_DEF (task, get_metric_score),
 	LUA_INTERFACE_DEF (task, get_metric_action),
+	LUA_INTERFACE_DEF (task, set_metric_score),
+	LUA_INTERFACE_DEF (task, set_metric_action),
 	LUA_INTERFACE_DEF (task, learn),
 	LUA_INTERFACE_DEF (task, set_settings),
 	LUA_INTERFACE_DEF (task, get_settings),
@@ -2883,6 +2900,77 @@ lua_task_get_metric_action (lua_State *L)
 	return 1;
 }
 
+static gint
+lua_task_set_metric_score (lua_State *L)
+{
+	struct rspamd_task *task = lua_check_task (L, 1);
+	const gchar *metric_name;
+	struct metric_result *metric_res;
+	gdouble nscore;
+
+	metric_name = luaL_checkstring (L, 2);
+	nscore = luaL_checknumber (L, 3);
+
+	if (metric_name == NULL) {
+		metric_name = DEFAULT_METRIC;
+	}
+
+	if (task && metric_name) {
+		if ((metric_res =
+			g_hash_table_lookup (task->results, metric_name)) != NULL) {
+			metric_res->score = nscore;
+			lua_pushboolean (L, true);
+		}
+		else {
+			lua_pushboolean (L, false);
+		}
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
+
+static gint
+lua_task_set_metric_action (lua_State *L)
+{
+	struct rspamd_task *task = lua_check_task (L, 1);
+	const gchar *metric_name, *action_name;
+	struct metric_result *metric_res;
+	gint action;
+
+	metric_name = luaL_checkstring (L, 2);
+
+	if (metric_name == NULL) {
+		metric_name = DEFAULT_METRIC;
+	}
+
+	action_name = luaL_checkstring (L, 3);
+
+	if (task && metric_name && action_name) {
+		if ((metric_res =
+			g_hash_table_lookup (task->results, metric_name)) != NULL) {
+
+			if (rspamd_action_from_str (action_name, &action)) {
+				metric_res->action = action;
+				lua_pushboolean (L, true);
+			}
+			else {
+				lua_pushboolean (L, false);
+			}
+		}
+		else {
+			lua_pushboolean (L, false);
+		}
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
+
 /* Image functions */
 static gint
 lua_image_get_width (lua_State *L)