content: /*
   * 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();
    }
  
  }
  