[Project] Various fixes and `rspamadm mime strip` commandvstakhov-strip-attachments

1 files changed, 54 insertions, 77 deletions

diff --git a/lualib/lua_mime.lua b/lualib/lua_mime.lua
index 71867fbed..795a803e5 100644
--- a/lualib/lua_mime.lua
+++ b/lualib/lua_mime.lua
@@ -770,9 +770,6 @@ end
 --   * max_text_size: number, maximum text part size to keep (default: inf)
 -- @return {table} modified message state similar to other modification functions:
 -- * 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)
--- * new_cte: new content-transfer encoding (string)
 --]]
 exports.remove_attachments = function(task, settings)
   local newline_s = newline(task)
@@ -797,7 +794,7 @@ exports.remove_attachments = function(task, settings)
   for i, part in ipairs(task:get_parts()) do
     local keep_part = false
 
-    if part:is_text() then
+    if part:is_text() and not part:is_attachment() then
       local length = part:get_length()
       if length >= min_text_size and length <= max_text_size then
         keep_part = true
@@ -822,100 +819,80 @@ exports.remove_attachments = function(task, settings)
     return false
   end
 
-  -- Prepare new message structure
-  local need_multipart = false
-  local text_parts_count = 0
-  for _, part in ipairs(parts_to_keep) do
-    if part:is_text() then
-      text_parts_count = text_parts_count + 1
-    end
-  end
-  need_multipart = text_parts_count > 1 or (keep_images and next(parts_to_keep))
-
-  -- Set content type
-  if need_multipart then
-    state.new_ct = {
-      type = 'multipart',
-      subtype = 'mixed'
-    }
-    cur_boundary = '--XXX'
-    boundaries[1] = cur_boundary
-
-    out[#out + 1] = {
-      string.format('Content-Type: multipart/mixed; boundary="%s"%s',
-          cur_boundary, newline_s),
-      true
-    }
-    out[#out + 1] = { '', true }
-  else
-    -- Single part message
-    for _, part in ipairs(parts_to_keep) do
-      if part:is_text() then
-        state.new_ct = {
-          type = 'text',
-          subtype = part:get_text():is_html() and 'html' or 'plain'
-        }
-        break
-      end
-    end
-  end
-
   -- Second pass: reconstruct message
   for i, part in ipairs(task:get_parts()) do
+    local boundary = part:get_boundary()
     if part:is_multipart() then
-      -- Skip multipart containers
-      local boundary = part:get_boundary()
+      if cur_boundary then
+        out[#out + 1] = { string.format('--%s',
+            boundaries[#boundaries]), true }
+      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
-          out[#out + 1] = {
-            string.format('--%s--', boundaries[#boundaries]),
-            true
-          }
+          -- Need to close boundary
+          out[#out + 1] = { string.format('--%s--',
+              boundaries[#boundaries]), true }
           table.remove(boundaries)
+          cur_boundary = nil
         end
-      end
-    elseif parts_indexes_to_keep[i] then
-      if need_multipart then
-        out[#out + 1] = {
-          string.format('--%s', cur_boundary),
-          true
-        }
+        out[#out + 1] = { string.format('--%s',
+            boundary), true }
       end
 
-      -- Add part headers
-      local headers = {}
-      for _, h in ipairs(part:get_header_array()) do
-        table.insert(headers, string.format('%s: %s', h.name, h.value))
-      end
+      out[#out + 1] = { part:get_raw_headers(), true }
+    else
+      if parts_indexes_to_keep[i] then
+        if boundary then
+          if cur_boundary and boundary ~= cur_boundary then
+            -- Need to close previous boundary
+            out[#out + 1] = { string.format('--%s--',
+                boundaries[#boundaries]), true }
+            table.remove(boundaries)
+            cur_boundary = boundary
+          end
+          out[#out + 1] = { string.format('--%s',
+              boundary), true }
+        end
+        -- Add part headers
+        local headers = part:get_raw_headers()
 
-      if #headers > 0 then
+        if headers then
+          out[#out + 1] = {
+            headers,
+            true
+          }
+        end
+
+        -- Add content
         out[#out + 1] = {
-          table.concat(headers, newline_s),
-          true
+          part:get_raw_content(),
+          false
         }
       end
 
-      -- Add empty line between headers and content
-      out[#out + 1] = { '', true }
-
-      -- Add content
-      out[#out + 1] = {
-        part:get_raw_content(),
-        false
-      }
     end
   end
 
   -- Close remaining boundaries
-  if need_multipart then
-    out[#out + 1] = {
-      string.format('--%s--', cur_boundary),
-      true
-    }
+  local b = table.remove(boundaries)
+  while b do
+    out[#out + 1] = { string.format('--%s--', b), true }
+    if #boundaries > 0 then
+      out[#out + 1] = { '', true }
+    end
+    b = table.remove(boundaries)
   end
 
   state.out = out
-  state.need_rewrite_ct = true
 
   return state
 end