--[[
Copyright (c) 2000-2020 Synology Inc. All rights reserved.
]]--

--[[
Rspamd filter priority:
* prefilter for Rspamd:
  * MCP(10)
  * extract MCP(9)
  * AntiVirus whitelist/DKIM spam/DMARC quarantine(6)
  * AntiVirus(5)
  * Spam blacklist/Spam whitelist(3)
  * Disable Spam(2)
  * cache(1)
* prefilter for BitDefender:
  * MCP(10)
  * extract MCP(9)
  * AntiVirus whitelist/Spam blacklist/Spam whitelist/DKIM spam/DMARC quarantine(6)
  * AntiVirus(5)
  * Disable Spam/BitDefender Clean(2)
  * cache(1)
* filter: spam, SpamAssassin rules
* postfilter: header tagging + cache(10)
]] --

local rspamd_logger = require "rspamd_logger"
local cache         = require "synology/cache"
local constants     = require "synology/constants"
local synology_util = require "synology/synology_util"

local function spam_check(task, add, del)
  local spam_opt = rspamd_config:get_all_opt('spam')
  local spam_enabled = spam_opt['enabled']
  local spam_learn_enabled = spam_opt['learn_enabled']
  local use_bitdefender = spam_opt['use_bitdefender']

  if not spam_enabled then
    return
  end

  local action = task:get_metric_action('default')
  local scores = task:get_metric_score('default')
  -- score = { [1] = current score, [2] = required score }
  local current_score = scores[1]
  local required_score = scores[2]

  local symbols = task:get_symbols_all()
  -- symbol: { [group] = group, [options] = {[1] = option1, [2] = option2}, [name] = name, [score] = score }
  -- whitelist
  for _, symbol in ipairs(symbols) do
    if synology_util.is_whitelist_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.NO
      if use_bitdefender then
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_status(symbol)
      else
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_ham_status(symbol, required_score)
      end
      rspamd_logger.infox(task, "Hit whitelist; symbol %s", symbol)
      return
    end
  end
  -- blacklist
  for _, symbol in ipairs(symbols) do
    if synology_util.is_blacklist_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.YES
      if use_bitdefender then
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_status(symbol)
      else
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_spam_status(symbol, required_score)
      end
      rspamd_logger.infox(task, "Hit blacklist; symbol %s", symbol)
      return
    end
  end
  -- DKIM fail
  for _, symbol in ipairs(symbols) do
    if synology_util.is_dkim_spam_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.YES
      if use_bitdefender then
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_status(symbol)
      else
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_spam_status(symbol, required_score)
      end
      rspamd_logger.infox(task, "Spammed by DKIM; symbol %s", symbol)
      return
    end
  end
  -- DMARC quarantine
  for _, symbol in ipairs(symbols) do
    if synology_util.is_dmarc_quarantine_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.YES
      if use_bitdefender then
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_status(symbol)
      else
        add[constants.SPAM_STATUS_HEADER] = synology_util.generate_spam_status(symbol, required_score)
      end
      rspamd_logger.infox(task, "Quarantined by DMARC; symbol %s", symbol)
      return
    end
  end
  -- BitDefender Spam
  for _, symbol in ipairs(symbols) do
    if synology_util.is_bitdefender_spam_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.YES
      add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_spam_status(symbol)
      rspamd_logger.infox(task, "Spammed by BitDefender; symbol %s", symbol)
      rspamd_logger.errx(task, "msgid=<%1>: Spam, %2", task:get_message_id(), add[constants.SPAM_STATUS_HEADER])
      return
    end
  end
  -- BitDefender skipped threat type
  for _, symbol in ipairs(symbols) do
    if synology_util.is_bitdefender_skip_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.NO
      add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_skip_status(symbol)
      rspamd_logger.infox(task, "Skipped by BitDefender; symbol %s", symbol)
      return
    end
  end
  -- BitDefender clean
  for _, symbol in ipairs(symbols) do
    if synology_util.is_bitdefender_clean_symbol(symbol) then
      add[constants.SPAM_FLAG_HEADER] = constants.NO
      add[constants.SPAM_STATUS_HEADER] = synology_util.generate_bitdefender_clean_status(symbol)
      rspamd_logger.infox(task, "Cleaned by BitDefender; symbol %s", symbol)
      return
    end
  end

  local spam_status
  local cache_hit = false
  if cache.enabled then
    local cache_symbol = task:get_symbol(cache.settings.symbol)
    if cache_symbol ~= nil then
      rspamd_logger.infox(task, "Cache hit; getting spam status from cache")
      cache_hit = true

      -- split cache_symbol into flag and status with semicolon
      local cache_spam_result = cache_symbol[1]['options'][1]
      local sep_index = string.find(cache_spam_result, ';')
      if sep_index then
        add[constants.SPAM_FLAG_HEADER] = string.sub(cache_spam_result, 1, sep_index - 1)
        spam_status = string.sub(cache_spam_result, sep_index + 1)
      end
    end
  end

  if spam_status == nil then
    if cache.enabled then
      rspamd_logger.infox(task, "Cache miss; calculate spam status")
    end
    spam_status = synology_util.generate_score_status(current_score, required_score)
    if spam_learn_enabled then
      local spam_learn_threshold_spam = spam_opt['learn_threshold_spam']
      local spam_learn_threshold_non_spam = spam_opt['learn_threshold_non_spam']

      if current_score >= spam_learn_threshold_spam then
        spam_status = spam_status .. ', autolearn=spam'
      elseif current_score <= spam_learn_threshold_non_spam then
        spam_status = spam_status .. ', autolearn=ham'
      end
    end

    for _, symbol in ipairs(symbols) do
      if synology_util.is_spam_symbol(symbol) then
        spam_status = spam_status .. ', ' .. symbol['name'] .. ' ' .. symbol['score']
      end
    end

    if action == constants.SYNOLOGY_SPAM_ACTION then
      add[constants.SPAM_FLAG_HEADER] = constants.YES
    else
      add[constants.SPAM_FLAG_HEADER] = constants.NO
    end
  end

  add[constants.SPAM_STATUS_HEADER] = spam_status

  rspamd_logger.infox(task, "Spam Flag: %1", add[constants.SPAM_FLAG_HEADER])
  rspamd_logger.infox(task, "Spam Status: %1", add[constants.SPAM_STATUS_HEADER])
  if add[constants.SPAM_FLAG_HEADER] == constants.YES then
    rspamd_logger.errx(task, "msgid=<%1>: Spam, %2", task:get_message_id(), add[constants.SPAM_STATUS_HEADER])
  end

  -- add spam result to cache if not in cache
  if cache.enabled and not cache_hit then
    -- join flag and status into cache_symbol with semicolon
    cache.set(task, add[constants.SPAM_FLAG_HEADER] .. ';' .. add[constants.SPAM_STATUS_HEADER])
  end
