Andrew Lewis
9 months ago
3 changed files with 354 additions and 48 deletions
@ -0,0 +1,302 @@ |
|||
--[[ |
|||
Copyright (c) 2025, Vsevolod Stakhov <vsevolod@rspamd.com> |
|||
|
|||
Licensed under the Apache License, Version 2.0 (the "License"); |
|||
you may not use this file except in compliance with the License. |
|||
You may obtain a copy of the License at |
|||
|
|||
http://www.apache.org/licenses/LICENSE-2.0 |
|||
|
|||
Unless required by applicable law or agreed to in writing, software |
|||
distributed under the License is distributed on an "AS IS" BASIS, |
|||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|||
See the License for the specific language governing permissions and |
|||
limitations under the License. |
|||
]]-- |
|||
|
|||
local E = {} |
|||
local N = 'contextal' |
|||
|
|||
if confighelp then |
|||
return |
|||
end |
|||
|
|||
local opts = rspamd_config:get_all_opt(N) |
|||
if not opts then |
|||
return |
|||
end |
|||
|
|||
local lua_redis = require "lua_redis" |
|||
local lua_util = require "lua_util" |
|||
local rspamd_http = require "rspamd_http" |
|||
local rspamd_logger = require "rspamd_logger" |
|||
local rspamd_util = require "rspamd_util" |
|||
local ucl = require "ucl" |
|||
|
|||
local contextal_actions = { |
|||
'ALERT', |
|||
'ALLOW', |
|||
'BLOCK', |
|||
'QUARANTINE', |
|||
'SPAM', |
|||
} |
|||
|
|||
local settings = { |
|||
action_symbol_prefix = 'CONTEXTAL_ACTION', |
|||
base_url = 'http://localhost:8080', |
|||
cache_ttl = 3600, |
|||
custom_actions = {}, |
|||
http_timeout = 2, |
|||
key_prefix = 'CXAL', |
|||
request_ttl = 4, |
|||
submission_symbol = 'CONTEXTAL_SUBMIT', |
|||
} |
|||
|
|||
local static_boundary = rspamd_util.random_hex(32) |
|||
|
|||
local function cache_key(task) |
|||
return string.format('%s_%s', settings.key_prefix, task:get_digest()) |
|||
end |
|||
|
|||
local function process_actions(task, obj, is_cached) |
|||
for _, match in ipairs((obj[1] or E).actions) do |
|||
local act = match.action |
|||
local scenario = match.scenario |
|||
if not (act and scenario) then |
|||
rspamd_logger.err(task, 'bad result: %s', match) |
|||
elseif contextal_actions[act] then |
|||
task:insert_result(settings.action_symbol_prefix .. '_' .. act, 1.0, scenario) |
|||
else |
|||
rspamd_logger.err(task, 'unknown action: %s', act) |
|||
end |
|||
end |
|||
|
|||
if not redis_params or is_cached then return end |
|||
|
|||
local cache_obj |
|||
if (obj[1] or E).actions then |
|||
cache_obj = {[1] = {["actions"] = obj[1].actions}} |
|||
elseif (obj[1] or E).work_id then |
|||
cache_obj = {[1] = {["work_id"] = obj[1].work_id}} |
|||
else |
|||
rspamd_logger.err(task, 'bad result: %s', obj) |
|||
return |
|||
end |
|||
|
|||
local key = cache_key(task) |
|||
local ret = lua_redis.redis_make_request(task, |
|||
redis_params, -- connect params |
|||
key, -- hash key |
|||
true, -- is write |
|||
redis_get_cb, --callback |
|||
'SET', -- command |
|||
{ key, ucl.to_format(cache_obj, 'json-compact') } -- arguments |
|||
) |
|||
|
|||
if not ret then |
|||
rspamd_logger.err(task, 'cannot make redis request to cache result') |
|||
return |
|||
end |
|||
end |
|||
|
|||
local function process_cached(task, txt) |
|||
local parser = ucl.parser() |
|||
local _, err = parser:parse_string(txt) |
|||
if err then |
|||
rspamd_logger.err(task, 'cannot parse JSON (cached): %s', err) |
|||
return |
|||
end |
|||
local obj = parser:get_object() |
|||
if (obj[1] or E).actions then |
|||
task:disable_symbol(settings.action_symbol_prefix) |
|||
return process_actions(task, obj, true) |
|||
elseif (obj[1] or E).work_id then |
|||
task:get_mempool():set_variable('contextal_work_id', obj.work_id) |
|||
else |
|||
rspamd_logger.err(task, 'bad result (cached): %s', obj) |
|||
end |
|||
end |
|||
|
|||
local function submit(task) |
|||
|
|||
local function http_callback(err, code, body, hdrs) |
|||
if err then |
|||
rspamd_logger.err(task, 'http error: %s', err) |
|||
return |
|||
end |
|||
if code ~= 201 then |
|||
rspamd_logger.err(task, 'bad http code: %s', code) |
|||
return |
|||
end |
|||
local parser = ucl.parser() |
|||
local _, err = parser:parse_string(body) |
|||
if err then |
|||
rspamd_logger.err(task, 'cannot parse JSON: %s', err) |
|||
return |
|||
end |
|||
local obj = parser:get_object() |
|||
local work_id = obj.work_id |
|||
if work_id then |
|||
task:get_mempool():set_variable('contextal_work_id', work_id) |
|||
end |
|||
task:insert_result(settings.submission_symbol, 1.0, |
|||
string.format('work_id=%s', work_id or 'nil')) |
|||
end |
|||
|
|||
local req = { |
|||
object_data = {['data'] = task:get_content()}, |
|||
} |
|||
if settings.request_ttl then |
|||
req.ttl = {['data'] = tostring(settings.request_ttl)} |
|||
end |
|||
if settings.max_recursion then |
|||
req.maxrec = {['data'] = tostring(settings.max_recursion)} |
|||
end |
|||
rspamd_http.request({ |
|||
task = task, |
|||
url = settings.submit_url, |
|||
body = lua_util.table_to_multipart_body(req, static_boundary), |
|||
callback = http_callback, |
|||
headers = { |
|||
['Content-Type'] = string.format('multipart/form-data; boundary="%s"', static_boundary) |
|||
}, |
|||
timeout = settings.http_timeout, |
|||
gzip = settings.gzip, |
|||
keepalive = settings.keepalive, |
|||
no_ssl_verify = settings.no_ssl_verify, |
|||
}) |
|||
end |
|||
|
|||
local function submit_cb(task) |
|||
if redis_params then |
|||
|
|||
local function redis_get_cb(err, data) |
|||
if err then |
|||
rspamd_logger.err(task, 'error querying redis: %s', err) |
|||
return |
|||
end |
|||
if type(data) == 'userdata' then |
|||
return submit(task) |
|||
end |
|||
process_cached(task, data) |
|||
end |
|||
|
|||
local key = cache_key(task) |
|||
local ret = lua_redis.redis_make_request(task, |
|||
redis_params, -- connect params |
|||
key, -- hash key |
|||
false, -- is write |
|||
redis_get_cb, --callback |
|||
'GET', -- command |
|||
{ key } -- arguments |
|||
) |
|||
|
|||
if not ret then |
|||
rspamd_logger.err(task, 'cannot make redis request to check results') |
|||
return |
|||
end |
|||
|
|||
else |
|||
return submit(task) |
|||
end |
|||
end |
|||
|
|||
local function action_cb(task) |
|||
local work_id = task:get_mempool():get_variable('contextal_work_id', 'string') |
|||
if not work_id then |
|||
rspamd_logger.err(task, 'no work id found in mempool') |
|||
return |
|||
end |
|||
|
|||
local function http_callback(err, code, body, hdrs) |
|||
if err then |
|||
rspamd_logger.err(task, 'http error: %s', err) |
|||
return |
|||
end |
|||
if code ~= 200 then |
|||
rspamd_logger.err(task, 'bad http code: %s', code) |
|||
return |
|||
end |
|||
local parser = ucl.parser() |
|||
local _, err = parser:parse_string(body) |
|||
if err then |
|||
rspamd_logger.err(task, 'cannot parse JSON: %s', err) |
|||
return |
|||
end |
|||
local obj = parser:get_object() |
|||
if (obj[1] or E).actions then |
|||
return process_actions(task, obj, false) |
|||
end |
|||
end |
|||
|
|||
rspamd_http.request({ |
|||
task = task, |
|||
url = settings.actions_url .. work_id, |
|||
callback = http_callback, |
|||
timeout = settings.http_timeout, |
|||
gzip = settings.gzip, |
|||
keepalive = settings.keepalive, |
|||
no_ssl_verify = settings.no_ssl_verify, |
|||
}) |
|||
end |
|||
|
|||
local function set_url_path(base, path) |
|||
local ts = base:sub(#base) == '/' and '' or '/' |
|||
return base .. ts .. path |
|||
end |
|||
|
|||
local opts = rspamd_config:get_all_opt(N) |
|||
if not opts then return end |
|||
|
|||
settings = lua_util.override_defaults(settings, opts) |
|||
|
|||
contextal_actions = lua_util.list_to_hash(contextal_actions) |
|||
for _, k in ipairs(settings.custom_actions) do |
|||
contextal_actions[k] = true |
|||
end |
|||
|
|||
if not settings.base_url then |
|||
if not (settings.submit_url and settings.actions_url) then |
|||
rspamd_logger.err(rspamd_config, 'no URL configured for contextal') |
|||
lua_util.disable_module(N, 'config') |
|||
return |
|||
end |
|||
else |
|||
if not settings.submit_url then |
|||
settings.submit_url = set_url_path(settings.base_url, 'api/v1/submit') |
|||
end |
|||
if not settings.actions_url then |
|||
settings.actions_url = set_url_path(settings.base_url, 'api/v1/actions/') |
|||
end |
|||
end |
|||
|
|||
redis_params = lua_redis.parse_redis_server(N) |
|||
if redis_params then |
|||
lua_redis.register_prefix(settings.key_prefix .. '_*', N, |
|||
'Cache for contextal plugin') |
|||
end |
|||
|
|||
rspamd_config:register_symbol({ |
|||
name = settings.submission_symbol, |
|||
priority = lua_util.symbols_priorities.top, |
|||
type = 'prefilter', |
|||
group = N, |
|||
callback = submit_cb |
|||
}) |
|||
|
|||
local id = rspamd_config:register_symbol({ |
|||
name = settings.action_symbol_prefix, |
|||
type = 'postfilter', |
|||
priority = lua_util.symbols_priorities.high - 1, |
|||
group = N, |
|||
callback = action_cb |
|||
}) |
|||
|
|||
for k in pairs(contextal_actions) do |
|||
rspamd_config:register_symbol({ |
|||
name = settings.action_symbol_prefix .. '_' .. k, |
|||
parent = id, |
|||
type = 'virtual', |
|||
group = N, |
|||
}) |
|||
end |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue