1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 package org.htmlunit.attachment;
16
17 import java.io.IOException;
18 import java.io.InputStream;
19 import java.nio.file.Files;
20 import java.nio.file.LinkOption;
21 import java.nio.file.Path;
22
23 import org.apache.commons.io.FileUtils;
24 import org.apache.commons.logging.Log;
25 import org.apache.commons.logging.LogFactory;
26 import org.htmlunit.Page;
27 import org.htmlunit.WebResponse;
28 import org.htmlunit.util.StringUtils;
29
30
31
32
33
34
35
36
37
38
39
40
41
42 public class DownloadingAttachmentHandler implements AttachmentHandler {
43
44 private static final Log LOG = LogFactory.getLog(DownloadingAttachmentHandler.class);
45
46 private final Path downloadFolder_;
47
48
49
50
51
52
53
54
55 public DownloadingAttachmentHandler(final Path downloadFolder) throws IOException {
56 downloadFolder_ = downloadFolder;
57 if (Files.notExists(downloadFolder)) {
58 throw new IOException("The provided download folder '"
59 + downloadFolder + "' does not exist");
60 }
61 if (!Files.isWritable(downloadFolder)) {
62 throw new IOException("Can't write to the download folder '"
63 + downloadFolder + "'");
64 }
65 }
66
67
68
69
70
71
72
73 public DownloadingAttachmentHandler() throws IOException {
74 this(Path.of(System.getProperty("java.io.tmpdir")));
75 }
76
77
78
79
80 @Override
81 public void handleAttachment(final Page page, final String attachmentFilename) {
82 final Path destination = determineDestionationFile(page, attachmentFilename);
83
84 final WebResponse webResponse = page.getWebResponse();
85 try (InputStream contentAsStream = webResponse.getContentAsStream()) {
86 FileUtils.copyToFile(contentAsStream, destination.toFile());
87 }
88 catch (final Exception e) {
89 LOG.error("Failed to write attachment response content to '"
90 + destination.toAbsolutePath() + "'");
91 return;
92 }
93
94 webResponse.cleanUp();
95 }
96
97 private Path determineDestionationFile(final Page page, final String attachmentFilename) {
98 String fileName = attachmentFilename;
99
100 if (StringUtils.isBlank(fileName)) {
101 final String file = page.getWebResponse().getWebRequest().getUrl().getFile();
102 fileName = file.substring(file.lastIndexOf('/') + 1);
103 }
104
105 if (StringUtils.isBlank(fileName)) {
106 fileName = "download";
107 }
108
109 Path newPath = downloadFolder_.resolve(fileName);
110 int count = 1;
111 while (Files.exists(newPath, LinkOption.NOFOLLOW_LINKS)) {
112 final String newFileName;
113 final int pos = fileName.lastIndexOf('.');
114 if (pos == -1) {
115 newFileName = fileName + "(" + count + ")";
116 }
117 else {
118 newFileName = fileName.substring(0, pos)
119 + "(" + count + ")"
120 + fileName.substring(pos);
121 }
122
123 newPath = downloadFolder_.resolve(newFileName);
124 count++;
125 }
126
127 return newPath.toAbsolutePath();
128 }
129 }