commit c1791ae: [Feature] Fuzzy_check: Add gen_hashes command

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2020-06-04 16:07:24 +0100
URL: https://github.com/rspamd/rspamd/commit/c1791ae64d0e73606dd4703f2ddd9d8a68a06f23 (HEAD -> master)

[Feature] Fuzzy_check: Add gen_hashes command

---
 src/plugins/fuzzy_check.c | 134 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 134 insertions(+)

diff --git a/src/plugins/fuzzy_check.c b/src/plugins/fuzzy_check.c
index 36e093836..e33735c5b 100644
--- a/src/plugins/fuzzy_check.c
+++ b/src/plugins/fuzzy_check.c
@@ -185,6 +185,7 @@ static gint fuzzy_attach_controller (struct module_ctx *ctx,
 	GHashTable *commands);
 static gint fuzzy_lua_learn_handler (lua_State *L);
 static gint fuzzy_lua_unlearn_handler (lua_State *L);
+static gint fuzzy_lua_gen_hashes_handler (lua_State *L);
 
 module_t fuzzy_check_module = {
 		"fuzzy_check",
@@ -1147,6 +1148,9 @@ fuzzy_check_module_config (struct rspamd_config *cfg)
 		lua_pushstring (L, "learn");
 		lua_pushcfunction (L, fuzzy_lua_learn_handler);
 		lua_settable (L, -3);
+		lua_pushstring (L, "gen_hashes");
+		lua_pushcfunction (L, fuzzy_lua_gen_hashes_handler);
+		lua_settable (L, -3);
 		/* Finish fuzzy_check key */
 		lua_settable (L, -3);
 	}
@@ -3673,6 +3677,136 @@ fuzzy_lua_unlearn_handler (lua_State *L)
 	return 1;
 }
 
+static gint
+fuzzy_lua_gen_hashes_handler (lua_State *L)
+{
+	struct rspamd_task *task = lua_check_task (L, 1);
+	guint flag = 0, weight = 1, send_flags = 0;
+	const gchar *symbol;
+	struct fuzzy_ctx *fuzzy_module_ctx = fuzzy_get_context (task->cfg);
+	struct fuzzy_rule *rule;
+	GPtrArray *commands;
+	gint cmd = FUZZY_WRITE;
+	gint i;
+
+	if (task) {
+		if (lua_type (L, 2) == LUA_TNUMBER) {
+			flag = lua_tonumber (L, 2);
+		}
+		else if (lua_type (L, 2) == LUA_TSTRING) {
+			struct fuzzy_rule *rule;
+			guint i;
+			GHashTableIter it;
+			gpointer k, v;
+			struct fuzzy_mapping *map;
+
+			symbol = lua_tostring (L, 2);
+
+			PTR_ARRAY_FOREACH (fuzzy_module_ctx->fuzzy_rules, i, rule) {
+				if (flag != 0) {
+					break;
+				}
+
+				g_hash_table_iter_init (&it, rule->mappings);
+
+				while (g_hash_table_iter_next (&it, &k, &v)) {
+					map = v;
+
+					if (g_ascii_strcasecmp (symbol, map->symbol) == 0) {
+						flag = map->fuzzy_flag;
+						break;
+					}
+				}
+			}
+		}
+
+		if (flag == 0) {
+			return luaL_error (L, "bad flag");
+		}
+
+		if (lua_type (L, 3) == LUA_TNUMBER) {
+			weight = lua_tonumber (L, 3);
+		}
+
+		/* Flags */
+		if (lua_type (L, 4) == LUA_TTABLE) {
+			const gchar *sf;
+
+			for (lua_pushnil (L); lua_next (L, -2); lua_pop (L, 1)) {
+				sf = lua_tostring (L, -1);
+
+				if (sf) {
+					if (g_ascii_strcasecmp (sf, "noimages") == 0) {
+						send_flags |= FUZZY_CHECK_FLAG_NOIMAGES;
+					}
+					else if (g_ascii_strcasecmp (sf, "noattachments") == 0) {
+						send_flags |= FUZZY_CHECK_FLAG_NOATTACHMENTS;
+					}
+					else if (g_ascii_strcasecmp (sf, "notext") == 0) {
+						send_flags |= FUZZY_CHECK_FLAG_NOTEXT;
+					}
+				}
+			}
+		}
+
+		/* Type */
+		if (lua_type (L, 5) == LUA_TSTRING) {
+			const gchar *cmd_name = lua_tostring (L, 5);
+
+			if (strcmp (cmd_name, "add") == 0 || strcmp (cmd_name, "write") == 0) {
+				cmd = FUZZY_WRITE;
+			}
+			else if (strcmp (cmd_name, "delete") == 0 || strcmp (cmd_name, "remove") == 0) {
+				cmd = FUZZY_DEL;
+			}
+			else {
+				return luaL_error (L, "invalid command: %s", cmd_name);
+			}
+		}
+
+		lua_createtable (L, 0, fuzzy_module_ctx->fuzzy_rules->len);
+
+		PTR_ARRAY_FOREACH (fuzzy_module_ctx->fuzzy_rules, i, rule) {
+			if (rule->read_only) {
+				continue;
+			}
+
+			/* Check for flag */
+			if (g_hash_table_lookup (rule->mappings,
+					GINT_TO_POINTER (flag)) == NULL) {
+				msg_info_task ("skip rule %s as it has no flag %d defined"
+							   " false", rule->name, flag);
+				continue;
+			}
+
+			commands = fuzzy_generate_commands (task, rule, FUZZY_WRITE, flag,
+					weight, send_flags);
+
+			if (commands != NULL) {
+				struct fuzzy_cmd_io *io;
+				gint j;
+
+				lua_pushstring (L, rule->name);
+				lua_createtable (L, commands->len, 0);
+
+				PTR_ARRAY_FOREACH (commands, j, io) {
+					lua_pushlstring (L, io->io.iov_base, io->io.iov_len);
+					lua_rawseti (L, -2, j + 1);
+				}
+
+				lua_settable (L, -3); /* ret[rule->name] = {raw_fuzzy1, ..., raw_fuzzyn} */
+			}
+
+			g_ptr_array_free (commands, TRUE);
+		}
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
+
 static gboolean
 fuzzy_add_handler (struct rspamd_http_connection_entry *conn_ent,
 	struct rspamd_http_message *msg, struct module_ctx *ctx)