From ab5767eb8b54344e6ab7546a022a108024e270a2 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Thu, 30 Apr 2020 12:11:37 +0100
Subject: [PATCH] [Fix] Unify selectors digest functions

---
 lualib/lua_selectors/extractors.lua | 29 ++++-------------------------
 lualib/lua_selectors/transforms.lua | 25 ++++---------------------
 2 files changed, 8 insertions(+), 46 deletions(-)

diff --git a/lualib/lua_selectors/extractors.lua b/lualib/lua_selectors/extractors.lua
index 6e023c780..11728c268 100644
--- a/lualib/lua_selectors/extractors.lua
+++ b/lualib/lua_selectors/extractors.lua
@@ -16,6 +16,7 @@ limitations under the License.
 
 local fun = require 'fun'
 local lua_util = require "lua_util"
+local common = require "lua_selectors/common"
 local ts = require("tableshape").types
 local E = {}
 
@@ -131,33 +132,13 @@ uses any type by default)]],
   -- Get list of all attachments digests
   ['attachments'] = {
     ['get_value'] = function(task, args)
-
-      local s
       local parts = task:get_parts() or E
       local digests = {}
 
       if #args > 0 then
-        local rspamd_cryptobox = require "rspamd_cryptobox_hash"
-        local encoding = args[1] or 'hex'
-        local ht = args[2] or 'blake2'
-
-        for _,p in ipairs(parts) do
-          if p:get_filename() then
-            local h = rspamd_cryptobox.create_specific(ht, p:get_content('raw_parsed'))
-            if encoding == 'hex' then
-              s = h:hex()
-            elseif encoding == 'base32' then
-              s = h:base32()
-            elseif encoding == 'base64' then
-              s = h:base64()
-            end
-            table.insert(digests, s)
-          end
-        end
-      else
         for _,p in ipairs(parts) do
           if p:get_filename() then
-            table.insert(digests, p:get_digest())
+            table.insert(digests, common.create_digest(p:get_content('raw_parsed'), args))
           end
         end
       end
@@ -169,11 +150,9 @@ uses any type by default)]],
       return nil
     end,
     ['description'] = [[Get list of all attachments digests.
-The first optional argument is encoding (`hex`, `base32`, `base64`),
+The first optional argument is encoding (`hex`, `base32` (and forms `bleach32`, `rbase32`), `base64`),
 the second optional argument is optional hash type (`blake2`, `sha256`, `sha1`, `sha512`, `md5`)]],
-
-    ['args_schema'] = {ts.one_of{'hex', 'base32', 'base64'}:is_optional(),
-                       ts.one_of{'blake2', 'sha256', 'sha1', 'sha512', 'md5'}:is_optional()}
+    ['args_schema'] = common.digest_schema()
 
   },
   -- Get all attachments files
diff --git a/lualib/lua_selectors/transforms.lua b/lualib/lua_selectors/transforms.lua
index 85a445f67..e47fdcec6 100644
--- a/lualib/lua_selectors/transforms.lua
+++ b/lualib/lua_selectors/transforms.lua
@@ -18,6 +18,7 @@ local fun = require 'fun'
 local lua_util = require "lua_util"
 local ts = require("tableshape").types
 local logger = require 'rspamd_logger'
+local common = require "lua_selectors/common"
 local M = "selectors"
 
 local maps = require "lua_selectors/maps"
@@ -133,31 +134,13 @@ local transform_function = {
     },
     ['map_type'] = 'hash',
     ['process'] = function(inp, _, args)
-      local hash = require 'rspamd_cryptobox_hash'
-      local encoding = args[1] or 'hex'
-      local ht = args[2] or 'blake2'
-      local h = hash:create_specific(ht):update(inp)
-      local s
-
-      if encoding == 'hex' then
-        s = h:hex()
-      elseif encoding == 'base32' then
-        s = h:base32()
-      elseif encoding == 'bleach32' then
-        s = h:base32('bleach')
-      elseif encoding == 'rbase32' then
-        s = h:base32('rfc')
-      elseif encoding == 'base64' then
-        s = h:base64()
-      end
 
-      return s,'string'
+      return common.create_digest(inp, args),'string'
     end,
     ['description'] = [[Create a digest from a string.
 The first argument is encoding (`hex`, `base32` (and forms `bleach32`, `rbase32`), `base64`),
 the second argument is optional hash type (`blake2`, `sha256`, `sha1`, `sha512`, `md5`)]],
-    ['args_schema'] = {ts.one_of{'hex', 'base32', 'base64'}:is_optional(),
-                       ts.one_of{'blake2', 'sha256', 'sha1', 'sha512', 'md5'}:is_optional()}
+    ['args_schema'] = common.digest_schema()
   },
   -- Extracts substring
   ['substring'] = {
@@ -430,7 +413,7 @@ Empty string comes the first argument or 'true', non-empty string comes nil]],
           function(s)
             return string.gsub(tostring(s), '[\128-\255]', args[1] or '?')
           end, inp), 'string_list'
-      else 
+      else
         return string.gsub(tostring(inp), '[\128-\255]', '?'), 'string'
       end
     end,
-- 
2.39.5