From afb5bf3cbf718f231e344947ce5fd3d660961139 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Tue, 19 Jun 2018 18:50:45 +0100 Subject: [PATCH] [Feature] Add method to get statistical tokens in Lua API --- src/libstat/stat_api.h | 9 +++ src/libstat/stat_internal.h | 9 --- src/lua/lua_task.c | 129 +++++++++++++++++++++++++++++++++++- 3 files changed, 137 insertions(+), 10 deletions(-) diff --git a/src/libstat/stat_api.h b/src/libstat/stat_api.h index 1ade4ccd5..84db8ee01 100644 --- a/src/libstat/stat_api.h +++ b/src/libstat/stat_api.h @@ -39,6 +39,15 @@ typedef struct rspamd_stat_token_s { guint flags; } rspamd_stat_token_t; +typedef struct token_node_s { + guint64 data; + guint window_idx; + guint flags; + rspamd_stat_token_t *t1; + rspamd_stat_token_t *t2; + gdouble values[]; +} rspamd_token_t; + struct rspamd_stat_ctx; /** diff --git a/src/libstat/stat_internal.h b/src/libstat/stat_internal.h index cab185f7c..44f48ae5a 100644 --- a/src/libstat/stat_internal.h +++ b/src/libstat/stat_internal.h @@ -51,15 +51,6 @@ struct rspamd_statfile { gpointer bkcf; }; -typedef struct token_node_s { - guint64 data; - guint window_idx; - guint flags; - rspamd_stat_token_t *t1; - rspamd_stat_token_t *t2; - gdouble values[]; -} rspamd_token_t; - struct rspamd_stat_async_elt; typedef void (*rspamd_stat_async_handler)(struct rspamd_stat_async_elt *elt, diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c index 46f480c02..94d4c05dc 100644 --- a/src/lua/lua_task.c +++ b/src/lua/lua_task.c @@ -21,8 +21,9 @@ #include "unix-std.h" #include "libmime/smtp_parsers.h" #include "libserver/mempool_vars_internal.h" +#include "libserver/task.h" +#include "libstat/stat_api.h" #include -#include /*** * @module rspamd_task @@ -864,6 +865,25 @@ LUA_FUNCTION_DEF (task, disable_action); */ LUA_FUNCTION_DEF (task, get_newlines_type); +/*** + * @method task:get_stat_tokens() + * Returns list of tables the statistical tokens: + * - `data`: 64 bit number encoded as a string + * - `t1`: the first token (if any) + * - `t2`: the second token (if any) + * - `win`: window index + * - `flag`: table of strings: + * - `text`: text token + * - `meta`: meta token + * - `lua`: lua meta token + * - `exception`: exception + * - `subject`: subject token + * - `unigram`: unigram token + * + * @return {table of tables} + */ +LUA_FUNCTION_DEF (task, get_stat_tokens); + static const struct luaL_reg tasklib_f[] = { LUA_INTERFACE_DEF (task, load_from_file), {NULL, NULL} @@ -959,6 +979,7 @@ static const struct luaL_reg tasklib_m[] = { LUA_INTERFACE_DEF (task, headers_foreach), LUA_INTERFACE_DEF (task, disable_action), LUA_INTERFACE_DEF (task, get_newlines_type), + LUA_INTERFACE_DEF (task, get_stat_tokens), {"__tostring", rspamd_lua_class_tostring}, {NULL, NULL} }; @@ -4392,6 +4413,112 @@ lua_task_get_newlines_type (lua_State *L) return 1; } +static void +lua_push_stat_token (lua_State *L, rspamd_token_t *tok) +{ + gchar numbuf[64]; + + /* Table values + * - `data`: 64 bit number encoded as a string + * - `t1`: the first token (if any) + * - `t2`: the second token (if any) + * - `win`: window index + * - `flag`: table of strings: + * - `text`: text token + * - `meta`: meta token + * - `lua`: lua meta token + * - `exception`: exception + * - `subject`: subject token + * - `unigram`: unigram token + */ + lua_createtable (L, 0, 5); + + rspamd_snprintf (numbuf, sizeof (numbuf), "%uL", tok->data); + lua_pushstring (L, "data"); + lua_pushstring (L, numbuf); + lua_settable (L, -3); + + if (tok->t1) { + lua_pushstring (L, "t1"); + lua_pushlstring (L, tok->t1->begin, tok->t1->len); + lua_settable (L, -3); + } + + if (tok->t2) { + lua_pushstring (L, "t2"); + lua_pushlstring (L, tok->t2->begin, tok->t2->len); + lua_settable (L, -3); + } + + lua_pushstring (L, "win"); + lua_pushnumber (L, tok->window_idx); + lua_settable (L, -3); + + lua_pushstring (L, "flags"); + lua_createtable (L, 0, 5); + + /* Flags */ + { + if (tok->flags & RSPAMD_STAT_TOKEN_FLAG_TEXT) { + lua_pushstring (L, "text"); + lua_pushboolean (L, true); + lua_settable (L, -3); + } + if (tok->flags & RSPAMD_STAT_TOKEN_FLAG_META) { + lua_pushstring (L, "meta"); + lua_pushboolean (L, true); + lua_settable (L, -3); + } + if (tok->flags & RSPAMD_STAT_TOKEN_FLAG_LUA_META) { + lua_pushstring (L, "lua"); + lua_pushboolean (L, true); + lua_settable (L, -3); + } + if (tok->flags & RSPAMD_STAT_TOKEN_FLAG_EXCEPTION) { + lua_pushstring (L, "exception"); + lua_pushboolean (L, true); + lua_settable (L, -3); + } + if (tok->flags & RSPAMD_STAT_TOKEN_FLAG_SUBJECT) { + lua_pushstring (L, "subject"); + lua_pushboolean (L, true); + lua_settable (L, -3); + } + } + lua_settable (L, -3); +} + +static gint +lua_task_get_stat_tokens (lua_State *L) +{ + struct rspamd_task *task = lua_check_task (L, 1); + guint i; + rspamd_token_t *tok; + + if (task) { + if (!task->tokens) { + rspamd_stat_process_tokenize (NULL, task); + } + + if (!task->tokens) { + lua_pushnil (L); + } + else { + lua_createtable (L, task->tokens->len, 0); + + PTR_ARRAY_FOREACH (task->tokens, i, tok) { + lua_push_stat_token (L, tok); + lua_rawseti (L, -2, i + 1); + } + } + } + else { + return luaL_error (L, "invalid arguments"); + } + + return 1; +} + static gint lua_task_set_metric_subject (lua_State *L) { -- 2.39.5