[Feature] Fuzzy_check: Add gen_hashes command

diff --git a/src/plugins/fuzzy_check.c b/src/plugins/fuzzy_check.c
index 36e093836962957a282c779d070f08118f6994da..e33735c5b24f9c52e72c4d9da48e86a4c6ace279 100644 (file)
--- 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)