diff --git a/stirling-pdf/src/main/java/stirling/software/SPDF/service/PDFAttachmentService.java b/stirling-pdf/src/main/java/stirling/software/SPDF/service/PDFAttachmentService.java index 4b13054aa..d7db7c696 100644 --- a/stirling-pdf/src/main/java/stirling/software/SPDF/service/PDFAttachmentService.java +++ b/stirling-pdf/src/main/java/stirling/software/SPDF/service/PDFAttachmentService.java @@ -14,6 +14,7 @@ import org.apache.pdfbox.pdmodel.PageMode; import org.apache.pdfbox.pdmodel.common.filespecification.PDComplexFileSpecification; import org.apache.pdfbox.pdmodel.common.filespecification.PDEmbeddedFile; import org.apache.pdfbox.pdmodel.encryption.AccessPermission; +import org.apache.pdfbox.pdmodel.encryption.StandardProtectionPolicy; import org.springframework.stereotype.Service; import org.springframework.web.multipart.MultipartFile; @@ -35,6 +36,7 @@ public class PDFAttachmentService implements PDFAttachmentServiceInterface { try { existingNames = embeddedFilesTree.getNames(); + if (existingNames == null) { log.debug("No existing embedded files found, creating new names map."); existingNames = new HashMap<>(); @@ -46,7 +48,9 @@ public class PDFAttachmentService implements PDFAttachmentServiceInterface { throw e; } + grantAccessPermissions(document); final Map existingEmbeddedFiles = existingNames; + attachments.forEach( attachment -> { String filename = attachment.getOriginalFilename(); @@ -58,7 +62,6 @@ public class PDFAttachmentService implements PDFAttachmentServiceInterface { embeddedFile.setCreationDate(new GregorianCalendar()); embeddedFile.setModDate(new GregorianCalendar()); String contentType = attachment.getContentType(); - if (StringUtils.isNotBlank(contentType)) { embeddedFile.setSubtype(contentType); } @@ -84,7 +87,6 @@ public class PDFAttachmentService implements PDFAttachmentServiceInterface { }); embeddedFilesTree.setNames(existingNames); - grantAccessPermissions(document); PDFAttachmentUtils.setCatalogViewerPreferences(document, PageMode.USE_ATTACHMENTS); ByteArrayOutputStream output = new ByteArrayOutputStream(); document.save(output); @@ -93,17 +95,31 @@ public class PDFAttachmentService implements PDFAttachmentServiceInterface { } private void grantAccessPermissions(PDDocument document) { - AccessPermission currentPermissions = document.getCurrentAccessPermission(); + try { + AccessPermission currentPermissions = document.getCurrentAccessPermission(); - currentPermissions.setCanAssembleDocument(true); - currentPermissions.setCanFillInForm(currentPermissions.canFillInForm()); - currentPermissions.setCanModify(true); - currentPermissions.setCanPrint(true); - currentPermissions.setCanPrintFaithful(true); + currentPermissions.setCanAssembleDocument(true); + currentPermissions.setCanFillInForm(currentPermissions.canFillInForm()); + currentPermissions.setCanModify(true); + currentPermissions.setCanPrint(true); + currentPermissions.setCanPrintFaithful(true); - // Ensure these permissions are enabled for embedded file access - currentPermissions.setCanExtractContent(true); - currentPermissions.setCanExtractForAccessibility(true); - currentPermissions.setCanModifyAnnotations(true); + // Ensure these permissions are enabled for embedded file access + currentPermissions.setCanExtractContent(true); + currentPermissions.setCanExtractForAccessibility(true); + currentPermissions.setCanModifyAnnotations(true); + + var protectionPolicy = new StandardProtectionPolicy(null, null, currentPermissions); + + if (!document.isAllSecurityToBeRemoved()) { + document.setAllSecurityToBeRemoved(true); + } + + document.protect(protectionPolicy); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + document.save(output); + } catch (IOException e) { + throw new RuntimeException(e); + } } } diff --git a/stirling-pdf/src/test/java/stirling/software/SPDF/service/PDFAttachmentServiceTest.java b/stirling-pdf/src/test/java/stirling/software/SPDF/service/PDFAttachmentServiceTest.java index 8b1b676c0..af40f444c 100644 --- a/stirling-pdf/src/test/java/stirling/software/SPDF/service/PDFAttachmentServiceTest.java +++ b/stirling-pdf/src/test/java/stirling/software/SPDF/service/PDFAttachmentServiceTest.java @@ -7,6 +7,8 @@ import java.util.List; import org.apache.pdfbox.pdmodel.PDDocument; import org.apache.pdfbox.pdmodel.PDEmbeddedFilesNameTreeNode; import org.apache.pdfbox.pdmodel.common.filespecification.PDComplexFileSpecification; +import org.apache.pdfbox.pdmodel.encryption.AccessPermission; +import org.apache.pdfbox.pdmodel.encryption.StandardProtectionPolicy; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.springframework.web.multipart.MultipartFile; @@ -126,4 +128,89 @@ class PDFAttachmentServiceTest { verify(embeddedFilesTree).setNames(anyMap()); } } + + @Test + void addAttachmentToPDF_WithProtectedDocument() throws IOException { + try (var document = new PDDocument()) { + // Create a document with restricted permissions (this simulates an encrypted/protected document) + AccessPermission ap = new AccessPermission(); + ap.setCanExtractContent(false); // Restrict content extraction initially + var spp = new StandardProtectionPolicy("owner", "user", ap); + document.protect(spp); + + var embeddedFilesTree = mock(PDEmbeddedFilesNameTreeNode.class); + var attachments = List.of(mock(MultipartFile.class)); + var existingNames = new HashMap(); + + when(embeddedFilesTree.getNames()).thenReturn(existingNames); + when(attachments.get(0).getOriginalFilename()).thenReturn("test.txt"); + when(attachments.get(0).getInputStream()).thenReturn( + new ByteArrayInputStream("Test content".getBytes())); + when(attachments.get(0).getSize()).thenReturn(12L); + when(attachments.get(0).getContentType()).thenReturn("text/plain"); + + byte[] result = pdfAttachmentService.addAttachment(document, embeddedFilesTree, attachments); + + assertNotNull(result); + assertTrue(result.length > 0); + verify(embeddedFilesTree).setNames(anyMap()); + } + } + + @Test + void addAttachmentToPDF_WithRestrictedPermissions() throws IOException { + try (var document = new PDDocument()) { + // Create a document with very restricted permissions that should block permission changes + AccessPermission ap = new AccessPermission(); + ap.setCanModify(false); + ap.setCanAssembleDocument(false); + ap.setCanExtractContent(false); + var spp = new StandardProtectionPolicy("owner", "user", ap); + document.protect(spp); + + var embeddedFilesTree = mock(PDEmbeddedFilesNameTreeNode.class); + var attachments = List.of(mock(MultipartFile.class)); + var existingNames = new HashMap(); + + when(embeddedFilesTree.getNames()).thenReturn(existingNames); + when(attachments.get(0).getOriginalFilename()).thenReturn("test.txt"); + when(attachments.get(0).getInputStream()).thenReturn( + new ByteArrayInputStream("Test content".getBytes())); + when(attachments.get(0).getSize()).thenReturn(12L); + when(attachments.get(0).getContentType()).thenReturn("text/plain"); + + byte[] result = pdfAttachmentService.addAttachment(document, embeddedFilesTree, attachments); + + assertNotNull(result); + assertTrue(result.length > 0); + verify(embeddedFilesTree).setNames(anyMap()); + } + } + + @Test + void addAttachmentToPDF_WithNonEncryptedDocument() throws IOException { + try (var document = new PDDocument()) { + var embeddedFilesTree = mock(PDEmbeddedFilesNameTreeNode.class); + var attachments = List.of(mock(MultipartFile.class)); + var existingNames = new HashMap(); + + when(embeddedFilesTree.getNames()).thenReturn(existingNames); + when(attachments.get(0).getOriginalFilename()).thenReturn("test.txt"); + when(attachments.get(0).getInputStream()).thenReturn( + new ByteArrayInputStream("Test content".getBytes())); + when(attachments.get(0).getSize()).thenReturn(12L); + when(attachments.get(0).getContentType()).thenReturn("text/plain"); + + byte[] result = pdfAttachmentService.addAttachment(document, embeddedFilesTree, attachments); + + assertNotNull(result); + assertTrue(result.length > 0); + // Verify permissions are set correctly for non-encrypted documents + AccessPermission permissions = document.getCurrentAccessPermission(); + assertTrue(permissions.canExtractContent()); + assertTrue(permissions.canExtractForAccessibility()); + assertTrue(permissions.canModifyAnnotations()); + verify(embeddedFilesTree).setNames(anyMap()); + } + } }