Add DKIM support to RBL module

1 files changed, 48 insertions, 0 deletions

diff --git a/src/plugins/lua/rbl.lua b/src/plugins/lua/rbl.lua
index 32e3af27a..a2dc7abbc 100644
--- a/src/plugins/lua/rbl.lua
+++ b/src/plugins/lua/rbl.lua
@@ -35,6 +35,11 @@ local private_ips = nil
 
 local rspamd_logger = require 'rspamd_logger'
 local rspamd_ip = require 'rspamd_ip'
+local rspamd_url = require 'rspamd_url'
+
+local symbols = {
+  dkim_allow_symbol = 'R_DKIM_ALLOW',
+}
 
 local function validate_dns(lstr)
   if lstr:match('%.%.') then
@@ -168,6 +173,35 @@ local function rbl_cb (task)
 	end)()
       end
 
+      if rbl['dkim'] then
+	(function()
+	  if notgot['dkim'] then
+	    return
+	  end
+	  if not havegot['dkim'] then
+            local das = task:get_symbol(symbols['dkim_allow_symbol'])
+            if das and das[1] and das[1]['options'] and das[1]['options'][0] then
+              havegot['dkim'] = das[1]['options']
+	    else
+	      notgot['dkim'] = true
+	      return
+	    end
+	  end
+          for _, d in pairs(havegot['dkim']) do
+            if rbl['dkim_domainonly'] then
+              local url_from = rspamd_url.create(task:get_mempool(), d)
+              if url_from then
+                d = url_from:get_tld()
+              else
+                return
+              end
+            end
+	    task:get_resolver():resolve_a(task:get_session(), task:get_mempool(),
+	      d .. '.' .. rbl['rbl'], rbl_dns_cb, k)
+          end
+	end)()
+      end
+
       if rbl['emails'] then
         (function()
           if notgot['emails'] then
@@ -290,6 +324,8 @@ if type(rspamd_config.get_api_version) ~= 'nil' then
     rspamd_config:register_module_option('rbl', 'default_from', 'string')
     rspamd_config:register_module_option('rbl', 'default_rdns', 'string')
     rspamd_config:register_module_option('rbl', 'default_helo', 'string')
+    rspamd_config:register_module_option('rbl', 'default_dkim', 'string')
+    rspamd_config:register_module_option('rbl', 'default_dkim_domainonly', 'string')
     rspamd_config:register_module_option('rbl', 'default_unknown', 'string')
     rspamd_config:register_module_option('rbl', 'default_exclude_users', 'string')
     rspamd_config:register_module_option('rbl', 'default_exclude_private_ips', 'string')
@@ -318,6 +354,8 @@ default_defaults = {
   ['default_unknown'] = {[1] = false, [2] = 'unknown'},
   ['default_rdns'] = {[1] = false, [2] = 'rdns'},
   ['default_helo'] = {[1] = false, [2] = 'helo'},
+  ['default_dkim'] = {[1] = false, [2] = 'dkim'},
+  ['default_dkim_domainonly'] = {[1] = true, [2] = 'dkim_domainonly'},
   ['default_emails'] = {[1] = false, [2] = 'emails'},
   ['default_exclude_users'] = {[1] = false, [2] = 'exclude_users'},
   ['default_exclude_private_ips'] = {[1] = true, [2] = 'exclude_private_ips'},
@@ -341,6 +379,7 @@ end
 
 local white_symbols = {}
 local black_symbols = {}
+local need_dkim = false
 
 local id = rspamd_config:register_callback_symbol_priority(1.0, 0, rbl_cb)
 
@@ -354,6 +393,9 @@ for key,rbl in pairs(opts['rbls']) do
     for s,_ in pairs(rbl['returncodes']) do
       if type(rspamd_config.get_api_version) ~= 'nil' then
         rspamd_config:register_virtual_symbol(s, 1, id)
+        if rbl['dkim'] then
+          need_dkim = true
+        end
         if(rbl['is_whitelist']) then
           if type(rbl['whitelist_exception']) == 'string' then
             if (rbl['whitelist_exception'] ~= s) then
@@ -388,6 +430,9 @@ for key,rbl in pairs(opts['rbls']) do
   end
   if type(rspamd_config.get_api_version) ~= 'nil' and rbl['symbol'] then
     rspamd_config:register_virtual_symbol(rbl['symbol'], 1, id)
+    if rbl['dkim'] then
+      need_dkim = true
+    end
     if(rbl['is_whitelist']) then
           if type(rbl['whitelist_exception']) == 'string' then
             if (rbl['whitelist_exception'] ~= rbl['symbol']) then
@@ -422,3 +467,6 @@ for _, w in pairs(white_symbols) do
     rspamd_config:add_composite(csymbol, w .. ' & ' .. b)
   end
 end
+if need_dkim then
+  rspamd_config:register_dependency(id, symbols['dkim_allow_symbol'])
+end