local rspamd_logger = require "rspamd_logger"
local rspamd_fann = require "rspamd_fann"
local rspamd_util = require "rspamd_util"
-local fann_symbol_spam = 'FANNR_SPAM'
-local fann_symbol_ham = 'FANNR_HAM'
local rspamd_redis = require "lua_redis"
local fun = require "fun"
local meta_functions = require "meta_functions"
+
-- Module vars
+local default_options = {
+ train = {
+ max_trains = 1000,
+ max_epoch = 1000,
+ max_usages = 10,
+ use_settings = false,
+ watch_interval = 60.0,
+ mse = 0.001,
+ autotrain = true,
+ },
+ nlayers = 4,
+ lock_expire = 600,
+ learning_spawned = false,
+ ann_expire = 60 * 60 * 24 * 2, -- 2 days
+ symbol_spam = 'FANNR_SPAM',
+ symbol_ham = 'FANNR_HAM',
+}
+
+local settings = {
+ rules = {
+ }
+}
+
-- ANNs indexed by settings id
local fanns = {
- ['0'] = {
- version = 0,
- }
}
local opts = rspamd_config:get_all_opt("fann_redis")
local redis_save_unlock_sha = nil
local redis_params
-redis_params = rspamd_parse_redis_server('fann_redis')
-
-local fann_prefix = 'RFANN'
-local max_trains = 1000
-local max_epoch = 1000
-local max_usages = 10
-local use_settings = false
-local watch_interval = 60.0
-local mse = 0.0001
-local nlayers = 4
-local lock_expire = 600
-local learning_spawned = false
-local ann_expire = 60 * 60 * 24 * 2 -- 2 days
local function load_scripts(cfg, ev_base, on_load_cb)
local function can_train_sha_cb(err, data)
)
end
-local function gen_fann_prefix(id)
+local function gen_fann_prefix(rule, id)
if id then
- return fann_prefix .. rspamd_config:get_symbols_cksum():hex() .. id,id
+ return rule.prefix .. rspamd_config:get_symbols_cksum():hex() .. id,
+ rule.prefix .. id
else
- return fann_prefix .. rspamd_config:get_symbols_cksum():hex(), nil
+ return rule.prefix .. rspamd_config:get_symbols_cksum():hex(), nil
end
end
-local function is_fann_valid(prefix, ann)
+local function is_fann_valid(rule, prefix, ann)
if ann then
local n = rspamd_config:get_symbols_count() + rspamd_count_metatokens()
end
local layers = ann:get_layers()
- if not layers or #layers ~= nlayers then
+ if not layers or #layers ~= rule.nlayers then
rspamd_logger.infox(rspamd_config, 'ANN %s has incorrect number of layers: %s',
prefix, #layers)
return false
end
local function fann_scores_filter(task)
- local id = '0'
- if use_settings then
- local sid = task:get_settings_id()
- if sid then
- id = tostring(sid)
- end
- end
+ for _,rule in settings.rules do
+ local id = rule.prefix .. '0'
+ if rule.use_settings then
+ local sid = task:get_settings_id()
+ if sid then
+ id = rule.prefix .. tostring(sid)
+ end
+ end
- if fanns[id].fann then
- local fann_data = task:get_symbols_tokens()
- local mt = meta_functions.rspamd_gen_metatokens(task)
- -- Add filtered meta tokens
- fun.each(function(e) table.insert(fann_data, e) end, mt)
-
- local out = fanns[id].fann:test(fann_data)
- local symscore = string.format('%.3f', out[1])
- rspamd_logger.infox(task, 'fann score: %s', symscore)
-
- if out[1] > 0 then
- local result = rspamd_util.normalize_prob(out[1] / 2.0, 0)
- task:insert_result(fann_symbol_spam, result, symscore, id)
- else
- local result = rspamd_util.normalize_prob((-out[1]) / 2.0, 0)
- task:insert_result(fann_symbol_ham, result, symscore, id)
+ if fanns[id].fann then
+ local fann_data = task:get_symbols_tokens()
+ local mt = meta_functions.rspamd_gen_metatokens(task)
+ -- Add filtered meta tokens
+ fun.each(function(e) table.insert(fann_data, e) end, mt)
+
+ local out = fanns[id].fann:test(fann_data)
+ local symscore = string.format('%.3f', out[1])
+ rspamd_logger.infox(task, 'fann score: %s', symscore)
+
+ if out[1] > 0 then
+ local result = rspamd_util.normalize_prob(out[1] / 2.0, 0)
+ task:insert_result(rule.symbol_spam, result, symscore, id)
+ else
+ local result = rspamd_util.normalize_prob((-out[1]) / 2.0, 0)
+ task:insert_result(rule.symbol_ham, result, symscore, id)
+ end
end
end
end
-local function create_train_fann(n, id)
- id = tostring(id)
- local prefix = gen_fann_prefix(id)
+local function create_train_fann(rule, n, id)
+ id = rule.prefix .. tostring(id)
+ local prefix = gen_fann_prefix(rule, id)
if not fanns[id] then
fanns[id] = {}
end
-
+ -- Fix that for flexibe layers number
if fanns[id].fann then
- if n ~= fanns[id].fann:get_inputs() or
+ if n ~= fanns[id].fann:get_inputs() or --
(fanns[id].fann_train and n ~= fanns[id].fann_train:get_inputs()) then
- rspamd_logger.infox(rspamd_config, 'recreate ANN %s as it has a wrong number of inputs, version %s', prefix,
+ rspamd_logger.infox(rspamd_config,
+ 'recreate ANN %s as it has a wrong number of inputs, version %s',
+ prefix,
fanns[id].version)
- fanns[id].fann_train = rspamd_fann.create(nlayers, n, n / 2, n / 4, 1)
+ fanns[id].fann_train = rspamd_fann.create(rule.nlayers, n, n / 2, n / 4, 1)
fanns[id].fann = nil
- elseif fanns[id].version % max_usages == 0 then
+ elseif fanns[id].version % rule.max_usages == 0 then
-- Forget last fann
rspamd_logger.infox(rspamd_config, 'recreate ANN %s, version %s', prefix,
fanns[id].version)
- fanns[id].fann_train = rspamd_fann.create(nlayers, n, n / 2, n / 4, 1)
+ fanns[id].fann_train = rspamd_fann.create(rule.nlayers, n, n / 2, n / 4, 1)
else
fanns[id].fann_train = fanns[id].fann
end
else
- fanns[id].fann_train = rspamd_fann.create(nlayers, n, n / 2, n / 4, 1)
+ fanns[id].fann_train = rspamd_fann.create(rule.nlayers, n, n / 2, n / 4, 1)
fanns[id].version = 0
end
end
-local function load_or_invalidate_fann(data, id, ev_base)
+local function load_or_invalidate_fann(rule, data, id, ev_base)
local ver = data[2]
- local prefix = gen_fann_prefix(id)
+ local prefix = gen_fann_prefix(rule, id)
if not ver or not tonumber(ver) then
rspamd_logger.errx(rspamd_config, 'cannot get version for ANN: %s', prefix)
ann = rspamd_fann.load_data(ann_data)
end
- if is_fann_valid(prefix, ann) then
+ if is_fann_valid(rule, prefix, ann) then
fanns[id].fann = ann
rspamd_logger.infox(rspamd_config, 'loaded ANN %s version %s from redis',
prefix, ver)
rspamd_logger.infox(rspamd_config, 'invalidate ANN %s', prefix)
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_invalidate_cb, --callback
end
end
-local function fann_train_callback(task, score, required_score, id)
- local train_opts = opts['train']
- local fname,suffix = gen_fann_prefix(id)
+local function fann_train_callback(rule, task, score, required_score, id)
+ local train_opts = rule['train']
+ local fname,suffix = gen_fann_prefix(rule, id)
local learn_spam, learn_ham
local str = rspamd_util.zstd_compress(table.concat(fann_data, ';'))
rspamd_redis.redis_make_request(task,
- redis_params,
+ rule.redis,
nil,
true, -- is write
learn_vec_cb, --callback
end
rspamd_redis.rspamd_redis_make_request(task,
- redis_params,
+ rule.redis,
nil,
true, -- is write
can_train_cb, --callback
'EVALSHA', -- command
- {redis_can_train_sha, '4', gen_fann_prefix(nil),
+ {redis_can_train_sha, '4', gen_fann_prefix(rule, nil),
suffix, k, tostring(max_trains)} -- arguments
)
end
end
-local function train_fann(_, ev_base, elt)
+local function train_fann(rule, _, ev_base, elt)
local spam_elts = {}
local ham_elts = {}
elt = tostring(elt)
- local prefix = gen_fann_prefix(elt)
+ local prefix = gen_fann_prefix(rule, elt)
local function redis_unlock_cb(err)
if err then
prefix, err)
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
redis_unlock_cb, --callback
end
local function ann_trained(errcode, errmsg, train_mse)
- learning_spawned = false
+ rule.learning_spawned = false
if errcode ~= 0 then
rspamd_logger.errx(rspamd_config, 'cannot train ANN %s: %s',
prefix, errmsg)
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_unlock_cb, --callback
fanns[elt].fann_train = nil
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_save_cb, --callback
prefix, err)
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_unlock_cb, --callback
if not fanns[elt] or not fanns[elt].fann_train
or n ~= fanns[elt].fann_train:get_inputs() then
-- Create fann if it does not exist
- create_train_fann(n, elt)
+ create_train_fann(rule, n, elt)
end
if #inputs < max_trains / 2 then
rspamd_logger.infox(rspamd_config, 'invalidate ANN %s: training data is invalid', prefix)
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_invalidate_cb, --callback
{redis_locked_invalidate_sha, 1, prefix}
)
else
- learning_spawned = true
+ rule.learning_spawned = true
rspamd_logger.infox(rspamd_config, 'start learning ANN %s', prefix)
- fanns[elt].fann_train:train_threaded(inputs, outputs, ann_trained, ev_base,
- {max_epochs = max_epoch, desired_mse = mse})
+ fanns[elt].fann_train:train_threaded(inputs, outputs, ann_trained,
+ ev_base, {
+ max_epochs = rule.train.max_epoch,
+ desired_mse = rule.train.mse
+ })
end
end
end
prefix, err)
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_unlock_cb, --callback
end, data))
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
redis_ham_cb, --callback
-- Can train ANN
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
redis_spam_cb, --callback
prefix)
end
end
- if learning_spawned then
+ if rule.learning_spawned then
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_lock_extend_cb, --callback
rspamd_logger.infox(rspamd_config, 'do not learn ANN %s, locked by another process', prefix)
end
end
- if learning_spawned then
+ if rule.learning_spawned then
rspamd_logger.infox(rspamd_config, 'do not learn ANN %s, already learning another ANN', prefix)
return
end
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
true, -- is write
redis_lock_cb, --callback
'EVALSHA', -- command
{redis_maybe_lock_sha, '4', prefix, tostring(os.time()),
- tostring(lock_expire), rspamd_util.get_hostname()}
+ tostring(rule.lock_expire), rspamd_util.get_hostname()}
)
end
-local function maybe_train_fanns(cfg, ev_base)
+local function maybe_train_fanns(rule, cfg, ev_base)
local function members_cb(err, data)
if err then
rspamd_logger.errx(rspamd_config, 'cannot get FANNS list from redis: %s', err)
elseif type(data) == 'table' then
fun.each(function(elt)
elt = tostring(elt)
- local prefix = gen_fann_prefix(elt)
+ local prefix = gen_fann_prefix(rule, elt)
local redis_len_cb = function(_err, _data)
if _err then
- rspamd_logger.errx(rspamd_config, 'cannot get FANN trains %s from redis: %s', prefix, _err)
+ rspamd_logger.errx(rspamd_config,
+ 'cannot get FANN trains %s from redis: %s', prefix, _err)
elseif _data and type(_data) == 'number' or type(_data) == 'string' then
if tonumber(_data) and tonumber(_data) >= max_trains then
- rspamd_logger.infox(rspamd_config, 'need to learn ANN %s after %s learn vectors (%s required)',
+ rspamd_logger.infox(rspamd_config,
+ 'need to learn ANN %s after %s learn vectors (%s required)',
prefix, tonumber(_data), max_trains)
- train_fann(cfg, ev_base, elt)
+ train_fann(rule, cfg, ev_base, elt)
end
end
end
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
redis_len_cb, --callback
-- First we need to get all fanns stored in our Redis
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
members_cb, --callback
'SMEMBERS', -- command
- {gen_fann_prefix(nil)} -- arguments
+ {gen_fann_prefix(rule, nil)} -- arguments
)
- return watch_interval
+ return rule.watch_interval
end
-local function check_fanns(_, ev_base)
+local function check_fanns(rule, _, ev_base)
local function members_cb(err, data)
if err then
rspamd_logger.errx(rspamd_config, 'cannot get FANNS list from redis: %s', err)
load_scripts(rspamd_config, ev_base, nil)
end
elseif _data and type(_data) == 'table' then
- load_or_invalidate_fann(_data, elt, ev_base)
+ load_or_invalidate_fann(rule, _data, elt, ev_base)
end
end
end
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
redis_update_cb, --callback
'EVALSHA', -- command
- {redis_maybe_load_sha, 2, gen_fann_prefix(elt), tostring(local_ver)}
+ {redis_maybe_load_sha, 2, gen_fann_prefix(rule, elt), tostring(local_ver)}
)
end,
data)
-- First we need to get all fanns stored in our Redis
rspamd_redis.redis_make_request_taskless(ev_base,
rspamd_config,
- redis_params,
+ rule.redis,
nil,
false, -- is write
members_cb, --callback
'SMEMBERS', -- command
- {gen_fann_prefix(nil)} -- arguments
+ {gen_fann_prefix(rule, nil)} -- arguments
)
- return watch_interval
+ return rule.watch_interval
end
local function ann_push_vector(task)
local scores = task:get_metric_score()
local sid = task:get_settings_id()
- if use_settings then
- fann_train_callback(task, scores[1], scores[2], tostring(sid))
- else
- fann_train_callback(task, scores[1], scores[2], "1")
+
+ for _,rule in settings.rules do
+ if rule.use_settings then
+ fann_train_callback(rule, task, scores[1], scores[2], tostring(sid))
+ else
+ fann_train_callback(rule, task, scores[1], scores[2], "0")
+ end
end
end
+redis_params = rspamd_parse_redis_server('fann_redis')
+
-- Initialization part
if not (opts and type(opts) == 'table') or not redis_params then
rspamd_logger.infox(rspamd_config, 'Module is unconfigured')
'module is eventually disabled')
return
else
- use_settings = opts['use_settings']
- if opts['spam_symbol'] then
- fann_symbol_spam = opts['spam_symbol']
- end
- if opts['ham_symbol'] then
- fann_symbol_ham = opts['ham_symbol']
- end
- if opts['prefix'] then
- fann_prefix = opts['prefix']
- end
- if opts['lock_expire'] then
- lock_expire = tonumber(opts['lock_expire'])
+ local rules = opts['rules']
+
+ if not rules then
+ -- Use legacy configuration
+ rules = {}
+ rules['RFANN'] = opts
end
- rspamd_config:set_metric_symbol({
- name = fann_symbol_spam,
- score = 3.0,
- description = 'Neural network SPAM',
- group = 'fann'
- })
+
local id = rspamd_config:register_symbol({
- name = fann_symbol_spam,
+ name = 'FANN_CHECK',
type = 'postfilter,nostat',
priority = 6,
callback = fann_scores_filter
})
- rspamd_config:set_metric_symbol({
- name = fann_symbol_ham,
- score = -2.0,
- description = 'Neural network HAM',
- group = 'fann'
- })
- rspamd_config:register_symbol({
- name = fann_symbol_ham,
- type = 'virtual,nostat',
- parent = id
- })
- if opts['train'] then
- if opts['train']['max_train'] then
- max_trains = opts['train']['max_train']
- end
- if opts['train']['max_epoch'] then
- max_epoch = opts['train']['max_epoch']
- end
- if opts['train']['max_usages'] then
- max_usages = opts['train']['max_usages']
+
+ for k,r in rules do
+ rules[k] = default_options
+ rules[k]['redis'] = redis_params
+ local cur = rules[k]
+ -- Override defaults
+ for sk,v in r do
+ cur[sk] = v
end
- if opts['train']['mse'] then
- mse = opts['train']['mse']
+ if not cur.prefix then
+ cur.prefix = k
end
+ rspamd_config:set_metric_symbol({
+ name = cur.symbol_spam,
+ score = 3.0,
+ description = 'Neural network SPAM',
+ group = 'fann'
+ })
+
+ rspamd_config:set_metric_symbol({
+ name = cur.symbol_ham,
+ score = -2.0,
+ description = 'Neural network HAM',
+ group = 'fann'
+ })
rspamd_config:register_symbol({
- name = 'FANN_VECTOR_PUSH',
- type = 'postfilter,nostat',
- priority = 5,
- callback = ann_push_vector
+ name = cur.symbol_ham,
+ type = 'virtual,nostat',
+ parent = id
})
end
+
+ rspamd_config:register_symbol({
+ name = 'FANN_VECTOR_PUSH',
+ type = 'postfilter,nostat',
+ priority = 5,
+ callback = ann_push_vector
+ })
+
+ settings.rules = rules
+
-- Add training scripts
- rspamd_config:add_on_load(function(cfg, ev_base, worker)
- load_scripts(cfg, ev_base, check_fanns)
-
- if worker:get_name() == 'normal' then
- -- We also want to train neural nets when they have enough data
- rspamd_config:add_periodic(ev_base, 0.0,
- function(_cfg, _ev_base)
- return maybe_train_fanns(_cfg, _ev_base)
- end)
- end
- end)
+ for k,rule in settings.rules do
+ rspamd_config:add_on_load(function(cfg, ev_base, worker)
+ load_scripts(cfg, ev_base, function(cfg, ev_base)
+ check_fanns(rule, cfg, ev_base)
+ end)
+
+ if worker:get_name() == 'normal' then
+ -- We also want to train neural nets when they have enough data
+ rspamd_config:add_periodic(ev_base, 0.0,
+ function(_cfg, _ev_base)
+ return maybe_train_fanns(rule, _cfg, _ev_base)
+ end)
+ end
+ end)
+ end
end
projects / thirdparty / rspamd.git / commitdiff
