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.l0.ProcessingService;
5 import com.trendmicro.grid.acl.l0.datatypes.*;
6 import com.trendmicro.grid.acl.metadata.Metadata;
7 import net.sf.tinyjee.streams.PositionOutputStream;
8
9 import java.io.File;
10 import java.io.FileInputStream;
11 import java.io.FileOutputStream;
12 import java.net.URI;
13 import java.util.*;
14 import java.util.zip.ZipEntry;
15 import java.util.zip.ZipOutputStream;
16
17 import static com.trendmicro.grid.acl.l0.datatypes.ProcessPackageDataSet.ProcessType.standard;
18 import static java.lang.System.out;
19
20
21
22
23
24
25
26 public class DataGenerator {
27
28 static final Random RANDOM = new Random();
29
30 static final List<String> FORBIDDEN_HOSTNAMES = Arrays.asList(System.getProperty("forbidden.hostnames",
31 "mygacl.trendmicro.com,gacl.trendmicro.com").split("[\\s,;]+"));
32
33 static final String[] HELP = {
34 "GRID Symmetric Data Generator - Version 1.0",
35 "",
36 "Generates profile based data, using a symmetric algorithm to build the hierachy.",
37 "Uses the processing result services offered by the GRID Access Layer (GACL) to ",
38 "store processing results.",
39 "",
40 };
41
42 static final CommandlineParser CLI = new CommandlineParser().
43 defineParameter("Dumps the built-in profile to the specified file,\n " +
44 "allowing its customization.",
45 false, File.class, null, "-d", "--dump").
46 defineParameter("Dumps the generated datasets to the specified ZIP file,\n " +
47 "allowing its their re-usage in inside the request-autotester.",
48 false, File.class, null, "--dump-datasets").
49 defineParameter("Use the specified profile instead of the built-in.",
50 false, File.class, null, "-p", "--profile").
51
52 defineParameter("Specifies the batch size used for sending\n" +
53 "packages.",
54 false, Integer.class, 5, "-b", "--batch").
55
56 defineSwitchParameter("Be verbose", "-v", "--verbose");
57
58 static final CommandlineClientFactory CLIENT_FACTORY = new CommandlineClientFactory(CLI);
59
60
61 public static void main(String[] args) throws Exception {
62 boolean canContinue = CLI.parse(args, out, HELP);
63 if (!canContinue)
64 return;
65
66 File profileDumpFile = CLI.getParameter("-d", File.class);
67 if (profileDumpFile != null) {
68 dumpProfile(profileDumpFile);
69 return;
70 }
71
72 Properties profile = loadProfile();
73 applyProfile(profile);
74
75 final ServiceClient client = CLIENT_FACTORY.newServiceClient(CLI);
76 if (FORBIDDEN_HOSTNAMES.contains(client.getDefaultHost().toLowerCase())) {
77 throw new IllegalArgumentException("The hostname '" +
78 client.getDefaultHost() + "' is not allowed to be used.");
79 }
80
81 final long time = System.currentTimeMillis();
82 try {
83 generateData(client);
84 } finally {
85 out.println();
86 out.printf("Time: %.3f sec %n", ((double) (System.currentTimeMillis() - time)) / 1000D);
87 }
88 }
89
90 static void dumpProfile(File to) throws Exception {
91 loadProfile().store(new FileOutputStream(to), "");
92 }
93
94 static Properties loadProfile() throws Exception {
95 Properties profile = new Properties();
96
97 File externalProfile = CLI.getParameter("-p", File.class);
98 if (externalProfile != null)
99 profile.load(new FileInputStream(externalProfile));
100 else
101 profile.load(DataGenerator.class.getResourceAsStream("/DataGenerator.properties"));
102
103 return profile;
104 }
105
106 static List<String> packageLocales = new ArrayList<String>();
107 static List<String> packageVersions = new ArrayList<String>();
108 static List<String> packageVendors = new ArrayList<String>();
109 static List<String> packageBaseNames = new ArrayList<String>();
110 static List<String> packageFolders = new ArrayList<String>();
111 static List<String> packageFiles = new ArrayList<String>();
112 static Set<String> detailedFiles = new HashSet<String>();
113 static String[] packageTags;
114 static Map<String, String[]> fileTags = new HashMap<String, String[]>();
115
116 static void applyProfile(Properties properties) {
117 for (Map.Entry<Object, Object> entry : properties.entrySet()) {
118 String key = entry.getKey().toString();
119 if (key.startsWith("fileTags."))
120 fileTags.put(key.substring(8), entry.getValue().toString().split("\\s+"));
121 }
122
123 detailedFiles.addAll(Arrays.asList(properties.getProperty("detailedFileExtensions").split("\\s+")));
124
125 packageFiles.addAll(Arrays.asList(properties.getProperty("packageFiles").split("\\s+")));
126 packageFolders.addAll(Arrays.asList(properties.getProperty("packageFolders").split("\\s+")));
127 packageTags = properties.getProperty("packageTags").split("\\s+");
128
129 packageVersions.addAll(Arrays.asList(properties.getProperty("packageVersions").split("\\s+")));
130 packageLocales.addAll(Arrays.asList(properties.getProperty("packageLocales").split("\\s+")));
131 packageVendors.addAll(Arrays.asList(properties.getProperty("packageVendors").split("\\s+")));
132 packageBaseNames.addAll(Arrays.asList(properties.getProperty("packageBaseNames").split("\\s+")));
133 }
134
135 static void generateData(ServiceClient client) throws Exception {
136 final boolean verbose = CLI.isParameterTrue("-v");
137 final int batchSize = CLI.getParameter("-b", Integer.class);
138 final ProcessingService processingService = client.getPort(ProcessingService.class);
139
140 final File dataSetsDumpFile = CLI.getParameter("--dump-datasets", File.class);
141 final ZipOutputStream dataSetDumpStream = dataSetsDumpFile == null ? null :
142 new ZipOutputStream(new FileOutputStream(dataSetsDumpFile));
143
144 int count = 0, filesCount = 0;
145 try {
146 List<ProcessPackageDataSet> dataSets = new ArrayList<ProcessPackageDataSet>();
147 for (String packageVendor : packageVendors) {
148 for (String packageBaseName : packageBaseNames) {
149 for (String version : packageVersions) {
150 for (String locale : packageLocales) {
151 String packageName = packageVendor + ":" + packageBaseName + ":" + version + ":" + locale;
152 dataSets.add(generateDataSet(packageName, packageVendor));
153
154 if (dataSets.size() > batchSize) {
155 processingService.storeProcessingResultsAndFinalizeJobs(dataSets);
156
157 for (ProcessPackageDataSet dataSet : dataSets) {
158 count++;
159 filesCount += dataSet.getPackageMember().size();
160 if (verbose) {
161 System.out.println("Stored: " + dataSet);
162 } else {
163 System.out.println("Stored: Package " +
164 dataSet.getProcessedPackage().getPackageInformation().getName());
165 }
166 }
167
168 if (dataSetDumpStream != null)
169 dumpDataSets(dataSetDumpStream, dataSets);
170
171 dataSets.clear();
172 }
173 }
174 }
175 }
176 }
177 } finally {
178 if (dataSetDumpStream != null)
179 dataSetDumpStream.close();
180 }
181
182 System.out.println("Added " + count + " packages, containing " + filesCount + " files");
183 }
184
185 static void dumpDataSets(ZipOutputStream zOut, List<ProcessPackageDataSet> results) throws Exception {
186 PositionOutputStream pOut = new PositionOutputStream(zOut, false);
187 for (ProcessPackageDataSet result : results) {
188 zOut.putNextEntry(new ZipEntry(result.getReferringJob().toString() + ".dataset.xml"));
189 ProcessPackageDataSet.getXmlSerializer().save(result, pOut);
190 pOut.close();
191 zOut.closeEntry();
192 }
193 }
194
195 static ProcessPackageDataSet generateDataSet(String packageName, String companyName) {
196 Collection<PackageMember> members = new ArrayList<PackageMember>();
197 for (String packageFolder : packageFolders) {
198 for (String packageFile : packageFiles)
199 members.add(generateMember(packageFolder + "/" + packageFile, companyName));
200 }
201
202 PackageDetails pd = generatePackage(packageName, companyName, members);
203 Source source = generateSource(packageName, pd.getFileMetadata().getIdentifier(), companyName);
204
205 ProcessPackageDataSet dataSet = new ProcessPackageDataSet(
206 UUID.randomUUID(), standard, Collections.singleton(source));
207 dataSet.setProcessedPackage(pd);
208 dataSet.getPackageMember().addAll(members);
209
210 return dataSet;
211 }
212
213 static PackageDetails generatePackage(String packageName, String companyName, Collection<PackageMember> members) {
214 Set<String> tags = new HashSet<String>();
215 for (PackageMember member : members) {
216 FileInformation info = member.getInformation();
217 if (info == null)
218 info = member.getDetails().getInformation();
219
220 tags.addAll(info.getTags());
221 }
222
223 tags.addAll(Arrays.asList(packageTags));
224
225 String vendorName = extractName(packageName, 1);
226 String familyName = extractName(packageName, 2);
227 FileDetails fileDetails = generateFile("setup.exe", companyName);
228 PackageInformation info = new PackageInformation(packageName, packageName, familyName, vendorName,
229 tags.toArray(new String[tags.size()]), fileDetails.getInformation());
230
231 tags.add(packageName);
232 tags.add(vendorName);
233 tags.add(familyName);
234
235 Metadata metadata = new Metadata();
236 metadata.add("scanResult.SCANNER_A").value("passed");
237 metadata.add("scanResult.SCANNER_B").value("passed");
238 metadata.add("scanResult.SCANNER_C").value("passed");
239 metadata.add("scanResult.SCANNER_D").value("passed");
240 metadata.add("scanResult.SCANNER_E").value("passed");
241
242 return new PackageDetails(null, info, metadata,
243 new FileMetadata(fileDetails.getIdentifier(), fileDetails.getMetadata()));
244 }
245
246 static Source generateSource(String packageName, FileIdentifier packageFileId, String companyName) {
247 URI remoteURI = URI.create("http://" + companyName + "/generated-data/" + packageName);
248 SourceInformation si = new SourceInformation(new SourceIdentifier(remoteURI, null), new Date(), null, false);
249 Metadata m = new Metadata();
250 m.add("contentName").value(packageName);
251 m.add("contentSHA1").value(packageFileId.getSHA1Hash());
252 m.add("contentMD5").value(packageFileId.getMD5Hash());
253 return new Source(remoteURI, null, si, null, m);
254 }
255
256 static PackageMember generateMember(String name, String companyName) {
257 FileDetails fd = generateFile(name, companyName);
258 NamedFileIdentifier id = new NamedFileIdentifier(fd.getIdentifier(), name);
259 if (detailedFiles.contains(extractExtension(name)))
260 return new PackageMember(id, fd);
261 return new PackageMember(id, fd.getInformation());
262 }
263
264 static FileDetails generateFile(String name, String companyName) {
265 FileIdentifier id = new FileIdentifier(create(20), create(16));
266 String[] tags = fileTags.get(extractExtension(name));
267 if (tags == null)
268 tags = fileTags.get("fileTags.");
269
270 Calendar cal = Calendar.getInstance();
271 Date firstSeen = cal.getTime();
272 cal.add(Calendar.HOUR, 1);
273 Date lastRetrieved = cal.getTime();
274 cal.add(Calendar.HOUR, 1);
275 Date lastProcessed = cal.getTime();
276
277 FileInformation info = new FileInformation(firstSeen, lastRetrieved, lastProcessed, tags);
278
279 Metadata metadata = new Metadata();
280 metadata.add("companyName").value(companyName);
281 metadata.add("originalFileName").value(name);
282 metadata.add("internalName").value(companyName);
283 metadata.add("fileSize").value(1024 * 6 + RANDOM.nextInt(1024 * 1024 * 64));
284 metadata.add("sha256").value(create(256 / 8));
285 metadata.add("sha512").value(create(512 / 8));
286
287 return new FileDetails(id, info, metadata);
288 }
289
290 static String extractExtension(String name) {
291 int idx = name.indexOf('.');
292 return idx == -1 ? "." : name.substring(idx);
293 }
294
295 static byte[] create(int len) {
296 byte[] b = new byte[len];
297 RANDOM.nextBytes(b);
298 return b;
299 }
300
301 static String extractName(String name, int delimiterCount) {
302 int endIdx = -1;
303 for (int i = 0; i < delimiterCount; i++)
304 endIdx = name.indexOf(':', endIdx + 1);
305
306 return endIdx == -1 ? name : name.substring(0, endIdx);
307 }
308 }