[Rework] Move mime modification functions to lua_mime library

2 files changed, 253 insertions, 187 deletions

diff --git a/lualib/lua_mime.lua b/lualib/lua_mime.lua
new file mode 100644
index 000000000..09c8bf82a
--- /dev/null
+++ b/lualib/lua_mime.lua
@@ -0,0 +1,239 @@
+--[[
+Copyright (c) 2019, Vsevolod Stakhov <vsevolod@highsecure.ru>
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+]]--
+
+--[[[
+-- @module lua_mime
+-- This module contains helper functions to modify mime parts
+--]]
+
+local rspamd_util = require "rspamd_util"
+
+local exports = {}
+
+local function newline(task)
+  local t = task:get_newlines_type()
+
+  if t == 'cr' then
+    return '\r'
+  elseif t == 'lf' then
+    return '\n'
+  end
+
+  return '\r\n'
+end
+
+
+--[[[
+-- @function lua_mime.add_text_footer(task, html_footer, text_footer)
+-- Adds a footer to all text parts in a message. It returns a table with the following
+-- fields:
+-- * out: new content (body only)
+-- * need_rewrite_ct: boolean field that means if we must rewrite content type
+-- * new_ct: new content type (type => string, subtype => string)
+--]]
+exports.add_text_footer = function(task, html_footer, text_footer)
+  local newline_s = newline(task)
+  local res = {}
+  local out = {}
+  local text_parts = task:get_text_parts()
+
+  if not (html_footer or text_footer) or not (text_parts and #text_parts > 0) then
+    return false
+  end
+
+  local function do_append_footer(part, footer, is_multipart)
+    local tp = part:get_text()
+    local ct = 'text/plain'
+    local cte = 'quoted-printable'
+
+    if tp:is_html() then
+      ct = 'text/html'
+    end
+
+    if part:get_cte() == '7bit' then
+      cte = '7bit'
+    end
+
+    if is_multipart then
+      out[#out + 1] = string.format('Content-Type: %s; charset=utf-8%s'..
+          'Content-Transfer-Encoding: %s',
+          ct, newline_s, cte)
+      out[#out + 1] = ''
+    end
+
+    local content = tostring(tp:get_content('raw_utf') or '')
+    local double_nline = newline_s .. newline_s
+    local nlen = #double_nline
+    -- Hack, if part ends with 2 newline, then we append it after footer
+    if content:sub(-(nlen), nlen + 1) == double_nline then
+      content = string.format('%s%s',
+          content:sub(-(#newline_s), #newline_s + 1), -- content without last newline
+          footer)
+      out[#out + 1] = {rspamd_util.encode_qp(content,
+          80, task:get_newlines_type()), true}
+      out[#out + 1] = ''
+    else
+      content = content .. footer
+      out[#out + 1] = {rspamd_util.encode_qp(content,
+          80, task:get_newlines_type()), true}
+      out[#out + 1] = ''
+    end
+
+  end
+
+  if html_footer or text_footer then
+    -- We need to take extra care about content-type and cte
+    local ct = task:get_header('Content-Type')
+    if ct then
+      ct = rspamd_util.parse_content_type(ct, task:get_mempool())
+    end
+
+    if ct then
+      if ct.type and ct.type == 'text' then
+        if ct.subtype then
+          if html_footer and (ct.subtype == 'html' or ct.subtype == 'htm') then
+            res.need_rewrite_ct = true
+          elseif text_footer and ct.subtype == 'plain' then
+            res.need_rewrite_ct = true
+          end
+        else
+          if text_footer then
+            res.need_rewrite_ct = true
+          end
+        end
+
+        res.new_ct = ct
+      end
+    else
+
+      if text_parts then
+
+        if #text_parts == 1 then
+          res.need_rewrite_ct = true
+          res.new_ct = {
+            type = 'text',
+            subtype = 'plain'
+          }
+        elseif #text_parts > 1 then
+          -- XXX: in fact, it cannot be
+          res.new_ct = {
+            type = 'multipart',
+            subtype = 'mixed'
+          }
+        end
+      end
+    end
+  end
+
+  local boundaries = {}
+  local cur_boundary
+
+  for _,part in ipairs(task:get_parts()) do
+    local boundary = part:get_boundary()
+    if part:is_multipart() then
+      if cur_boundary then
+        out[#out + 1] = string.format('--%s',
+            boundaries[#boundaries])
+      end
+
+      boundaries[#boundaries + 1] = boundary or '--XXX'
+      cur_boundary = boundary
+
+      local rh = part:get_raw_headers()
+      if #rh > 0 then
+        out[#out + 1] = {rh, true}
+      end
+    elseif part:is_message() then
+      if boundary then
+        if cur_boundary and boundary ~= cur_boundary then
+          -- Need to close boundary
+          out[#out + 1] = string.format('--%s--%s',
+              boundaries[#boundaries], newline_s)
+          table.remove(boundaries)
+          cur_boundary = nil
+        end
+        out[#out + 1] = string.format('--%s',
+            boundary)
+      end
+
+      out[#out + 1] = {part:get_raw_headers(), true}
+    else
+      local append_footer = false
+      local skip_footer = part:is_attachment()
+
+      local parent = part:get_parent()
+      if parent then
+        local t,st = parent:get_type()
+
+        if t == 'multipart' and st == 'signed' then
+          -- Do not modify signed parts
+          skip_footer = true
+        end
+      end
+      if text_footer and part:is_text() then
+        local tp = part:get_text()
+
+        if not tp:is_html() then
+          append_footer = text_footer
+        end
+      end
+
+      if html_footer and part:is_text() then
+        local tp = part:get_text()
+
+        if tp:is_html() then
+          append_footer = html_footer
+        end
+      end
+
+      if boundary then
+        if cur_boundary and boundary ~= cur_boundary then
+          -- Need to close boundary
+          out[#out + 1] = string.format('--%s--%s',
+              boundaries[#boundaries], newline_s)
+          table.remove(boundaries)
+          cur_boundary = boundary
+        end
+        out[#out + 1] = string.format('--%s',
+            boundary)
+      end
+
+      if append_footer and not skip_footer then
+        do_append_footer(part, append_footer,
+            parent and parent:is_multipart())
+      else
+        out[#out + 1] = {part:get_raw_headers(), true}
+        out[#out + 1] = {part:get_raw_content(), false}
+      end
+    end
+  end
+
+  -- Close remaining
+  local b = table.remove(boundaries)
+  while b do
+    out[#out + 1] = string.format('--%s--', b)
+    if #boundaries > 0 then
+      out[#out + 1] = ''
+    end
+    b = table.remove(boundaries)
+  end
+
+  res.out = out
+
+  return res
+end
+
+return exports
\ No newline at end of file
diff --git a/lualib/rspamadm/mime.lua b/lualib/rspamadm/mime.lua
index 91bc06993..ae29f0eb1 100644
--- a/lualib/rspamadm/mime.lua
+++ b/lualib/rspamadm/mime.lua
@@ -22,6 +22,7 @@ local rspamd_logger = require "rspamd_logger"
 local lua_meta = require "lua_meta"
 local rspamd_url = require "rspamd_url"
 local lua_util = require "lua_util"
+local lua_mime = require "lua_mime"
 local ucl = require "ucl"
 
 -- Define command line options
@@ -642,47 +643,6 @@ local function modify_handler(opts)
     return content
   end
 
-  local function do_append_footer(task, part, footer, is_multipart, out)
-    local newline_s = newline(task)
-    local tp = part:get_text()
-    local ct = 'text/plain'
-    local cte = 'quoted-printable'
-
-    if tp:is_html() then
-      ct = 'text/html'
-    end
-
-    if part:get_cte() == '7bit' then
-      cte = '7bit'
-    end
-
-    if is_multipart then
-      out[#out + 1] = string.format('Content-Type: %s; charset=utf-8%s'..
-          'Content-Transfer-Encoding: %s',
-          ct, newline_s, cte)
-      out[#out + 1] = ''
-    end
-
-    local content = tostring(tp:get_content('raw_utf') or '')
-    local double_nline = newline_s .. newline_s
-    local nlen = #double_nline
-    -- Hack, if part ends with 2 newline, then we append it after footer
-    if content:sub(-(nlen), nlen + 1) == double_nline then
-      content = string.format('%s%s',
-          content:sub(-(#newline_s), #newline_s + 1), -- content without last newline
-          footer)
-      out[#out + 1] = {rspamd_util.encode_qp(content,
-          80, task:get_newlines_type()), true}
-      out[#out + 1] = ''
-    else
-      content = content .. footer
-      out[#out + 1] = {rspamd_util.encode_qp(content,
-          80, task:get_newlines_type()), true}
-      out[#out + 1] = ''
-    end
-
-  end
-
   local text_footer, html_footer
 
   if opts['text_footer'] then
@@ -696,13 +656,12 @@ local function modify_handler(opts)
   for _,fname in ipairs(opts.file) do
     local task = load_task(opts, fname)
     local newline_s = newline(task)
-    local need_rewrite_ct = false
-    local parsed_ct
-    local seen_cte = false
-    local out = {}
+    local seen_cte
 
-    local function process_headers_cb(name, hdr)
+    local rewrite = lua_mime.add_text_footer(task, html_footer, text_footer) or {}
+    local out = {} -- Start with headers
 
+    local function process_headers_cb(name, hdr)
       for _,h in ipairs(opts['remove_header']) do
         if name:match(h) then
           return
@@ -723,16 +682,16 @@ local function modify_handler(opts)
         end
       end
 
-      if need_rewrite_ct then
+      if rewrite.need_rewrite_ct then
         if name:lower() == 'content-type' then
           local nct = string.format('%s: %s/%s; charset=utf-8',
-              'Content-Type', parsed_ct.type, parsed_ct.subtype)
+              'Content-Type', rewrite.new_ct.type, rewrite.new_ct.subtype)
           out[#out + 1] = nct
           return
         elseif name:lower() == 'content-transfer-encoding' then
-          seen_cte = true
           out[#out + 1] = string.format('%s: %s',
               'Content-Transfer-Encoding', 'quoted-printable')
+          seen_cte = true
           return
         end
       end
@@ -740,50 +699,6 @@ local function modify_handler(opts)
       out[#out + 1] = hdr.raw:gsub('\r?\n?$', '')
     end
 
-    if html_footer or text_footer then
-      -- We need to take extra care about content-type and cte
-      local ct = task:get_header('Content-Type')
-      if ct then
-        ct = rspamd_util.parse_content_type(ct, task:get_mempool())
-      end
-
-      if ct then
-        if ct.type and ct.type == 'text' then
-          if ct.subtype then
-            if html_footer and (ct.subtype == 'html' or ct.subtype == 'htm') then
-              need_rewrite_ct = true
-            elseif text_footer and ct.subtype == 'plain' then
-              need_rewrite_ct = true
-            end
-          else
-            if text_footer then
-              need_rewrite_ct = true
-            end
-          end
-
-          parsed_ct = ct
-        end
-      else
-        local text_parts = task:get_text_parts()
-        if text_parts then
-
-          if #text_parts == 1 then
-            need_rewrite_ct = true
-            parsed_ct = {
-              type = 'text',
-              subtype = 'plain'
-            }
-          elseif #text_parts > 1 then
-            -- XXX: in fact, it cannot be
-            parsed_ct = {
-              type = 'multipart',
-              subtype = 'mixed'
-            }
-          end
-        end
-      end
-    end
-
     task:headers_foreach(process_headers_cb, {full = true})
 
     for _,h in ipairs(opts['add_header']) do
@@ -795,108 +710,20 @@ local function modify_handler(opts)
       end
     end
 
-    if not seen_cte and need_rewrite_ct then
+    if not seen_cte and rewrite.need_rewrite_ct then
       out[#out + 1] = string.format('%s: %s',
           'Content-Transfer-Encoding', 'quoted-printable')
     end
 
     -- End of headers
-    --local eoh_pos = #out
     out[#out + 1] = ''
 
-    local boundaries = {}
-    local cur_boundary
-
-    for _,part in ipairs(task:get_parts()) do
-      local boundary = part:get_boundary()
-      if part:is_multipart() then
-        if cur_boundary then
-          out[#out + 1] = string.format('--%s',
-              boundaries[#boundaries])
-        end
-
-        boundaries[#boundaries + 1] = boundary or '--XXX'
-        cur_boundary = boundary
-
-        local rh = part:get_raw_headers()
-        if #rh > 0 then
-          out[#out + 1] = {rh, true}
-        end
-      elseif part:is_message() then
-        if boundary then
-          if cur_boundary and boundary ~= cur_boundary then
-            -- Need to close boundary
-            out[#out + 1] = string.format('--%s--%s',
-                boundaries[#boundaries], newline_s)
-            table.remove(boundaries)
-            cur_boundary = nil
-          end
-          out[#out + 1] = string.format('--%s',
-              boundary)
-        end
-
-        out[#out + 1] = {part:get_raw_headers(), true}
-      else
-        local append_footer = false
-        local skip_footer = part:is_attachment()
-
-        local parent = part:get_parent()
-        if parent then
-          local t,st = parent:get_type()
-
-          if t == 'multipart' and st == 'signed' then
-            -- Do not modify signed parts
-            skip_footer = true
-          end
-        end
-        if text_footer and part:is_text() then
-          local tp = part:get_text()
-
-          if not tp:is_html() then
-            append_footer = text_footer
-          end
-        end
-
-        if html_footer and part:is_text() then
-          local tp = part:get_text()
-
-          if tp:is_html() then
-            append_footer = html_footer
-          end
-        end
-
-        if boundary then
-          if cur_boundary and boundary ~= cur_boundary then
-            -- Need to close boundary
-            out[#out + 1] = string.format('--%s--%s',
-                boundaries[#boundaries], newline_s)
-            table.remove(boundaries)
-            cur_boundary = boundary
-          end
-          out[#out + 1] = string.format('--%s',
-              boundary)
-        end
-
-        io.flush()
-
-        if append_footer and not skip_footer then
-          do_append_footer(task, part, append_footer,
-              parent and parent:is_multipart(), out)
-        else
-          out[#out + 1] = {part:get_raw_headers(), true}
-          out[#out + 1] = {part:get_raw_content(), false}
-        end
+    if rewrite.out then
+      for _,o in ipairs(rewrite.out) do
+        out[#out + 1] = o
       end
-    end
-
-    -- Close remaining
-    local b = table.remove(boundaries)
-    while b do
-      out[#out + 1] = string.format('--%s--', b)
-      if #boundaries > 0 then
-        out[#out + 1] = ''
-      end
-      b = table.remove(boundaries)
+    else
+      out[#out + 1] = task:get_rawbody()
     end
 
     for _,o in ipairs(out) do