Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import jakarta.activation.DataSource;
import jakarta.mail.Address;
import jakarta.mail.BodyPart;
import jakarta.mail.Header;
import jakarta.mail.MessagingException;
import jakarta.mail.Multipart;
import jakarta.mail.Session;
Expand Down Expand Up @@ -46,6 +47,7 @@
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand All @@ -61,7 +63,8 @@
@WritesAttribute(attribute = "filename ", description = "The filename of the attachment"),
@WritesAttribute(attribute = "email.attachment.parent.filename ", description = "The filename of the parent FlowFile"),
@WritesAttribute(attribute = "email.attachment.parent.uuid", description = "The UUID of the original FlowFile."),
@WritesAttribute(attribute = "mime.type", description = "The mime type of the attachment.")})
@WritesAttribute(attribute = "mime.type", description = "The mime type of the attachment."),
@WritesAttribute(attribute = ExtractEmailAttachments.ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + "<attachment header name>", description = "Attachment header.")})

public class ExtractEmailAttachments extends AbstractProcessor {
public static final String ATTACHMENT_ORIGINAL_FILENAME = "email.attachment.parent.filename";
Expand All @@ -80,6 +83,7 @@ public class ExtractEmailAttachments extends AbstractProcessor {
.description("FlowFiles that could not be parsed")
.build();

static final String ATTACHMENT_HEADER_ATTRIBUTE_PREFIX = "email.attachment.header.";
private static final String ATTACHMENT_DISPOSITION = "attachment";

private static final Set<Relationship> RELATIONSHIPS = Set.of(
Expand Down Expand Up @@ -117,12 +121,13 @@ public void onTrigger(final ProcessContext context, final ProcessSession session

final String originalFlowFileName = originalFlowFile.getAttribute(CoreAttributes.FILENAME.key());
try {
final List<DataSource> attachments = new ArrayList<>();
final List<Attachment> attachments = new ArrayList<>();
parseAttachments(attachments, originalMessage, 0);

for (final DataSource data : attachments) {
for (final Attachment attachment : attachments) {
FlowFile split = session.create(originalFlowFile);
final Map<String, String> attributes = new HashMap<>();
final DataSource data = attachment.dataSource();
final String name = data.getName();
if (name != null && !name.isBlank()) {
attributes.put(CoreAttributes.FILENAME.key(), name);
Expand All @@ -131,6 +136,13 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
if (contentType != null && !contentType.isBlank()) {
attributes.put(CoreAttributes.MIME_TYPE.key(), contentType);
}

for (Map.Entry<String, String> entry : attachment.headers().entrySet()) {
final String headerAttributeName = ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + entry.getKey().toLowerCase();
final String headerAttributeValue = entry.getValue();
attributes.put(headerAttributeName, headerAttributeValue);
}

String parentUuid = originalFlowFile.getAttribute(CoreAttributes.UUID.key());
attributes.put(ATTACHMENT_ORIGINAL_UUID, parentUuid);
attributes.put(ATTACHMENT_ORIGINAL_FILENAME, originalFlowFileName);
Expand Down Expand Up @@ -176,7 +188,7 @@ public Set<Relationship> getRelationships() {
return RELATIONSHIPS;
}

private void parseAttachments(final List<DataSource> attachments, final MimePart parentPart, final int depth) throws MessagingException, IOException {
private void parseAttachments(final List<Attachment> attachments, final MimePart parentPart, final int depth) throws MessagingException, IOException {
final String disposition = parentPart.getDisposition();

final Object parentContent = parentPart.getContent();
Expand All @@ -191,7 +203,24 @@ private void parseAttachments(final List<DataSource> attachments, final MimePart
}
} else if (ATTACHMENT_DISPOSITION.equalsIgnoreCase(disposition) || depth > 0) {
final DataSource dataSource = parentPart.getDataHandler().getDataSource();
attachments.add(dataSource);
final Map<String, String> extractedHeaders = new HashMap<>();

if (parentPart instanceof final MimeBodyPart mimeBodyPart) {
final Enumeration<Header> headers = mimeBodyPart.getAllHeaders();
while (headers.hasMoreElements()) {
final Header header = headers.nextElement();
final String name = header.getName();
if (name != null && !name.isBlank()) {
final String value = header.getValue();
extractedHeaders.put(name, value);
}
}
}

final Attachment attachment = new Attachment(dataSource, extractedHeaders);
attachments.add(attachment);
}
}
}

record Attachment(DataSource dataSource, Map<String, String> headers) { }
Original file line number Diff line number Diff line change
Expand Up @@ -17,18 +17,26 @@

package org.apache.nifi.processors.email;

import jakarta.mail.Session;
import jakarta.mail.internet.InternetAddress;
import jakarta.mail.internet.MimeBodyPart;
import jakarta.mail.internet.MimeMessage;
import jakarta.mail.internet.MimeMultipart;
import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.io.ByteArrayOutputStream;
import java.util.List;

import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.Map;
import java.util.Properties;

public class TestExtractEmailAttachments {
private static final String EXPECTED_CONTENT_TYPE_KEY = ExtractEmailAttachments.ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + "content-type";
private static final String EXPECTED_CONTENT_DISPOSITION_KEY = ExtractEmailAttachments.ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + "content-disposition";

final String from = "Alice <alice@nifi.apache.org>";
final String to = "bob@nifi.apache.org";
final String subject = "Just a test email";
Expand All @@ -37,10 +45,15 @@ public class TestExtractEmailAttachments {

final GenerateAttachment attachmentGenerator = new GenerateAttachment(from, to, subject, message, hostName);

TestRunner runner;

@BeforeEach
void setUp() {
runner = TestRunners.newTestRunner(ExtractEmailAttachments.class);
}

@Test
public void testValidEmailWithAttachments() {
final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailAttachments());

byte[] withAttachment = attachmentGenerator.withAttachments(1);

runner.enqueue(withAttachment);
Expand All @@ -51,13 +64,17 @@ public void testValidEmailWithAttachments() {
runner.assertTransferCount(ExtractEmailAttachments.REL_ATTACHMENTS, 1);
// Have a look at the attachments...
final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(ExtractEmailAttachments.REL_ATTACHMENTS);
splits.get(0).assertAttributeEquals("filename", "pom.xml-0");
final MockFlowFile split = splits.getFirst();
split.assertAttributeEquals("filename", "pom.xml-0");
final Map<String, String> expected = Map.of(
EXPECTED_CONTENT_DISPOSITION_KEY, "attachment; filename=\"pom.xml-0\"",
EXPECTED_CONTENT_TYPE_KEY, "text/plain; charset=utf-8"
);
assertAttachmentHeaderAttributes(split, expected);
}

@Test
public void testValidEmailWithMultipleAttachments() {
final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailAttachments());

int amount = 3;
byte[] withAttachment = attachmentGenerator.withAttachments(amount);

Expand All @@ -69,19 +86,22 @@ public void testValidEmailWithMultipleAttachments() {
runner.assertTransferCount(ExtractEmailAttachments.REL_ATTACHMENTS, amount);

final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(ExtractEmailAttachments.REL_ATTACHMENTS);

List<String> filenames = new ArrayList<>();
for (int a = 0; a < amount; a++) {
filenames.add(splits.get(a).getAttribute("filename"));
final String expectedContentType = "text/plain; charset=utf-8";
final List<Map<String, String>> expectedHeaderAttachmentAttributes = List.of(
Map.of(EXPECTED_CONTENT_DISPOSITION_KEY, "attachment; filename=\"pom.xml-0\"", EXPECTED_CONTENT_TYPE_KEY, expectedContentType),
Map.of(EXPECTED_CONTENT_DISPOSITION_KEY, "attachment; filename=\"pom.xml-1\"", EXPECTED_CONTENT_TYPE_KEY, expectedContentType),
Map.of(EXPECTED_CONTENT_DISPOSITION_KEY, "attachment; filename=\"pom.xml-2\"", EXPECTED_CONTENT_TYPE_KEY, expectedContentType)
);

for (int index = 0; index < amount; index++) {
final MockFlowFile split = splits.get(index);
split.assertAttributeEquals("filename", "pom.xml-" + index);
assertAttachmentHeaderAttributes(split, expectedHeaderAttachmentAttributes.get(index));
}

assertTrue(filenames.containsAll(Arrays.asList("pom.xml-0", "pom.xml-1", "pom.xml-2")));
}

@Test
public void testValidEmailWithoutAttachments() {
final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailAttachments());

byte[] simpleEmail = attachmentGenerator.simpleMessage();

runner.enqueue(simpleEmail);
Expand All @@ -94,12 +114,108 @@ public void testValidEmailWithoutAttachments() {

@Test
public void testInvalidEmail() {
final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailAttachments());
runner.enqueue("test test test chocolate".getBytes());
runner.run();

runner.assertTransferCount(ExtractEmailAttachments.REL_ORIGINAL, 0);
runner.assertTransferCount(ExtractEmailAttachments.REL_FAILURE, 1);
runner.assertTransferCount(ExtractEmailAttachments.REL_ATTACHMENTS, 0);
}

@Test
public void testDeeplyNestedMultipartMimeMessage() throws Exception {
final byte[] deeplyNestedMultipartMimeMessage = generateDeeplyNestedMultipartMimeMessage();
runner.enqueue(deeplyNestedMultipartMimeMessage);
runner.run();

runner.assertTransferCount(ExtractEmailAttachments.REL_ORIGINAL, 1);
runner.assertTransferCount(ExtractEmailAttachments.REL_FAILURE, 0);
runner.assertTransferCount(ExtractEmailAttachments.REL_ATTACHMENTS, 4);

final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(ExtractEmailAttachments.REL_ATTACHMENTS);
final String expectedContentTransferEncodingKey = ExtractEmailAttachments.ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + "content-transfer-encoding";

final List<Map<String, String>> expectedHeaderAttachmentAttributes = List.of(
Map.of(expectedContentTransferEncodingKey, "quoted-printable", EXPECTED_CONTENT_TYPE_KEY, "text/plain; charset=iso-8859-1"),
Map.of(expectedContentTransferEncodingKey, "quoted-printable", EXPECTED_CONTENT_TYPE_KEY, "text/html; charset=iso-8859-1"),
Map.of(EXPECTED_CONTENT_DISPOSITION_KEY, "inline; filename=\"inline_image.png\"", expectedContentTransferEncodingKey, "base64",
EXPECTED_CONTENT_TYPE_KEY, "image/png; name=inline_image.png",
ExtractEmailAttachments.ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + "content-id", "<0011223344556677@8899AABBCCDDEEFF>"),
Map.of(EXPECTED_CONTENT_DISPOSITION_KEY, getPdfContentDisposition(), expectedContentTransferEncodingKey, "base64", EXPECTED_CONTENT_TYPE_KEY, "application/pdf; name=my-attachment.pdf",
ExtractEmailAttachments.ATTACHMENT_HEADER_ATTRIBUTE_PREFIX + "content-description", "my-attachment.pdf")
);

for (int index = 0; index < splits.size(); index++) {
MockFlowFile split = splits.get(index);
assertAttachmentHeaderAttributes(split, expectedHeaderAttachmentAttributes.get(index));
}
}

private byte[] generateDeeplyNestedMultipartMimeMessage() throws Exception {
final Properties props = new Properties();
final Session session = Session.getDefaultInstance(props, null);
final MimeMessage message = new MimeMessage(session);
message.setFrom(new InternetAddress("sender@example.com"));
message.addRecipient(MimeMessage.RecipientType.TO, new InternetAddress("receiver@example.com"));
message.setSubject("Deeply Nested Multipart Test");

final MimeMultipart mixedMultipart = new MimeMultipart("mixed");
final MimeMultipart relatedMultipart = new MimeMultipart("related");
relatedMultipart.setSubType("related; type=\"multipart/alternative\"");
final MimeMultipart alternativeMultipart = new MimeMultipart("alternative");

final MimeBodyPart plainTextPart = new MimeBodyPart();
plainTextPart.setContent("Hello World! This is plain text body.", "text/plain; charset=iso-8859-1");
plainTextPart.setHeader("Content-Transfer-Encoding", "quoted-printable");
alternativeMultipart.addBodyPart(plainTextPart);

final MimeBodyPart htmlTextPart = new MimeBodyPart();
htmlTextPart.setContent("<html><body><h1>Hello World!</h1> This is HTML body.</body></html>", "text/html; charset=iso-8859-1");
htmlTextPart.setHeader("Content-Transfer-Encoding", "quoted-printable");
alternativeMultipart.addBodyPart(htmlTextPart);

final MimeBodyPart alternativeWrapperPart = new MimeBodyPart();
alternativeWrapperPart.setContent(alternativeMultipart);
relatedMultipart.addBodyPart(alternativeWrapperPart);

final MimeBodyPart inlineImagePart = new MimeBodyPart();
inlineImagePart.setContent("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=", "image/png; name=\"inline_image.png\"");
inlineImagePart.setDisposition("inline; filename=\"inline_image.png\"");
inlineImagePart.setHeader("Content-ID", "<0011223344556677@8899AABBCCDDEEFF>");
inlineImagePart.setHeader("Content-Transfer-Encoding", "base64");
relatedMultipart.addBodyPart(inlineImagePart);

final MimeBodyPart relatedWrapperPart = new MimeBodyPart();
relatedWrapperPart.setContent(relatedMultipart);
mixedMultipart.addBodyPart(relatedWrapperPart);

final MimeBodyPart pdfAttachmentPart = new MimeBodyPart();
pdfAttachmentPart.setContent("JVBERi0xLjQKJdPr6gkwChMKMSAwIG9iagogIDw8IC9UeXBlIC9DYXRhbG9n...", "application/pdf; name=\"my-attachment.pdf\"");
pdfAttachmentPart.setDescription("my-attachment.pdf");
pdfAttachmentPart.setDisposition("attachment; filename=\"my-attachment.pdf\"");
pdfAttachmentPart.setHeader("Content-Transfer-Encoding", "base64");
pdfAttachmentPart.setHeader("Content-Disposition", getPdfContentDisposition());
mixedMultipart.addBodyPart(pdfAttachmentPart);
message.setContent(mixedMultipart);

final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
message.writeTo(outputStream);

return outputStream.toByteArray();
}

private String getPdfContentDisposition() {
return """
attachment;
filename="my-attachment.pdf"; size=71521;
creation-date="Thu, 13 Aug 2026 11:02:50 GMT";
modification-date="Thu, 13 Aug 2026 11:01:24 GMT\"""";
}

private void assertAttachmentHeaderAttributes(MockFlowFile split, Map<String, String> expected) {
for (Map.Entry<String, String> entry : expected.entrySet()) {
// Must account for jakarta.mail.internet.MimeBodyPart writing MIME headers using canonical CRLF (\r\n)
split.assertAttributeEquals(entry.getKey(), entry.getValue().replace("\n", "\r\n"));
}
}
}
Loading