1/* 2 * Copyright (C) 2014 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17package dexfuzz.listeners; 18 19import dexfuzz.Log; 20import dexfuzz.Options; 21import dexfuzz.executors.Executor; 22 23import java.io.File; 24import java.io.FileInputStream; 25import java.io.FileNotFoundException; 26import java.io.FileOutputStream; 27import java.io.IOException; 28import java.io.ObjectInputStream; 29import java.io.ObjectOutputStream; 30import java.security.MessageDigest; 31import java.security.NoSuchAlgorithmException; 32import java.util.HashMap; 33import java.util.List; 34import java.util.Map; 35 36/** 37 * Tracks unique programs and outputs. Also saves divergent programs! 38 */ 39public class UniqueProgramTrackerListener extends BaseListener { 40 /** 41 * Map of unique program MD5 sums, mapped to times seen. 42 */ 43 private Map<String, Integer> uniquePrograms; 44 45 /** 46 * Map of unique program outputs (MD5'd), mapped to times seen. 47 */ 48 private Map<String, Integer> uniqueOutputs; 49 50 /** 51 * Used to remember the seed used to fuzz the fuzzed file, so we can save it with this 52 * seed as a name, if we find a divergence. 53 */ 54 private long currentSeed; 55 56 /** 57 * Used to remember the name of the file we've fuzzed, so we can save it if we 58 * find a divergence. 59 */ 60 private String fuzzedFile; 61 62 private MessageDigest digest; 63 private String databaseFile; 64 65 /** 66 * Save the database every X number of iterations. 67 */ 68 private static final int saveDatabasePeriod = 20; 69 70 public UniqueProgramTrackerListener(String databaseFile) { 71 this.databaseFile = databaseFile; 72 } 73 74 @Override 75 public void handleSeed(long seed) { 76 currentSeed = seed; 77 } 78 79 /** 80 * Given a program filename, calculate the MD5sum of 81 * this program. 82 */ 83 private String getMD5SumOfProgram(String programName) { 84 byte[] buf = new byte[256]; 85 try { 86 FileInputStream stream = new FileInputStream(programName); 87 boolean done = false; 88 while (!done) { 89 int bytesRead = stream.read(buf); 90 if (bytesRead == -1) { 91 done = true; 92 } else { 93 digest.update(buf); 94 } 95 } 96 stream.close(); 97 } catch (FileNotFoundException e) { 98 e.printStackTrace(); 99 } catch (IOException e) { 100 e.printStackTrace(); 101 } 102 return new String(digest.digest()); 103 } 104 105 private String getMD5SumOfOutput(String output) { 106 digest.update(output.getBytes()); 107 return new String(digest.digest()); 108 } 109 110 @SuppressWarnings("unchecked") 111 private void loadUniqueProgsData() { 112 File file = new File(databaseFile); 113 if (!file.exists()) { 114 uniquePrograms = new HashMap<String, Integer>(); 115 uniqueOutputs = new HashMap<String, Integer>(); 116 return; 117 } 118 119 try { 120 ObjectInputStream objectStream = 121 new ObjectInputStream(new FileInputStream(databaseFile)); 122 uniquePrograms = (Map<String, Integer>) objectStream.readObject(); 123 uniqueOutputs = (Map<String, Integer>) objectStream.readObject(); 124 objectStream.close(); 125 } catch (FileNotFoundException e) { 126 e.printStackTrace(); 127 } catch (IOException e) { 128 e.printStackTrace(); 129 } catch (ClassNotFoundException e) { 130 e.printStackTrace(); 131 } 132 133 } 134 135 private void saveUniqueProgsData() { 136 // Since we could potentially stop the program while writing out this DB, 137 // copy the old file beforehand, and then delete it if we successfully wrote out the DB. 138 boolean oldWasSaved = false; 139 File file = new File(databaseFile); 140 if (file.exists()) { 141 try { 142 Process process = 143 Runtime.getRuntime().exec(String.format("cp %1$s %1$s.old", databaseFile)); 144 // Shouldn't block, cp shouldn't produce output. 145 process.waitFor(); 146 oldWasSaved = true; 147 } catch (IOException exception) { 148 exception.printStackTrace(); 149 } catch (InterruptedException exception) { 150 exception.printStackTrace(); 151 } 152 } 153 154 // Now write out the DB. 155 boolean success = false; 156 try { 157 ObjectOutputStream objectStream = 158 new ObjectOutputStream(new FileOutputStream(databaseFile)); 159 objectStream.writeObject(uniquePrograms); 160 objectStream.writeObject(uniqueOutputs); 161 objectStream.close(); 162 success = true; 163 } catch (FileNotFoundException e) { 164 e.printStackTrace(); 165 } catch (IOException e) { 166 e.printStackTrace(); 167 } 168 169 // If we get here, and we successfully wrote out the DB, delete the saved one. 170 if (oldWasSaved && success) { 171 try { 172 Process process = 173 Runtime.getRuntime().exec(String.format("rm %s.old", databaseFile)); 174 // Shouldn't block, rm shouldn't produce output. 175 process.waitFor(); 176 } catch (IOException exception) { 177 exception.printStackTrace(); 178 } catch (InterruptedException exception) { 179 exception.printStackTrace(); 180 } 181 } else if (oldWasSaved && !success) { 182 Log.error("Failed to successfully write out the unique programs DB!"); 183 Log.error("Old DB should be saved in " + databaseFile + ".old"); 184 } 185 } 186 187 private void addToMap(String md5sum, Map<String, Integer> map) { 188 if (map.containsKey(md5sum)) { 189 map.put(md5sum, map.get(md5sum) + 1); 190 } else { 191 map.put(md5sum, 1); 192 } 193 } 194 195 private void saveDivergentProgram() { 196 File before = new File(fuzzedFile); 197 File after = new File(String.format("divergent_programs/%d.dex", currentSeed)); 198 boolean success = before.renameTo(after); 199 if (!success) { 200 Log.error("Failed to save divergent program! Does divergent_programs/ exist?"); 201 } 202 } 203 204 @Override 205 public void setup() { 206 try { 207 digest = MessageDigest.getInstance("MD5"); 208 loadUniqueProgsData(); 209 } catch (NoSuchAlgorithmException e) { 210 e.printStackTrace(); 211 } 212 } 213 214 @Override 215 public void handleIterationFinished(int iteration) { 216 if ((iteration % saveDatabasePeriod) == (saveDatabasePeriod - 1)) { 217 saveUniqueProgsData(); 218 } 219 } 220 221 @Override 222 public void handleSuccessfullyFuzzedFile(String programName) { 223 String md5sum = getMD5SumOfProgram(programName); 224 addToMap(md5sum, uniquePrograms); 225 226 fuzzedFile = programName; 227 } 228 229 @Override 230 public void handleDivergences(Map<String, List<Executor>> outputMap) { 231 // Just use the first one. 232 String output = (String) outputMap.keySet().toArray()[0]; 233 String md5sum = getMD5SumOfOutput(output); 234 addToMap(md5sum, uniqueOutputs); 235 236 saveDivergentProgram(); 237 } 238 239 @Override 240 public void handleSuccess(Map<String, List<Executor>> outputMap) { 241 // There's only one, use it. 242 String output = (String) outputMap.keySet().toArray()[0]; 243 String md5sum = getMD5SumOfOutput(output); 244 addToMap(md5sum, uniqueOutputs); 245 } 246 247 @Override 248 public void handleSummary() { 249 if (Options.reportUnique) { 250 Log.always("-- UNIQUE PROGRAM REPORT --"); 251 Log.always("Unique Programs Seen: " + uniquePrograms.size()); 252 Log.always("Unique Outputs Seen: " + uniqueOutputs.size()); 253 Log.always("---------------------------"); 254 } 255 256 saveUniqueProgsData(); 257 } 258 259} 260