From 413d624f765a49d29d8e4ab8607121f2a1e9d049 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Tue, 23 Apr 2019 15:48:53 +0100
Subject: [PATCH] [Feature] Lua_util: Allow to obfuscate different fields

---
 lualib/lua_util.lua               | 28 +++++++++++++++-------------
 src/plugins/lua/clickhouse.lua    |  5 +++--
 src/plugins/lua/history_redis.lua |  2 +-
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/lualib/lua_util.lua b/lualib/lua_util.lua
index 4f185ecab..27df2c72e 100644
--- a/lualib/lua_util.lua
+++ b/lualib/lua_util.lua
@@ -934,29 +934,31 @@ exports.get_task_verdict = function(task)
 end
 
 ---[[[
--- @function lua_util.maybe_obfuscate_subject(subject, settings)
--- Obfuscate subject if enabled in settings. Also checks utf8 validity.
+-- @function lua_util.maybe_obfuscate_string(subject, settings, prefix)
+-- Obfuscate string if enabled in settings. Also checks utf8 validity.
 -- Supported settings:
--- * subject_privacy = false - subject privacy is off
--- * subject_privacy_alg = 'blake2' - default hash-algorithm to obfuscate subject
--- * subject_privacy_prefix = 'obf' - prefix to show it's obfuscated
--- * subject_privacy_length = 16 - cut the length of the hash
+-- * <prefix>_privacy = false - subject privacy is off
+-- * <prefix>_privacy_alg = 'blake2' - default hash-algorithm to obfuscate subject
+-- * <prefix>_privacy_prefix = 'obf' - prefix to show it's obfuscated
+-- * <prefix>_privacy_length = 16 - cut the length of the hash
 -- @return obfuscated or validated subject
 --]]
 
-exports.maybe_obfuscate_subject = function(subject, settings)
+exports.maybe_obfuscate_string = function(subject, settings, prefix)
   local hash = require 'rspamd_cryptobox_hash'
   if subject and not rspamd_util.is_valid_utf8(subject) then
     subject = '???'
-  elseif settings.subject_privacy then
-    local hash_alg = settings.subject_privacy_alg or 'blake2'
+  elseif settings[prefix .. '_privacy'] then
+    local hash_alg = settings[prefix .. '_privacy_alg'] or 'blake2'
     local subject_hash = hash.create_specific(hash_alg, subject)
+    local strip_len = settings[prefix .. '_privacy_length']
+    local privacy_prefix = settings[prefix .. '_privacy_prefix'] or ''
 
-    if settings.subject_privacy_length then
-      subject = (settings.subject_privacy_prefix or 'obf') .. ':' ..
-          subject_hash:hex():sub(1, settings.subject_privacy_length)
+    if strip_len then
+      subject = privacy_prefix .. ':' ..
+          subject_hash:hex():sub(1, strip_len)
     else
-      subject = (settings.subject_privacy_prefix or '') .. ':' ..
+      subject = privacy_prefix .. ':' ..
           subject_hash:hex()
     end
   end
diff --git a/src/plugins/lua/clickhouse.lua b/src/plugins/lua/clickhouse.lua
index c3be9e2d9..9559ef5cb 100644
--- a/src/plugins/lua/clickhouse.lua
+++ b/src/plugins/lua/clickhouse.lua
@@ -438,7 +438,8 @@ local function clickhouse_collect(task)
   end
 
   local list_id = task:get_header('List-Id') or ''
-  local message_id = task:get_message_id() or ''
+  local message_id = lua_util.maybe_obfuscate_string(task:get_message_id() or '',
+      settings, 'mid')
 
   local score = task:get_metric_score('default')[1];
   local bayes = 'unknown';
@@ -589,7 +590,7 @@ local function clickhouse_collect(task)
 
   local subject = ''
   if settings.insert_subject then
-    subject = lua_util.maybe_obfuscate_subject(task:get_subject() or '', settings)
+    subject = lua_util.maybe_obfuscate_string(task:get_subject() or '', settings, 'subject')
   end
 
   local scan_real,scan_virtual = task:get_scan_time()
diff --git a/src/plugins/lua/history_redis.lua b/src/plugins/lua/history_redis.lua
index c18ea736c..ed97cb7fe 100644
--- a/src/plugins/lua/history_redis.lua
+++ b/src/plugins/lua/history_redis.lua
@@ -207,7 +207,7 @@ local function handle_history_request(task, conn, from, to, reset)
         collectgarbage()
         t1 = rspamd_util:get_ticks()
         fun.each(function(e)
-          e.subject = lua_util.maybe_obfuscate_subject(e.subject, settings)
+          e.subject = lua_util.maybe_obfuscate_string(e.subject, settings, 'subject')
         end, data)
         reply.rows = data
         conn:send_ucl(reply)
-- 
2.39.5