* Make trie plugin and fix trie lua API

1 files changed, 102 insertions, 0 deletions

diff --git a/src/plugins/lua/trie.lua b/src/plugins/lua/trie.lua
new file mode 100644
index 000000000..d4bafe943
--- /dev/null
+++ b/src/plugins/lua/trie.lua
@@ -0,0 +1,102 @@
+-- Trie is rspamd module designed to define and operate with suffix trie
+
+local tries = {}
+
+local function split(str, delim, maxNb)
+	-- Eliminate bad cases...
+	if string.find(str, delim) == nil then
+		return { str }
+	end
+	if maxNb == nil or maxNb < 1 then
+		maxNb = 0    -- No limit
+	end
+	local result = {}
+	local pat = "(.-)" .. delim .. "()"
+	local nb = 0
+	local lastPos
+	for part, pos in string.gfind(str, pat) do
+		nb = nb + 1
+		result[nb] = part
+		lastPos = pos
+		if nb == maxNb then break end
+	end
+	-- Handle the last field
+	if nb ~= maxNb then
+		result[nb + 1] = string.sub(str, lastPos)
+	end
+	return result
+end
+
+local function add_trie(params)
+	local symbol = params[1]
+	
+	file = io.open(params[2])
+	if file then
+		local trie = {}
+		trie['trie'] = rspamd_trie:create(true)
+		num = 0
+		for line in file:lines() do
+			trie['trie']:add_pattern(line, num)
+			num = num + 1
+		end
+		
+		if type(rspamd_config.get_api_version) ~= 'nil' then
+			rspamd_config:register_virtual_symbol(symbol, 1.0)
+		end
+		file:close()
+		trie['symbol'] = symbol
+		table.insert(tries, trie)
+	else
+		local patterns = split(params[2], ',')
+		local trie = {}
+		trie['trie'] = rspamd_trie:create(true)
+		print (type(trie['trie']))
+		for num,pattern in ipairs(patterns) do
+			trie['trie']:add_pattern(pattern, num)
+		end
+		if type(rspamd_config.get_api_version) ~= 'nil' then
+			rspamd_config:register_virtual_symbol(symbol, 1.0)
+		end
+		trie['symbol'] = symbol
+		table.insert(tries, trie)
+	end
+end
+
+function check_trie(task)
+	for _,trie in ipairs(tries) do
+		print (type(trie['trie']))
+		if trie['trie']:search_task(task) then
+			task:insert_result(trie['symbol'], 1)
+		end
+	end
+end
+
+-- Registration
+if type(rspamd_config.get_api_version) ~= 'nil' then
+	if rspamd_config:get_api_version() >= 1 then
+		rspamd_config:register_module_option('trie', 'rule', 'string')
+	end
+end
+
+local opts =  rspamd_config:get_all_opt('trie')
+if opts then
+	local strrules = opts['rule']
+	if strrules then
+		if type(strrules) == 'table' then 
+			for _,value in ipairs(strrules) do
+				local params = split(value, ':')
+				add_trie(params)
+			end
+		elseif type(strrules) == 'string' then
+			local params = split(strrules, ':')
+			add_trie (params)
+		end
+	end
+	if table.maxn(tries) then
+		if type(rspamd_config.get_api_version) ~= 'nil' then
+			rspamd_config:register_callback_symbol('TRIE', 1.0, 'check_trie')
+		else
+			rspamd_config:register_symbol('TRIE', 1.0, 'check_trie')
+		end
+	end
+end