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.program;
  
  import dexfuzz.Log;
  import dexfuzz.MutationStats;
  import dexfuzz.Options;
  import dexfuzz.listeners.BaseListener;
  import dexfuzz.program.mutators.ArithOpChanger;
  import dexfuzz.program.mutators.BranchShifter;
  import dexfuzz.program.mutators.CmpBiasChanger;
  import dexfuzz.program.mutators.CodeMutator;
  import dexfuzz.program.mutators.ConstantValueChanger;
  import dexfuzz.program.mutators.ConversionRepeater;
  import dexfuzz.program.mutators.FieldFlagChanger;
  import dexfuzz.program.mutators.InstructionDeleter;
  import dexfuzz.program.mutators.InstructionDuplicator;
  import dexfuzz.program.mutators.InstructionSwapper;
  import dexfuzz.program.mutators.NewMethodCaller;
  import dexfuzz.program.mutators.NonsenseStringPrinter;
  import dexfuzz.program.mutators.PoolIndexChanger;
  import dexfuzz.program.mutators.RandomInstructionGenerator;
  import dexfuzz.program.mutators.SwitchBranchShifter;
  import dexfuzz.program.mutators.TryBlockShifter;
  import dexfuzz.program.mutators.ValuePrinter;
  import dexfuzz.program.mutators.VRegChanger;
  import dexfuzz.rawdex.ClassDataItem;
  import dexfuzz.rawdex.ClassDefItem;
  import dexfuzz.rawdex.CodeItem;
  import dexfuzz.rawdex.DexRandomAccessFile;
  import dexfuzz.rawdex.EncodedField;
  import dexfuzz.rawdex.EncodedMethod;
  import dexfuzz.rawdex.FieldIdItem;
  import dexfuzz.rawdex.MethodIdItem;
  import dexfuzz.rawdex.ProtoIdItem;
  import dexfuzz.rawdex.RawDexFile;
  import dexfuzz.rawdex.TypeIdItem;
  import dexfuzz.rawdex.formats.ContainsPoolIndex.PoolIndexKind;
  
  import java.io.BufferedReader;
  import java.io.BufferedWriter;
  import java.io.FileReader;
  import java.io.FileWriter;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Random;
  
  /**
   * After the raw DEX file has been parsed, it is passed into this class
   * that represents the program in a mutatable form.
   * The class uses a CodeTranslator to translate between the raw DEX form
   * for a method, and the mutatable form. It also controls all CodeMutators,
   * deciding which ones should be applied to each CodeItem.
   */
  public class Program {
    /**
     * The RNG used during mutation.
     */
    private Random rng;
  
    /**
     * The seed that was given to the RNG.
     */
    public long rngSeed;
  
    /**
     * The parsed raw DEX file.
     */
    private RawDexFile rawDexFile;
  
    /**
     * The system responsible for translating from CodeItems to MutatableCode and vice-versa.
     */
    private CodeTranslator translator;
  
    /**
     * Responsible for adding new class ID items, method ID items, etc.
     */
    private IdCreator idCreator;
  
    /**
     * A list of all the MutatableCode that the CodeTranslator produced from
     * CodeItems that are acceptable to mutate.
     */
    private List<MutatableCode> mutatableCodes;
  
    /**
     * A list of all MutatableCode items that were mutated when mutateTheProgram()
     * was called. updateRawDexFile() will update the relevant CodeItems when called,
     * and then clear this list.
     */
    private List<MutatableCode> mutatedCodes;
  
    /**
     * A list of all registered CodeMutators that this Program can use to mutate methods.
     */
    private List<CodeMutator> mutators;
  
    /**
     * Used if we're loading mutations from a file, so we can find the correct mutator.
     */
    private Map<Class<? extends CodeMutator>, CodeMutator> mutatorsLookupByClass;
  
    /**
     * Tracks mutation stats.
     */
    private MutationStats mutationStats;
  
    /**
     * A list of all mutations used for loading/dumping mutations from/to a file.
     */
    private List<Mutation> mutations;
  
    /**
     * The listener who is interested in events.
     * May be a listener that is responsible for multiple listeners.
     */
    private BaseListener listener;
  
    /**
     * Given a maximum number of mutations that can be performed on a method, n,
     * give up after attempting (n * this value) mutations for any method.
     */
    private static final int MAXIMUM_MUTATION_ATTEMPT_FACTOR = 10;
  
    /**
     * Construct the mutatable Program based on the raw DEX file that was parsed initially.
     */
    public Program(RawDexFile rawDexFile, List<Mutation> previousMutations,
        BaseListener listener) {
      this.listener = listener;
  
      idCreator = new IdCreator(rawDexFile);
  
      // Set up the RNG.
      rng = new Random();
      if (Options.usingProvidedSeed) {
        rng.setSeed(Options.rngSeed);
        rngSeed = Options.rngSeed;
      } else {
        long seed = System.currentTimeMillis();
        listener.handleSeed(seed);
        rng.setSeed(seed);
        rngSeed = seed;
      }
  
      if (previousMutations != null) {
        mutations = previousMutations;
      } else {
        // Allocate the mutations list.
        mutations = new ArrayList<Mutation>();
  
        // Read in the mutations if we need to.
        if (Options.loadMutations) {
          // Allocate the mutators lookup table.
          mutatorsLookupByClass = new HashMap<Class<? extends CodeMutator>, CodeMutator>();
          loadMutationsFromDisk(Options.loadMutationsFile);
        }
      }
  
      // Allocate the mutators list.
      mutators = new ArrayList<CodeMutator>();
  
      this.rawDexFile = rawDexFile;
  
      mutatableCodes = new ArrayList<MutatableCode>();
      mutatedCodes = new ArrayList<MutatableCode>();
  
      translator = new CodeTranslator();
  
      mutationStats = new MutationStats();
  
      // Register all the code mutators here.
      registerMutator(new ArithOpChanger(rng, mutationStats, mutations));
      registerMutator(new BranchShifter(rng, mutationStats, mutations));
      registerMutator(new CmpBiasChanger(rng, mutationStats, mutations));
      registerMutator(new ConstantValueChanger(rng, mutationStats, mutations));
      registerMutator(new ConversionRepeater(rng, mutationStats, mutations));
      registerMutator(new FieldFlagChanger(rng, mutationStats, mutations));
      registerMutator(new InstructionDeleter(rng, mutationStats, mutations));
      registerMutator(new InstructionDuplicator(rng, mutationStats, mutations));
      registerMutator(new InstructionSwapper(rng, mutationStats, mutations));
      registerMutator(new NewMethodCaller(rng, mutationStats, mutations));
      registerMutator(new NonsenseStringPrinter(rng, mutationStats, mutations));
      registerMutator(new PoolIndexChanger(rng, mutationStats, mutations));
      registerMutator(new RandomInstructionGenerator(rng, mutationStats, mutations));
      registerMutator(new SwitchBranchShifter(rng, mutationStats, mutations));
      registerMutator(new TryBlockShifter(rng, mutationStats, mutations));
      registerMutator(new ValuePrinter(rng, mutationStats, mutations));
      registerMutator(new VRegChanger(rng, mutationStats, mutations));
  
      associateClassDefsAndClassData();
      associateCodeItemsWithMethodNames();
  
      int codeItemIdx = 0;
      for (CodeItem codeItem : rawDexFile.codeItems) {
        if (legalToMutate(codeItem)) {
          Log.debug("Legal to mutate code item " + codeItemIdx);
          int mutatableCodeIdx = mutatableCodes.size();
          mutatableCodes.add(translator.codeItemToMutatableCode(this, codeItem,
              codeItemIdx, mutatableCodeIdx));
        } else {
          Log.debug("Not legal to mutate code item " + codeItemIdx);
        }
        codeItemIdx++;
      }
    }
  
    private void registerMutator(CodeMutator mutator) {
      if (mutator.canBeTriggered()) {
        Log.debug("Registering mutator " + mutator.getClass().getSimpleName());
        mutators.add(mutator);
      }
      if (Options.loadMutations) {
        mutatorsLookupByClass.put(mutator.getClass(), mutator);
      }
    }
  
    /**
     * Associate ClassDefItem to a ClassDataItem and vice-versa.
     * This is so when we're associating method names with code items,
     * we can find the name of the class the method belongs to.
     */
    private void associateClassDefsAndClassData() {
      for (ClassDefItem classDefItem : rawDexFile.classDefs) {
        if (classDefItem.classDataOff.pointsToSomething()) {
          ClassDataItem classDataItem = (ClassDataItem)
              classDefItem.classDataOff.getPointedToItem();
          classDataItem.meta.classDefItem = classDefItem;
          classDefItem.meta.classDataItem = classDataItem;
        }
      }
    }
  
    /**
     * For each CodeItem, find the name of the method the item represents.
     * This is done to allow the filtering of mutating methods based on if
     * they have the name *_MUTATE, but also for debugging info.
     */
    private void associateCodeItemsWithMethodNames() {
      // Associate method names with codeItems.
      for (ClassDataItem classDataItem : rawDexFile.classDatas) {
  
        String className = "";
        if (classDataItem.meta.classDefItem != null) {
          int typeIdx = classDataItem.meta.classDefItem.classIdx;
          TypeIdItem typeIdItem = rawDexFile.typeIds.get(typeIdx);
          className = rawDexFile.stringDatas.get(typeIdItem.descriptorIdx).getString() + ".";
        }
  
        // Do direct methods...
        // Track the current method index with this value, since the encoding in
        // each EncodedMethod is the absolute index for the first EncodedMethod,
        // and then relative index for the rest...
        int methodIdx = 0;
        for (EncodedMethod method : classDataItem.directMethods) {
          methodIdx = associateMethod(method, methodIdx, className);
        }
        // Reset methodIdx for virtual methods...
        methodIdx = 0;
        for (EncodedMethod method : classDataItem.virtualMethods) {
          methodIdx = associateMethod(method, methodIdx, className);
        }
      }
    }
  
    /**
     * Associate the name of the provided method with its CodeItem, if it
     * has one.
     *
     * @param methodIdx The method index of the last EncodedMethod that was handled in this class.
     * @return The method index of the EncodedMethod that has just been handled in this class.
     */
    private int associateMethod(EncodedMethod method, int methodIdx, String className) {
      if (!method.codeOff.pointsToSomething()) {
        // This method doesn't have a code item, so we won't encounter it later.
        return methodIdx;
      }
  
      // First method index is an absolute index.
      // The rest are relative to the previous.
      // (so if methodIdx is initialised to 0, this single line works)
      methodIdx = methodIdx + method.methodIdxDiff;
  
      // Get the name.
      MethodIdItem methodIdItem = rawDexFile.methodIds.get(methodIdx);
      ProtoIdItem protoIdItem = rawDexFile.protoIds.get(methodIdItem.protoIdx);
      String shorty = rawDexFile.stringDatas.get(protoIdItem.shortyIdx).getString();
      String methodName = className
          + rawDexFile.stringDatas.get(methodIdItem.nameIdx).getString();
  
      // Get the codeItem.
      if (method.codeOff.getPointedToItem() instanceof CodeItem) {
        CodeItem codeItem = (CodeItem) method.codeOff.getPointedToItem();
        codeItem.meta.methodName = methodName;
        codeItem.meta.shorty = shorty;
        codeItem.meta.isStatic = method.isStatic();
      } else {
        Log.errorAndQuit("You've got an EncodedMethod that points to an Offsettable"
            + " that does not contain a CodeItem");
      }
  
      return methodIdx;
    }
  
    /**
     * Determine, based on the current options supplied to dexfuzz, as well as
     * its capabilities, if the provided CodeItem can be mutated.
     * @param codeItem The CodeItem we may wish to mutate.
     * @return If the CodeItem can be mutated.
     */
    private boolean legalToMutate(CodeItem codeItem) {
      if (!Options.mutateLimit) {
        Log.debug("Mutating everything.");
        return true;
      }
      if (codeItem.meta.methodName.endsWith("_MUTATE")) {
        Log.debug("Code item marked with _MUTATE.");
        return true;
      }
      Log.debug("Code item not marked with _MUTATE, but not mutating all code items.");
      return false;
    }
  
    private int getNumberOfMutationsToPerform() {
      // We want n mutations to be twice as likely as n+1 mutations.
      //
      // So if we have max 3,
      // then 0 has 8 chances ("tickets"),
      //      1 has 4 chances
      //      2 has 2 chances
      //  and 3 has 1 chance
  
      // Allocate the tickets
      // n mutations need (2^(n+1) - 1) tickets
      // e.g.
      // 3 mutations => 15 tickets
      // 4 mutations => 31 tickets
      int tickets = (2 << Options.methodMutations) - 1;
  
      // Pick the lucky ticket
      int luckyTicket = rng.nextInt(tickets);
  
      // The tickets are put into buckets with accordance with log-base-2.
      // have to make sure it's luckyTicket + 1, because log(0) is undefined
      // so:
      // log_2(1) => 0
      // log_2(2) => 1
      // log_2(3) => 1
      // log_2(4) => 2
      // log_2(5) => 2
      // log_2(6) => 2
      // log_2(7) => 2
      // log_2(8) => 3
      // ...
      // so to make the highest mutation value the rarest,
      //   subtract log_2(luckyTicket+1) from the maximum number
      // log2(x) <=> 31 - Integer.numberOfLeadingZeros(x)
      int luckyMutation = Options.methodMutations
          - (31 - Integer.numberOfLeadingZeros(luckyTicket + 1));
  
      return luckyMutation;
    }
  
    /**
     * Returns true if we completely failed to mutate this method's mutatable code after
     * attempting to.
     */
    private boolean mutateAMutatableCode(MutatableCode mutatableCode) {
      int mutations = getNumberOfMutationsToPerform();
  
      Log.info("Attempting " + mutations + " mutations for method " + mutatableCode.name);
  
      int mutationsApplied = 0;
  
      int maximumMutationAttempts = Options.methodMutations * MAXIMUM_MUTATION_ATTEMPT_FACTOR;
      int mutationAttempts = 0;
      boolean hadToBail = false;
  
      while (mutationsApplied < mutations) {
        int mutatorIdx = rng.nextInt(mutators.size());
        CodeMutator mutator = mutators.get(mutatorIdx);
        Log.info("Running mutator " + mutator.getClass().getSimpleName());
        if (mutator.attemptToMutate(mutatableCode)) {
          mutationsApplied++;
        }
        mutationAttempts++;
        if (mutationAttempts > maximumMutationAttempts) {
          Log.info("Bailing out on mutation for this method, tried too many times...");
          hadToBail = true;
          break;
        }
      }
  
      // If any of them actually mutated it, excellent!
      if (mutationsApplied > 0) {
        Log.info("Method was mutated.");
        mutatedCodes.add(mutatableCode);
      } else {
        Log.info("Method was not mutated.");
      }
  
      return ((mutationsApplied == 0) && hadToBail);
    }
  
    /**
     * Go through each mutatable method in turn, and attempt to mutate it.
     * Afterwards, call updateRawDexFile() to apply the results of mutation to the
     * original code.
     */
    public void mutateTheProgram() {
      if (Options.loadMutations) {
        applyMutationsFromList();
        return;
      }
  
      // Typically, this is 2 to 10...
      int methodsToMutate = Options.minMethods
          + rng.nextInt((Options.maxMethods - Options.minMethods) + 1);
  
      // Check we aren't trying to mutate more methods than we have.
      if (methodsToMutate > mutatableCodes.size()) {
        methodsToMutate = mutatableCodes.size();
      }
  
      // Check if we're going to end up mutating all the methods.
      if (methodsToMutate == mutatableCodes.size()) {
        // Just do them all in order.
        Log.info("Mutating all possible methods.");
        for (MutatableCode mutatableCode : mutatableCodes) {
          if (mutatableCode == null) {
            Log.errorAndQuit("Why do you have a null MutatableCode?");
          }
          mutateAMutatableCode(mutatableCode);
        }
        Log.info("Finished mutating all possible methods.");
      } else {
        // Pick them at random.
        Log.info("Randomly selecting " + methodsToMutate + " methods to mutate.");
        while (mutatedCodes.size() < methodsToMutate) {
          int randomMethodIdx = rng.nextInt(mutatableCodes.size());
          MutatableCode mutatableCode = mutatableCodes.get(randomMethodIdx);
          if (mutatableCode == null) {
            Log.errorAndQuit("Why do you have a null MutatableCode?");
          }
          if (!mutatedCodes.contains(mutatableCode)) {
            boolean completelyFailedToMutate = mutateAMutatableCode(mutatableCode);
            if (completelyFailedToMutate) {
              methodsToMutate--;
            }
          }
        }
        Log.info("Finished mutating the methods.");
      }
  
      listener.handleMutationStats(mutationStats.getStatsString());
  
      if (Options.dumpMutations) {
        writeMutationsToDisk(Options.dumpMutationsFile);
      }
    }
  
    private void writeMutationsToDisk(String fileName) {
      Log.debug("Writing mutations to disk.");
      try {
        BufferedWriter writer = new BufferedWriter(new FileWriter(fileName));
        for (Mutation mutation : mutations) {
          MutationSerializer.writeMutation(writer, mutation);
        }
        writer.close();
      } catch (IOException e) {
        Log.errorAndQuit("IOException while writing mutations to disk...");
      }
    }
  
    private void loadMutationsFromDisk(String fileName) {
      Log.debug("Loading mutations from disk.");
      try {
        BufferedReader reader = new BufferedReader(new FileReader(fileName));
        while (reader.ready()) {
          Mutation mutation = MutationSerializer.readMutation(reader);
          mutations.add(mutation);
        }
        reader.close();
      } catch (IOException e) {
        Log.errorAndQuit("IOException while loading mutations from disk...");
      }
    }
  
    private void applyMutationsFromList() {
      Log.info("Applying preloaded list of mutations...");
      for (Mutation mutation : mutations) {
        // Repopulate the MutatableCode field from the recorded index into the Program's list.
        mutation.mutatableCode = mutatableCodes.get(mutation.mutatableCodeIdx);
  
        // Get the right mutator.
        CodeMutator mutator = mutatorsLookupByClass.get(mutation.mutatorClass);
  
        // Apply the mutation.
        mutator.forceMutate(mutation);
  
        // Add this mutatable code to the list of mutated codes, if we haven't already.
        if (!mutatedCodes.contains(mutation.mutatableCode)) {
          mutatedCodes.add(mutation.mutatableCode);
        }
      }
      Log.info("...finished applying preloaded list of mutations.");
    }
  
    public List<Mutation> getMutations() {
      return mutations;
    }
  
    /**
     * Updates any CodeItems that need to be updated after mutation.
     */
    public boolean updateRawDexFile() {
      boolean anythingMutated = !(mutatedCodes.isEmpty());
      for (MutatableCode mutatedCode : mutatedCodes) {
        translator.mutatableCodeToCodeItem(rawDexFile.codeItems
            .get(mutatedCode.codeItemIdx), mutatedCode);
      }
      mutatedCodes.clear();
      return anythingMutated;
    }
  
    public void writeRawDexFile(DexRandomAccessFile file) throws IOException {
      rawDexFile.write(file);
    }
  
    public void updateRawDexFileHeader(DexRandomAccessFile file) throws IOException {
      rawDexFile.updateHeader(file);
    }
  
    /**
     * Used by the CodeMutators to determine legal index values.
     */
    public int getTotalPoolIndicesByKind(PoolIndexKind poolIndexKind) {
      switch (poolIndexKind) {
        case Type:
          return rawDexFile.typeIds.size();
        case Field:
          return rawDexFile.fieldIds.size();
        case String:
          return rawDexFile.stringIds.size();
        case Method:
          return rawDexFile.methodIds.size();
        case Invalid:
          return 0;
        default:
      }
      return 0;
    }
  
    /**
     * Used by the CodeMutators to lookup and/or create Ids.
     */
    public IdCreator getNewItemCreator() {
      return idCreator;
    }
  
    /**
     * Used by FieldFlagChanger, to find an EncodedField for a specified field in an insn,
     * if that field is actually defined in this DEX file. If not, null is returned.
     */
    public EncodedField getEncodedField(int fieldIdx) {
      if (fieldIdx >= rawDexFile.fieldIds.size()) {
        Log.debug(String.format("Field idx 0x%x specified is not defined in this DEX file.",
            fieldIdx));
        return null;
      }
      FieldIdItem fieldId = rawDexFile.fieldIds.get(fieldIdx);
  
      for (ClassDefItem classDef : rawDexFile.classDefs) {
        if (classDef.classIdx == fieldId.classIdx) {
          ClassDataItem classData = classDef.meta.classDataItem;
          return classData.getEncodedFieldWithIndex(fieldIdx);
        }
      }
  
      Log.debug(String.format("Field idx 0x%x specified is not defined in this DEX file.",
          fieldIdx));
      return null;
    }
  }
  