1   import com.trendmicro.grid.acl.client.ServiceClient;
2   import com.trendmicro.grid.acl.client.util.CommandlineClientFactory;
3   import com.trendmicro.grid.acl.client.util.CommandlineParser;
4   import com.trendmicro.grid.acl.client.util.ProgressIndicator;
5   import com.trendmicro.grid.acl.l0.AuthenticationException;
6   import com.trendmicro.grid.acl.l0.BatchCollection;
7   import com.trendmicro.grid.acl.l0.FileRequestService;
8   import com.trendmicro.grid.acl.l0.FileService;
9   import com.trendmicro.grid.acl.l0.datatypes.*;
10  import com.trendmicro.grid.acl.metadata.Meta;
11  import com.trendmicro.grid.acl.metadata.Metadata;
12  import net.sf.tinyjee.util.Hex;
13  
14  import java.io.*;
15  import java.net.HttpURLConnection;
16  import java.net.URL;
17  import java.text.DateFormat;
18  import java.text.SimpleDateFormat;
19  import java.util.*;
20  import java.util.concurrent.ExecutorService;
21  import java.util.concurrent.Executors;
22  import java.util.concurrent.TimeUnit;
23  
24  import static com.trendmicro.grid.acl.l0.datatypes.DaysRange.RangeType;
25  import static java.lang.System.out;
26  
27  /**
28   * Queries a list of files that are tagged with the specified tags and are
29   * included inside the optional time range.
30   *
31   * @author juergen_kellerer, 2010-07-01
32   * @version 1.0
33   */
34  public class FileQuery {
35  
36  	static {
37  		final String hexStringKey = "gacl.meta.print.binary.as.hex.string";
38  		if (System.getProperty(hexStringKey) == null)
39  			System.setProperty(hexStringKey, Boolean.TRUE.toString());
40  	}
41  
42  	static final ExecutorService executorService =
43  			Executors.newFixedThreadPool(Integer.getInteger("max.connections", 4));
44  
45  	static final String[] FILE_NAME_META_NAMES = {
46  			"originalFileName", "contentFileName", "downloadName", "filename", "name"};
47  	static final String[] FILE_SIZE_META_NAMES = {
48  			"contentLength", "originalFileSize", "fileSize", "size"};
49  
50  	static final String[] HELP = {
51  			"GRID File Query - Version 1.0",
52  			"",
53  			"Queries a list of files that match the specified tag query and are",
54  			"included inside an optional time range.",
55  			"",
56  			"Optionally transfers the files and meta information to a dedicated",
57  			"target folder.",
58  			"",
59  	};
60  
61  	static final CommandlineParser CLI = new CommandlineParser().
62  
63  			defineParameter("Defines the tag query to use for identifying the files.\n" +
64  					"Tag queries are of the format:\n" +
65  					"  \"tag01 tag01 -excludedTag01 -excludedTag02\"\n" +
66  					"  \"(g1tag -g1extag) (g2tag -g2extag)\"",
67  					true, String.class, null, "-q", "--query").
68  
69  			defineParameter("Defines an optional 'from' date using YYYY-MM-DD.",
70  					false, String.class, null, "-f", "--from").
71  			defineParameter("Defines an optional 'to' date using YYYY-MM-DD.",
72  					false, String.class, null, "-t", "--to").
73  			defineParameter("Defines how to apply dates, can be one of: " +
74  					"'FIRST_SEEN', 'LAST_RETRIEVED' or 'LAST_PROCESSED'",
75  					false, String.class, "FIRST_SEEN", "--range-type").
76  
77  			defineParameter("Copies the file content into the specified target folder.",
78  					false, File.class, null, "-c", "--copyto").
79  			defineParameter("Creates a hash list of all matched files and writes it to the given text file.",
80  					false, File.class, null, "--dump-hashes-to").
81  			defineParameter("Append the specified extension to copied files.", false, String.class,
82  					null, "-e", "--extension").
83  			defineParameter("Defines the depth of the folder hierachy in the output " +
84  					"directory (0 forces a flat hierachy).", false, Integer.class,
85  					6, "-d", "--depth").
86  			defineSwitchParameter("Copy with original filenames (if available in the file details).",
87  					"-n", "--original-names").
88  
89  			defineSwitchParameter("Dumps the file details along with the content. " +
90  					"When enabled, details are stored inside a file called like the content file " +
91  					"but extended with '.xml' and located in the same folder.",
92  					"--dump-details").
93  
94  			defineSwitchParameter("Be verbose.", "-v", "--verbose");
95  
96  	static final CommandlineClientFactory CLIENT_FACTORY = new CommandlineClientFactory(CLI);
97  	static PrintWriter hashListWriter;
98  
99  	public static void main(String[] args) throws Exception {
100 		boolean canContinue = CLI.parse(args, out, HELP);
101 		if (!canContinue)
102 			return;
103 
104 		final ServiceClient client = CLIENT_FACTORY.newServiceClient(CLI);
105 
106 		final File hashListFile = CLI.getParameter("--dump-hashes-to", File.class);
107 		if (hashListFile != null)
108 			hashListWriter = new PrintWriter(hashListFile);
109 
110 		long count = 0;
111 		final long time = System.currentTimeMillis();
112 		final ProgressIndicator indicator = new ProgressIndicator("");
113 		new Thread(indicator).start();
114 		try {
115 			final DaysRange range;
116 			String fromDate = CLI.getParameter("-f", String.class), toDate = CLI.getParameter("-t", String.class);
117 			if (fromDate != null || toDate != null) {
118 				final DateFormat df = new SimpleDateFormat("yyyy-MM-dd");
119 				final RangeType type = RangeType.valueOf(CLI.getParameter("--range-type", String.class));
120 				range = new DaysRange(type,
121 						fromDate == null ? null : df.parse(fromDate),
122 						toDate == null ? null : df.parse(toDate));
123 			} else
124 				range = null;
125 
126 			count = queryFiles(client, CLI.getParameter("-q", String.class), range, indicator);
127 		} finally {
128 			executorService.shutdown();
129 			while (!executorService.awaitTermination(30, TimeUnit.SECONDS))
130 				indicator.println("Waiting on pending transfers.");
131 
132 			indicator.close();
133 
134 			out.println();
135 			out.printf("Files: %d, Time: %.3f sec %n", count, ((double) (System.currentTimeMillis() - time)) / 1000D);
136 
137 			if (hashListWriter != null)
138 				hashListWriter.close();
139 		}
140 	}
141 
142 	static long queryFiles(ServiceClient client, String tagQuery, DaysRange range, ProgressIndicator indicator)
143 			throws AuthenticationException {
144 		final FileService fileService = client.getPort(FileService.class);
145 
146 		int pageNumber = 0;
147 		long totalCount = 0;
148 		FileIdentiferListPage idPage;
149 		int maxClusterSize = Integer.getInteger("max.chunksize", 16);
150 
151 		do {
152 			idPage = range == null ?
153 					fileService.getMatchingFiles(tagQuery, "1.0", pageNumber) :
154 					fileService.getMatchingFilesInRange(tagQuery, "1.0", range, pageNumber);
155 			if (idPage == null)
156 				break;
157 			pageNumber++;
158 
159 			IdHandler fh = null;
160 			for (FileIdentifier identifier : idPage.getElements()) {
161 				while (fh == null || !fh.addIfSpaceAvailable(identifier, maxClusterSize)) {
162 					if (fh != null)
163 						executorService.execute(fh);
164 					fh = new IdHandler(client, indicator);
165 				}
166 			}
167 
168 			if (fh != null)
169 				executorService.execute(fh);
170 
171 			totalCount += idPage.getElements().size();
172 
173 		} while (!idPage.isLastPage());
174 
175 		return totalCount;
176 	}
177 
178 	private static class IdHandler implements Runnable {
179 
180 		static final File targetFolder = CLI.getParameter("-c", File.class);
181 		static final String fileExtension = CLI.getParameter("-e", String.class);
182 		static final int pathLength = CLI.getParameter("-d", Integer.class);
183 		static final boolean printDetails = CLI.isParameterTrue("-v"),
184 				dumpDetails = CLI.isParameterTrue("--dump-details");
185 		static final boolean useOriginalFilenames = CLI.isParameterTrue("-n");
186 
187 		final private List<FileIdentifier> ids = new ArrayList<FileIdentifier>();
188 
189 		final ServiceClient client;
190 		final ProgressIndicator indicator;
191 
192 		private FileService fileService;
193 		private FileRequestService fileRequestService;
194 
195 		private IdHandler(ServiceClient client, ProgressIndicator indicator) {
196 			this.client = client;
197 			this.indicator = indicator;
198 		}
199 
200 		boolean addIfSpaceAvailable(FileIdentifier id, int maxSize) {
201 			return ids.size() < maxSize && ids.add(id);
202 		}
203 
204 		public FileService getFileService() {
205 			if (fileService == null) fileService = client.getPort(FileService.class);
206 			return fileService;
207 		}
208 
209 		public FileRequestService getRequestService() {
210 			if (fileRequestService == null) fileRequestService = client.getPort(FileRequestService.class);
211 			return fileRequestService;
212 		}
213 
214 		/**
215 		 * {@inheritDoc}
216 		 */
217 		@SuppressWarnings("unchecked")
218 		public void run() {
219 			DateFormat dateFormat = null;
220 
221 			try {
222 				final Iterator<FileDetails> detailsIterator;
223 				if (printDetails || dumpDetails || useOriginalFilenames)
224 					detailsIterator = getFileService().getFileDetailsList(BatchCollection.of(ids)).iterator();
225 				else
226 					detailsIterator = ((List) Collections.nCopies(ids.size(), null)).iterator();
227 
228 				final Iterator<URL> urlIterator;
229 				if (targetFolder != null)
230 					urlIterator = getRequestService().getRemoteTransferURLs(BatchCollection.of(ids)).iterator();
231 				else
232 					urlIterator = ((List) Collections.nCopies(ids.size(), null)).iterator();
233 
234 				for (FileIdentifier id : ids) {
235 					final URL url = urlIterator.next();
236 					final FileDetails details = detailsIterator.next();
237 
238 					String filename = details == null || !useOriginalFilenames ? null :
239 							findFirst(details.getMetadata(), new Meta("").value((String) null),
240 									FILE_NAME_META_NAMES).getValue();
241 
242 					if (filename == null)
243 						filename = fileExtension == null ? "content.data" : "content" + fileExtension;
244 					else if (fileExtension != null)
245 						filename += fileExtension;
246 
247 					dumpIdentifierToList(id, filename);
248 
249 					synchronized (indicator) {
250 						indicator.println("File sha1=" + hexEncode(id.getSHA1Hash()) +
251 								", md5=" + hexEncode(id.getMD5Hash()) + ", filename=" + filename);
252 
253 						if (printDetails && details != null) {
254 							if (dateFormat == null)
255 								dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSZ");
256 							printDetails(details, dateFormat, indicator);
257 						}
258 					}
259 
260 					final File target = targetFolder == null ? null :
261 							computeOutputFile(targetFolder, id.getSHA1Hash(), pathLength, filename);
262 					if (target != null)
263 						transfer(url, details, target);
264 				}
265 			} catch (AuthenticationException e) {
266 				indicator.println(" " + e.getMessage());
267 				throw new RuntimeException(e);
268 			} catch (IOException e) {
269 				indicator.println(" " + e.getMessage());
270 				throw new RuntimeException(e);
271 			}
272 		}
273 
274 		void transfer(URL url, FileDetails details, File target) throws IOException {
275 			if (url == null) {
276 				indicator.println("Failed to find the underlying content for the requested file.");
277 				return;
278 			}
279 
280 			long totalLength = details == null ? 0 : asLong(
281 					findFirst(details.getMetadata(), new Meta("").value(0), FILE_SIZE_META_NAMES));
282 			transferIfNewer(url, target, totalLength, indicator, "Downloading: '" + target.getName() + "'");
283 
284 			if (dumpDetails)
285 				dumpDetails(details, target);
286 		}
287 	}
288 
289 	static void dumpIdentifierToList(FileIdentifier identifier, String filename) {
290 		if (hashListWriter != null)
291 			try {
292 				StringBuilder b = Hex.encode(new StringBuilder(128).append("SHA1:"), identifier.getSHA1Hash());
293 				byte[] md5Hash = identifier.getMD5Hash();
294 				if (md5Hash != null)
295 					Hex.encode(b.append(" MD5:"), md5Hash);
296 				b.append(" */").append(filename);
297 				hashListWriter.println(b.toString());
298 			} catch (IOException e) {
299 				throw new RuntimeException(e);
300 			}
301 	}
302 
303 	static void dumpDetails(FileDetails details, File forTarget) throws IOException {
304 		try {
305 			File target = new File(forTarget + ".xml");
306 			FileDetails.getXmlSerializer().save(details, target, false, 4);
307 		} catch (IOException e) {
308 			throw e;
309 		} catch (Exception e) {
310 			throw new IOException(e);
311 		}
312 	}
313 
314 	static void printDetails(FileDetails details, DateFormat dateFormat, ProgressIndicator out) {
315 		FileInformation info = details.getInformation();
316 		out.println("  - " + details.getIdentifier());
317 		out.println("  - Tags: " + info.getTags());
318 		out.println("  - Located at '" + info.getSourceSiteCount() +
319 				"' Sites. Contained in '" + info.getSourcePackageCount() + "' packages.");
320 		out.println("  - First Seen: " + formatDate(dateFormat, info.getFirstSeen()));
321 		out.println("  - Last Retrieved & Processed: " +
322 				formatDate(dateFormat, info.getLastRetrieved()) + " / " +
323 				formatDate(dateFormat, info.getLastProcessed()));
324 
325 		out.println("  - Metadata:");
326 		if (details.getMetadata() == null)
327 			out.println("      N/A");
328 		else
329 			for (Meta meta : details.getMetadata().getMetaElements())
330 				out.println("      " + meta);
331 	}
332 
333 	static String formatDate(DateFormat dateFormat, Date date) {
334 		return date == null ? "N/A" : dateFormat.format(date);
335 	}
336 
337 	static String hexEncode(byte[] b) {
338 		return b == null ? "N/A" : Hex.encode(b);
339 	}
340 
341 	static long asLong(Meta meta) {
342 		double[] nv = meta.getNumericValues();
343 		if (nv != null && nv.length > 0)
344 			return (long) nv[0];
345 		return Long.parseLong(meta.getValue());
346 	}
347 
348 	static Meta findFirst(Metadata metadata, Meta defaultValue, String... names) {
349 		if (metadata != null)
350 			for (String name : names) {
351 				Meta m = metadata.get(name);
352 				if (m != null)
353 					return m;
354 			}
355 		return defaultValue;
356 	}
357 
358 	static File computeOutputFile(File targetFolder, byte[] hash, int pathLen, String filename) {
359 		final int increment = hash.length / Math.max(1, pathLen);
360 		final StringBuilder expectedPath = new StringBuilder((hash.length * 2) + filename.length());
361 		try {
362 			for (int i = 0; i < hash.length; i += increment) {
363 				if (i != 0)
364 					expectedPath.append('/');
365 				Hex.encode(expectedPath, hash, i, Math.min(increment, hash.length - i));
366 			}
367 			expectedPath.append('/').append(filename);
368 		} catch (IOException e) {
369 			throw new RuntimeException(e);
370 		}
371 
372 		final File target = new File(targetFolder, expectedPath.toString());
373 		final File parentFile = target.getParentFile();
374 		if (!parentFile.isDirectory() && !parentFile.mkdirs())
375 			throw new RuntimeException("Failed creating directory: " + parentFile);
376 
377 		return target;
378 	}
379 
380 	static void transferIfNewer(URL source, File target,
381 								long totalLength, ProgressIndicator indicator, String msg) throws IOException {
382 		HttpURLConnection huc = (HttpURLConnection) source.openConnection();
383 		huc.setIfModifiedSince(target.lastModified());
384 
385 		switch (huc.getResponseCode()) {
386 			case HttpURLConnection.HTTP_NOT_MODIFIED:
387 				indicator.println(msg + " (not modified)");
388 				break;
389 
390 			case HttpURLConnection.HTTP_OK:
391 				if (totalLength <= 0)
392 					totalLength = huc.getContentLength();
393 				long bytesCount = copy(huc.getInputStream(),
394 						new BufferedOutputStream(new FileOutputStream(target)),
395 						totalLength, indicator, msg);
396 
397 				indicator.setProgress(1, 1);
398 				indicator.println(String.format("%s (ok -  %.2f mb)", msg, (bytesCount / 1024D / 1024D)));
399 				break;
400 
401 			case HttpURLConnection.HTTP_NOT_FOUND:
402 				indicator.println(msg + " ERROR: (not found in repository)");
403 				break;
404 
405 			default:
406 				throw new IOException("Server replied with " + huc.getResponseCode() + " " +
407 						huc.getResponseMessage() + " on the attempt to transfer " + source);
408 		}
409 	}
410 
411 	static long copy(InputStream in, OutputStream out,
412 					 long totalLength, ProgressIndicator indicator, String msg) throws IOException {
413 		int r;
414 		long bytesCount = 0;
415 		byte[] buf = new byte[1024 * 8];
416 		try {
417 			while ((r = in.read(buf)) != -1) {
418 				if (out != null)
419 					out.write(buf, 0, r);
420 				bytesCount += r;
421 				indicator.setMessage(msg);
422 				if (totalLength != -1)
423 					indicator.setProgress(bytesCount, totalLength);
424 			}
425 
426 			return bytesCount;
427 		} finally {
428 			in.close();
429 			if (out != null)
430 				out.close();
431 		}
432 	}
433 }