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