| /* |
| * Copyright (C) 2014 The Android Open Source Project |
| * |
| * Licensed under the Apache License, Version 2.0 (the "License"); |
| * you may not use this file except in compliance with the License. |
| * You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| package dexfuzz.listeners; |
| |
| import dexfuzz.Log; |
| import dexfuzz.Options; |
| import dexfuzz.executors.Executor; |
| |
| import java.io.File; |
| import java.io.FileInputStream; |
| import java.io.FileNotFoundException; |
| import java.io.FileOutputStream; |
| import java.io.IOException; |
| import java.io.ObjectInputStream; |
| import java.io.ObjectOutputStream; |
| import java.security.MessageDigest; |
| import java.security.NoSuchAlgorithmException; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| |
| /** |
| * Tracks unique programs and outputs. Also saves divergent programs! |
| */ |
| public class UniqueProgramTrackerListener extends BaseListener { |
| /** |
| * Map of unique program MD5 sums, mapped to times seen. |
| */ |
| private Map<String, Integer> uniquePrograms; |
| |
| /** |
| * Map of unique program outputs (MD5'd), mapped to times seen. |
| */ |
| private Map<String, Integer> uniqueOutputs; |
| |
| /** |
| * Used to remember the seed used to fuzz the fuzzed file, so we can save it with this |
| * seed as a name, if we find a divergence. |
| */ |
| private long currentSeed; |
| |
| /** |
| * Used to remember the name of the file we've fuzzed, so we can save it if we |
| * find a divergence. |
| */ |
| private String fuzzedFile; |
| |
| private MessageDigest digest; |
| private String databaseFile; |
| |
| /** |
| * Save the database every X number of iterations. |
| */ |
| private static final int saveDatabasePeriod = 20; |
| |
| public UniqueProgramTrackerListener(String databaseFile) { |
| this.databaseFile = databaseFile; |
| } |
| |
| @Override |
| public void handleSeed(long seed) { |
| currentSeed = seed; |
| } |
| |
| /** |
| * Given a program filename, calculate the MD5sum of |
| * this program. |
| */ |
| private String getMD5SumOfProgram(String programName) { |
| byte[] buf = new byte[256]; |
| try { |
| FileInputStream stream = new FileInputStream(programName); |
| boolean done = false; |
| while (!done) { |
| int bytesRead = stream.read(buf); |
| if (bytesRead == -1) { |
| done = true; |
| } else { |
| digest.update(buf); |
| } |
| } |
| stream.close(); |
| } catch (FileNotFoundException e) { |
| e.printStackTrace(); |
| } catch (IOException e) { |
| e.printStackTrace(); |
| } |
| return new String(digest.digest()); |
| } |
| |
| private String getMD5SumOfOutput(String output) { |
| digest.update(output.getBytes()); |
| return new String(digest.digest()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void loadUniqueProgsData() { |
| File file = new File(databaseFile); |
| if (!file.exists()) { |
| uniquePrograms = new HashMap<String, Integer>(); |
| uniqueOutputs = new HashMap<String, Integer>(); |
| return; |
| } |
| |
| try { |
| ObjectInputStream objectStream = |
| new ObjectInputStream(new FileInputStream(databaseFile)); |
| uniquePrograms = (Map<String, Integer>) objectStream.readObject(); |
| uniqueOutputs = (Map<String, Integer>) objectStream.readObject(); |
| objectStream.close(); |
| } catch (FileNotFoundException e) { |
| e.printStackTrace(); |
| } catch (IOException e) { |
| e.printStackTrace(); |
| } catch (ClassNotFoundException e) { |
| e.printStackTrace(); |
| } |
| |
| } |
| |
| private void saveUniqueProgsData() { |
| // Since we could potentially stop the program while writing out this DB, |
| // copy the old file beforehand, and then delete it if we successfully wrote out the DB. |
| boolean oldWasSaved = false; |
| File file = new File(databaseFile); |
| if (file.exists()) { |
| try { |
| Process process = |
| Runtime.getRuntime().exec(String.format("cp %1$s %1$s.old", databaseFile)); |
| // Shouldn't block, cp shouldn't produce output. |
| process.waitFor(); |
| oldWasSaved = true; |
| } catch (IOException exception) { |
| exception.printStackTrace(); |
| } catch (InterruptedException exception) { |
| exception.printStackTrace(); |
| } |
| } |
| |
| // Now write out the DB. |
| boolean success = false; |
| try { |
| ObjectOutputStream objectStream = |
| new ObjectOutputStream(new FileOutputStream(databaseFile)); |
| objectStream.writeObject(uniquePrograms); |
| objectStream.writeObject(uniqueOutputs); |
| objectStream.close(); |
| success = true; |
| } catch (FileNotFoundException e) { |
| e.printStackTrace(); |
| } catch (IOException e) { |
| e.printStackTrace(); |
| } |
| |
| // If we get here, and we successfully wrote out the DB, delete the saved one. |
| if (oldWasSaved && success) { |
| try { |
| Process process = |
| Runtime.getRuntime().exec(String.format("rm %s.old", databaseFile)); |
| // Shouldn't block, rm shouldn't produce output. |
| process.waitFor(); |
| } catch (IOException exception) { |
| exception.printStackTrace(); |
| } catch (InterruptedException exception) { |
| exception.printStackTrace(); |
| } |
| } else if (oldWasSaved && !success) { |
| Log.error("Failed to successfully write out the unique programs DB!"); |
| Log.error("Old DB should be saved in " + databaseFile + ".old"); |
| } |
| } |
| |
| private void addToMap(String md5sum, Map<String, Integer> map) { |
| if (map.containsKey(md5sum)) { |
| map.put(md5sum, map.get(md5sum) + 1); |
| } else { |
| map.put(md5sum, 1); |
| } |
| } |
| |
| private void saveDivergentProgram() { |
| File before = new File(fuzzedFile); |
| File after = new File(String.format("divergent_programs/%d.dex", currentSeed)); |
| boolean success = before.renameTo(after); |
| if (!success) { |
| Log.error("Failed to save divergent program! Does divergent_programs/ exist?"); |
| } |
| } |
| |
| @Override |
| public void setup() { |
| try { |
| digest = MessageDigest.getInstance("MD5"); |
| loadUniqueProgsData(); |
| } catch (NoSuchAlgorithmException e) { |
| e.printStackTrace(); |
| } |
| } |
| |
| @Override |
| public void handleIterationFinished(int iteration) { |
| if ((iteration % saveDatabasePeriod) == (saveDatabasePeriod - 1)) { |
| saveUniqueProgsData(); |
| } |
| } |
| |
| @Override |
| public void handleSuccessfullyFuzzedFile(String programName) { |
| String md5sum = getMD5SumOfProgram(programName); |
| addToMap(md5sum, uniquePrograms); |
| |
| fuzzedFile = programName; |
| } |
| |
| @Override |
| public void handleDivergences(Map<String, List<Executor>> outputMap) { |
| // Just use the first one. |
| String output = (String) outputMap.keySet().toArray()[0]; |
| String md5sum = getMD5SumOfOutput(output); |
| addToMap(md5sum, uniqueOutputs); |
| |
| saveDivergentProgram(); |
| } |
| |
| @Override |
| public void handleSuccess(Map<String, List<Executor>> outputMap) { |
| // There's only one, use it. |
| String output = (String) outputMap.keySet().toArray()[0]; |
| String md5sum = getMD5SumOfOutput(output); |
| addToMap(md5sum, uniqueOutputs); |
| } |
| |
| @Override |
| public void handleSummary() { |
| if (Options.reportUnique) { |
| Log.always("-- UNIQUE PROGRAM REPORT --"); |
| Log.always("Unique Programs Seen: " + uniquePrograms.size()); |
| Log.always("Unique Outputs Seen: " + uniqueOutputs.size()); |
| Log.always("---------------------------"); |
| } |
| |
| saveUniqueProgsData(); |
| } |
| |
| } |