end

local function antivirus_check(task, add, del)
  local antivirus_opt = rspamd_config:get_all_opt('antivirus')
  local antivirus_enabled = antivirus_opt['enabled']

  if not antivirus_enabled then
    return
  end

  local action = task:get_metric_action('default')
  if action ~= constants.SYNOLOGY_ANTIVIRUS_ACTION then
    add[constants.VIRUS_HEADER] = constants.NO
    return
  end

  local virus_names
  for _, antivirus_symbol in ipairs(synology_util.antivirus_symbols()) do
    local virus = task:get_symbol(antivirus_symbol)
    if virus then
      virus_names = virus[1]['options']
      break
    end
  end

  if action == constants.SYNOLOGY_ANTIVIRUS_ACTION then
    virus_names = table.concat(virus_names, ', ')
    add[constants.VIRUS_HEADER] = constants.YES .. ', ' .. virus_names
    rspamd_logger.errx(task, "msgid=<%1>, Virus, %2", task:get_message_id(), virus_names)
  end
end

local function mcp_check(task, add, del)
  local mcp_opt = rspamd_config:get_all_opt('mcp')
  local mcp_enabled = mcp_opt['enabled']

  if not mcp_enabled then
    return
  end

  local mcp_score_required = mcp_opt['actions'][constants.SYNOLOGY_MCP_ACTION]
  local mcp_score = 0
  local mcp_symbols = {}
  local symbols = task:get_symbols_all()

  for _, symbol in ipairs(symbols) do
    if synology_util.is_mcp_symbol(symbol) then
      if symbol['name'] == constants.MCP_SCORE_SYMBOL then
        mcp_score = tonumber(symbol['options'][1])
      else
        mcp_symbols[#mcp_symbols + 1] = symbol['name']
      end
    end
  end

  if mcp_score >= mcp_score_required then
    local mcp_symbol_string = table.concat(mcp_symbols, ', ')
    rspamd_logger.errx(task, "msgid=<%1>, MCP, %2", task:get_message_id(), mcp_symbol_string)

    add[constants.MCP_HEADER] = string.format('%s, %s', constants.YES,
      synology_util.generate_full_score_status(mcp_score, mcp_score_required, mcp_symbol_string))
  else
    add[constants.MCP_HEADER] = constants.NO
  end
end

-- Multiple Headers
rspamd_config:register_symbol({
  name = 'SYNOLOGY_MCP_HEADER',
  type = 'prefilter',
  priority = 9,
  callback = function(task)
    -- Set total MCP score to MCP_SCORE_SYMBOL's option
    -- Remove total MCP score from default metric
    local mcp_opt = rspamd_config:get_all_opt('mcp')
    local mcp_enabled = mcp_opt['enabled']

    if not mcp_enabled then
      return
    end

    local mcp_score = 0
    local symbols = task:get_symbols_all()

    for _, symbol in ipairs(symbols) do
      if synology_util.is_mcp_symbol(symbol) then
        mcp_score = mcp_score + symbol['score']
      end
    end

    task:insert_result(constants.MCP_SCORE_SYMBOL, 0.0, tostring(mcp_score))

    -- remove mcp score
    local score = task:get_metric_score('default') -- score = {[1] = current score, [2] = required score }
    task:set_metric_score('default', score[1] - mcp_score)
  end
})

-- Multiple Headers
rspamd_config:register_symbol({
  name = 'SYNOLOGY_HEADERS',
  type = 'postfilter',
  priority = 10,
  callback = function(task)
    local add = {}
    local del = {
      ['X-Spam'] = 1,
      ['X-Virus'] = 1,
    }

    -- mcp
    mcp_check(task, add, del)

    -- spam
    spam_check(task, add, del)

    -- virus
    antivirus_check(task, add, del)

    task:set_milter_reply({
      add_headers = add,
      remove_headers = del,
    })
  end
})