1 import net.sf.tinyjee.util.Hex;
2
3 import java.io.*;
4 import java.net.URL;
5 import java.security.DigestInputStream;
6 import java.security.MessageDigest;
7 import java.util.concurrent.ExecutorService;
8 import java.util.concurrent.Executors;
9 import java.util.concurrent.TimeUnit;
10 import java.util.concurrent.atomic.AtomicLong;
11
12 import static java.lang.System.err;
13 import static java.lang.System.out;
14
15
16
17
18
19
20
21
22 public class TrivialFileChecker {
23
24 public static final String TITLE = "GRID - Trivial File Checker, Version 1.0";
25
26 public static final String PROTOCOL = System.getProperty("protocol", "https");
27 public static final String HOSTNAME = System.getProperty("host", "gacl.trendmicro.com");
28 public static final int CONNECTIONS = Integer.getInteger("connections", 16);
29
30 private static String endpointAddress = PROTOCOL + "://" + HOSTNAME + "/rs/level-0/files/isKnownGood/";
31 private static ExecutorService executorService = Executors.newFixedThreadPool(CONNECTIONS);
32 private static AtomicLong checkCount = new AtomicLong(0);
33
34 public static void main(String[] args) throws Exception {
35 out.println(TITLE);
36 out.println(" using endpoint: " + endpointAddress);
37 out.println(" (change using -Dhost=ipOrHost or -Dprotocol=http)");
38 out.println();
39
40 System.setProperty("http.maxConnections", Integer.toString(CONNECTIONS));
41
42 checkCount.set(0);
43 long time = System.currentTimeMillis();
44 try {
45 if (args.length == 0)
46 out.println("No files or directories specified.");
47 else {
48 for (String arg : args) {
49 File path = new File(arg);
50 if (!checkPath(path))
51 err.println("Ignoring invalid path: " + path);
52 }
53 }
54 } finally {
55 executorService.shutdown();
56 executorService.awaitTermination(60, TimeUnit.MINUTES);
57 time = System.currentTimeMillis() - time;
58 }
59
60 out.printf("%nChecked %d files in %.2f minutes (%d files/s)%n",
61 checkCount.get(), (double) time / 60D / 1000D, Math.round(checkCount.get() / (time / 1000D)));
62 }
63
64 private static boolean checkPath(File path) throws Exception {
65 if (path.isFile()) {
66 executorService.execute(new Worker(path));
67 } else if (path.isDirectory()) {
68 final File[] children = path.listFiles();
69 if (children != null) {
70 for (File child : children)
71 checkPath(child);
72 }
73 } else
74 return false;
75 return true;
76 }
77
78 private static class Worker implements Runnable {
79
80 private static ThreadLocal<byte[]> bufferPool = new ThreadLocal<byte[]>() {
81 @Override
82 protected byte[] initialValue() {
83 return new byte[512 * 1024];
84 }
85 };
86
87 private static byte[] calculateSHA1(File file) throws Exception {
88 final byte[] buffer = bufferPool.get();
89 final DigestInputStream in = new DigestInputStream(
90 new FileInputStream(file), MessageDigest.getInstance("SHA1"));
91 try {
92 while (in.read(buffer) != -1) ;
93 return in.getMessageDigest().digest();
94 } finally {
95 in.close();
96 }
97 }
98
99 private static Boolean checkFile(byte[] sha1) throws Exception {
100 URL ep = new URL(endpointAddress + Hex.encode(sha1));
101 BufferedReader reader = new BufferedReader(new InputStreamReader((InputStream) ep.getContent()), 16);
102 try {
103 String result = reader.readLine();
104 return result == null ? null : Boolean.valueOf(result);
105 } finally {
106 reader.close();
107 }
108 }
109
110 private File file;
111
112 private Worker(File file) {
113 this.file = file;
114 }
115
116 public void run() {
117 try {
118 Boolean result = checkFile(calculateSHA1(file));
119 checkCount.incrementAndGet();
120 out.printf("%7s> %s%n", (result == null ? "?" : (result ? " GOOD" : "BAD ")), file);
121 } catch (Exception e) {
122 synchronized (err) {
123 err.println("Failed checking: " + file);
124 e.printStackTrace();
125 }
126 }
127 }
128 }
129 }