Ejemplo n.º 1
0
int main (int argc, char **argv)
{

    if(!parseOptions(argc,argv)) {
        return EXIT_FAILURE;
    }

    BPatch bpatch;

    BPatch_binaryEdit *appBin = bpatch.openBinary (originalBinary, !instrumentLibraries.empty());
    if (appBin == NULL) {
        cerr << "Failed to open binary" << endl;
        return EXIT_FAILURE;
    }

    if (!appBin->loadLibrary (instLibrary)) {
        cerr << "Failed to open instrumentation library." << endl;
        cerr << "It needs to be located in the current working directory." << endl;
        return EXIT_FAILURE;
    }

    BPatch_image *appImage = appBin->getImage ();

    /* Find code coverage functions in the instrumentation library */
    BPatch_function *initAflForkServer =
        findFuncByName (appImage, (char *) "initAflForkServer");
    BPatch_function *bbCallback =
        findFuncByName (appImage, (char *) "bbCallback");
    if (!initAflForkServer || !bbCallback ) {
        cerr << "Instrumentation library lacks callbacks!" << endl;
        return EXIT_FAILURE;
    }

    //get and iterate over all modules, instrumenting only the default and manualy specified ones
    vector < BPatch_module * >*modules = appImage->getModules ();
    vector < BPatch_module * >::iterator moduleIter;
    BPatch_module *defaultModule = NULL;
    string defaultModuleName;
    for (moduleIter = modules->begin (); moduleIter != modules->end (); ++moduleIter) {
    //find default module name
        char moduleName[1024];
        (*moduleIter)->getName (moduleName, 1024);    
        if (string (moduleName).find ("DEFAULT_MODULE") != string::npos) {
            defaultModuleName = "DEFAULT_MODULE";
        }
    }
    if(defaultModuleName.empty()) 
        defaultModuleName = string(originalBinary).substr(string(originalBinary).find_last_of("\\/")+1);
    int bbIndex = 0;
    for (moduleIter = modules->begin (); moduleIter != modules->end (); ++moduleIter) {
        char moduleName[1024];
        (*moduleIter)->getName (moduleName, 1024);

        if ((*moduleIter)->isSharedLib ()) {
            if (instrumentLibraries.find (moduleName) == instrumentLibraries.end ()) {
                cout << "Skipping library: " << moduleName << endl;
                continue;
            }
        }

        if (string (moduleName).find (defaultModuleName) != string::npos) {
            defaultModule = (*moduleIter);
            if(skipMainModule) continue;
        }
        cout << "Instrumenting module: " << moduleName << endl;
        vector < BPatch_function * >*allFunctions =
            (*moduleIter)->getProcedures ();
        vector < BPatch_function * >::iterator funcIter;

        // iterate over all functions in the module
        for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
             ++funcIter) {
            BPatch_function *curFunc = *funcIter;
            char funcName[1024];
            curFunc->getName (funcName, 1024);
            if(string (funcName) == string("_start")) continue; // here's a bug on hlt
            insertBBCallback (appBin, curFunc, funcName, bbCallback, &bbIndex);
        }

    }

    //if entrypoint set ,find function  , else find _init
    BPatch_function *funcToPatch = NULL;
    if(!entryPoint) {
        BPatch_Vector<BPatch_function*> funcs;
        defaultModule->findFunction("_init", funcs);
        if(!funcs.size()) {
            cerr << "Couldn't locate _init, specify entry point manualy. "<< endl;
            return EXIT_FAILURE;
        }
        // there should really be only one
        funcToPatch = funcs[0];
    } else {
        funcToPatch = defaultModule->findFunctionByEntry(entryPoint);
    }
    if(!funcToPatch) {
        cerr << "Couldn't locate function at given entry point. "<< endl;
        return EXIT_FAILURE;
    }
    if(!insertCallToInit (appBin,  initAflForkServer,defaultModule,funcToPatch)){
        cerr << "Could not insert init callback at given entry point." << endl;
        return EXIT_FAILURE;
    }

    cout << "Saving the instrumented binary to " << instrumentedBinary << "..." << endl;
    // Output the instrumented binary
    if (!appBin->writeFile (instrumentedBinary)) {
        cerr << "Failed to write output file: " << instrumentedBinary << endl;
        return EXIT_FAILURE;
    }

    if(!runtimeLibraries.empty()) {
        cout << "Instrumenting runtime libraries." << endl;
        set<string>::iterator rtLibIter ;
        for(rtLibIter = runtimeLibraries.begin(); rtLibIter != runtimeLibraries.end(); rtLibIter++) {
            BPatch_binaryEdit *libBin = bpatch.openBinary ((*rtLibIter).c_str(), false);
            if (libBin == NULL) {
                cerr << "Failed to open binary "<< *rtLibIter << endl;
                return EXIT_FAILURE;
            }
            libBin->loadLibrary (instLibrary);
            BPatch_image *libImg = libBin->getImage ();
            vector < BPatch_module * >*modules = libImg->getModules ();
            moduleIter = modules->begin ();
            ++moduleIter;
            for ( ; moduleIter != modules->end (); ++moduleIter) {
                char moduleName[1024];
                (*moduleIter)->getName (moduleName, 1024);
                cout << "Instrumenting module: " << moduleName << endl;
                vector < BPatch_function * >*allFunctions =
                    (*moduleIter)->getProcedures ();
                vector < BPatch_function * >::iterator funcIter;
                // iterate over all functions in the module
                for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
                     ++funcIter) {
                    BPatch_function *curFunc = *funcIter;
                    char funcName[1024];
                    curFunc->getName (funcName, 1024);
                    if(string (funcName) == string("_start")) continue;
                    insertBBCallback (libBin, curFunc, funcName, bbCallback, &bbIndex);
                }
            }
            if (!libBin->writeFile ((*rtLibIter + ".ins").c_str())) {
                cerr << "Failed to write output file: " <<(*rtLibIter + ".ins").c_str() << endl;
                return EXIT_FAILURE;
            } else {
                cout << "Saved the instrumented library to " << (*rtLibIter + ".ins").c_str() << "." << endl;
            }
        }
    }

    cout << "All done! Happy fuzzing!" << endl;
    return EXIT_SUCCESS;

}
Ejemplo n.º 2
0
int launch_mutator()
{
   unsigned i = 0, j = 0;
   const char *reason = NULL;
   char buf[STRING_MAX] = {0};
   BPatch_Vector< BPatch_module * > *appModules = NULL;
   BPatch_Vector< BPatch_function * > *appFunctions = NULL;
   BPatch_flowGraph *appCFG = NULL;

   dynHandle *dh = mutatorInit();
   if (!dh) return(-1);

   /**************************************************************
    * Parsing Phase
    */
   // Is this if redundant?  Will anybody run this without parsing the modules?
   if (config.parse_level >= PARSE_MODULE) {
      sendMsg(config.outfd, ID_PARSE_MODULE, INFO);
      appModules = dh->image->getModules();
      if (!appModules) {
         sendMsg(config.outfd, ID_PARSE_MODULE, INFO, ID_FAIL,
                 "Failure in BPatch_image::getModules()");
         return(-1);
      } else
         sendMsg(config.outfd, ID_PARSE_MODULE, INFO, ID_PASS,
                 appModules->size());
   }
   if (config.parse_level >= PARSE_FUNC) {
      for (i = 0; i < appModules->size(); ++i) {
         (*appModules)[i]->getName(buf, sizeof(buf));

         // Check module inclusion/exclusion regular expression rules.
         reason = config.mod_rules.getReason(buf);
         if (reason) sendMsg(config.outfd, ID_DATA_STRING, INFO, ID_INFO, reason);
         if (!config.mod_rules.isValid(buf)) continue;

         sendMsg(config.outfd, ID_PARSE_FUNC, INFO, ID_TEST, buf);
         appFunctions = (*appModules)[i]->getProcedures();
         if (!appFunctions) {
            sendMsg(config.outfd, ID_PARSE_FUNC, INFO, ID_FAIL,
                    "Failure in BPatch_module::getProcedures()");
            return(-1);
         } else
            sendMsg(config.outfd, ID_PARSE_FUNC, INFO, ID_PASS,
                    appFunctions->size());

         if (config.parse_level >= PARSE_CFG) {
            sendMsg(config.outfd, ID_PARSE_MODULE_CFG, INFO, ID_TEST,
                    (*appModules)[i]->getName(buf, sizeof(buf)));

            int cfg_warn_cnt = 0, cfg_pass_cnt = 0;
            for (j = 0; j < appFunctions->size(); ++j) {
               (*appFunctions)[j]->getName(buf, sizeof(buf));

               // Check function inclusion/exclusion regular expression rules.
               reason = config.func_rules.getReason(buf);
               if (reason) sendMsg(config.outfd, ID_DATA_STRING, INFO, ID_INFO, reason);
               if (!config.func_rules.isValid(buf)) continue;

               sendMsg(config.outfd, ID_PARSE_FUNC_CFG, VERB1, ID_TEST, buf);
               appCFG = (*appFunctions)[j]->getCFG();
               if (!appCFG) {
                  sendMsg(config.outfd, ID_PARSE_FUNC_CFG, VERB1, ID_WARN,
                          "Failure in BPatch_function::getCFG()");
                  ++cfg_warn_cnt;
               } else {
                  sendMsg(config.outfd, ID_PARSE_FUNC_CFG, VERB1, ID_PASS);
                  ++cfg_pass_cnt;
               }
            }
            if (cfg_warn_cnt)
               sendMsg(config.outfd, ID_PARSE_MODULE_CFG, INFO, ID_WARN,
                       sprintf_static("%d warning(s), %d passed.", cfg_warn_cnt, cfg_pass_cnt));
            else
               sendMsg(config.outfd, ID_PARSE_MODULE_CFG, INFO, ID_PASS);
         }
      }
   }

   /**************************************************************
    * Instrumentation Phase
    */

   /* Load the library, if specified */
   if (config.instType == USER_FUNC) {
       char instLibrary[1024];
       int offset = strcspn(config.inst_function, ":");
       strncpy (instLibrary, config.inst_function, offset);
       instLibrary[offset] = '\0';
       if(! dh->addSpace->loadLibrary(instLibrary)) {
          sendMsg(config.outfd, ID_INST_FIND_POINTS, VERB3, ID_FAIL,
                  "Failure in loading library");
          return false;
       }
   }

   if (config.trace_inst) 
   {
      errno = 0;
      sendMsg(config.outfd, ID_TRACE_INIT, INFO);
      sendMsg(config.outfd, ID_TRACE_OPEN_READER, VERB1);
	
      if (config.use_process)
      {
         config.pipefd = open(config.pipe_filename, O_RDONLY | O_RSYNC | O_NONBLOCK);
         if (config.pipefd < 0) {
            sendMsg(config.outfd, ID_TRACE_OPEN_READER, VERB1, ID_FAIL,
                    sprintf_static("Mutator could not open trace pipe (%s) for read: %s\n",
                                   config.pipe_filename, strerror(errno)));
            config.pipefd = -1;

         } else {
            sendMsg(config.outfd, ID_TRACE_OPEN_READER, VERB1, ID_PASS);
	      
            // Run mutatee side of trace initialization.
            sendMsg(config.outfd, ID_TRACE_INIT_MUTATEE, VERB1);
	      
            if (!initTraceInMutatee(dh)) {
               sendMsg(config.outfd, ID_TRACE_INIT_MUTATEE, VERB1, ID_FAIL);
               config.pipefd = -1;
            } else
               sendMsg(config.outfd, ID_TRACE_INIT_MUTATEE, VERB1, ID_PASS);
         }

         if (config.pipefd == -1) {
            sendMsg(config.outfd, ID_TRACE_INIT, INFO, ID_FAIL,
                    "Disabling instrumentation tracing.");
         } else
            sendMsg(config.outfd, ID_TRACE_INIT, INFO, ID_PASS);

         // We should remove the named pipe now, pass or fail.
         unlink(config.pipe_filename);
      }
   }

   if (config.inst_level >= INST_FUNC_ENTRY) {
      
      if (config.transMode == TRANS_PROCESS)
         if (!dynStartTransaction(dh)) return(-1);

      for (i = 0; i < appModules->size(); ++i) {
         int mod_warn_cnt = 0, mod_pass_cnt = 0;
         (*appModules)[i]->getName(buf, sizeof(buf));
         // Check module inclusion/exclusion regular expression rules.
         reason = config.mod_rules.getReason(buf);
         if (reason) 
            sendMsg(config.outfd, ID_DATA_STRING, INFO, ID_INFO, reason);
         if (!config.mod_rules.isValid(buf)) {
            continue;
         }

         if (config.transMode == TRANS_MODULE)
            if (!dynStartTransaction(dh)) {
               continue;
            }

         sendMsg(config.outfd, ID_INST_MODULE, INFO, ID_TEST, buf);

         sendMsg(config.outfd, ID_INST_GET_FUNCS, VERB1);
         appFunctions = (*appModules)[i]->getProcedures();
         if (!appFunctions) {
            sendMsg(config.outfd, ID_INST_GET_FUNCS, VERB1, ID_FAIL,
                    "Failure in BPatch_module::getProcedures()");
            return(-1);
         } else
            sendMsg(config.outfd, ID_INST_GET_FUNCS, VERB1, ID_PASS,
                    appFunctions->size());

         for (j = 0; j < appFunctions->size(); ++j) {
            int func_warn_cnt = 0, func_pass_cnt = 0;
            (*appFunctions)[j]->getName(buf, sizeof(buf));

            // Check function inclusion/exclusion regular expression rules.
            reason = config.func_rules.getReason(buf);
            if (reason) sendMsg(config.outfd, ID_DATA_STRING, INFO, ID_INFO, reason);
            if (!config.func_rules.isValid(buf)) continue;

            if (config.transMode == TRANS_FUNCTION)
               if (!dynStartTransaction(dh)) continue;

            sendMsg(config.outfd, ID_INST_FUNC, VERB1, ID_TEST, buf);

            if (config.inst_level == INST_FUNC_ENTRY) {
               if (!instrumentFunctionEntry(dh, (*appFunctions)[j]))
                  ++func_warn_cnt;
               else
                  ++func_pass_cnt;
            }

            if (config.inst_level == INST_FUNC_EXIT) {
               if (!instrumentFunctionExit(dh, (*appFunctions)[j]))
                  ++func_warn_cnt;
               else
                  ++func_pass_cnt;
            }

            if (config.inst_level == INST_BASIC_BLOCK) {
               if (!instrumentBasicBlocks(dh, (*appFunctions)[j]))
                  ++func_warn_cnt;
               else
                  ++func_pass_cnt;
            }

            if (config.inst_level == INST_MEMORY_READ) {
               if (!instrumentMemoryReads(dh, (*appFunctions)[j]))
                  ++func_warn_cnt;
               else
                  ++func_pass_cnt;
            }

            if (config.inst_level == INST_MEMORY_WRITE) {
               if (!instrumentMemoryWrites(dh, (*appFunctions)[j]))
                  ++func_warn_cnt;
               else
                  ++func_pass_cnt;
            }

            if (config.transMode == TRANS_FUNCTION)
               if (!dynEndTransaction(dh)) {
                  func_warn_cnt = func_pass_cnt;
                  func_pass_cnt = 0;
               }

            if (func_warn_cnt) {
               sendMsg(config.outfd, ID_INST_FUNC, VERB1, ID_WARN,
                       sprintf_static("%d warning(s), %d passed.", func_warn_cnt, func_pass_cnt));
               ++mod_warn_cnt;
            } else {
               sendMsg(config.outfd, ID_INST_FUNC, VERB1, ID_PASS);
               ++mod_pass_cnt;
            }
         }

         if (config.transMode == TRANS_MODULE)
            if (!dynEndTransaction(dh)) {
               mod_warn_cnt = mod_pass_cnt;
               mod_pass_cnt = 0;
            }

         if (mod_warn_cnt)
            sendMsg(config.outfd, ID_INST_MODULE, INFO, ID_WARN,
                    sprintf_static("%d warning(s), %d passed.", mod_warn_cnt, mod_pass_cnt));
         else
            sendMsg(config.outfd, ID_INST_MODULE, INFO, ID_PASS);
      }

      if (config.transMode == TRANS_PROCESS)
         dynEndTransaction(dh);
   }

   if (config.hunt_crashes) {
      sendMsg(config.outfd, ID_CRASH_HUNT_NUM_ACTIONS, INFO, ID_INFO, numInstsAllowed);
   }

   if (!config.use_process)
   {
      BPatch_binaryEdit *writeBE = dynamic_cast<BPatch_binaryEdit *>(dh->addSpace);

      // Load symbol libraries (if necessary)
      if( writeBE != NULL ) {
          deque<string>::iterator symLib_iter;
          for(symLib_iter = config.symbol_libraries.begin();
              symLib_iter != config.symbol_libraries.end();
              ++symLib_iter)
          {
              if( !writeBE->loadLibrary((*symLib_iter).c_str()) ) {
                  sendMsg(config.outfd, ID_INST_FIND_POINTS, VERB3, ID_FAIL,
                          "Unable to load symbol library");
                  return false;
              }
          }
      }

      writeBE->writeFile(config.writeFilePath);
   }

   if (config.use_process)
   {
      sendMsg(config.outfd, ID_RUN_CHILD, INFO);
      if (!dh->proc->continueExecution()) {
         sendMsg(config.outfd, ID_RUN_CHILD, INFO, ID_FAIL,
                 "Failure in BPatch_thread::continueExecution()");
         return false;
      }
      sendMsg(config.outfd, ID_RUN_CHILD, INFO, ID_PASS);
      
    
      //
      // Child continued.  Start reading from trace pipe, if enabled.
      //
      if (config.trace_inst) {
         while (config.pipefd != -1) {
            sendMsg(config.outfd, ID_POLL_STATUS_CHANGE, DEBUG);
            bool change = dh->bpatch->pollForStatusChange();
            sendMsg(config.outfd, ID_POLL_STATUS_CHANGE, DEBUG, ID_PASS);

            // Recheck conditional if a change was detected.
            if (change) continue;

            readTracePipe();

            // Eeeew.  I know.  But how else do you wait on a file descriptor,
            // AND BPatch::pollforStatusChange() == true at the same time?

            // We should probably have BPatch::registerStatusChangeCallback()
            // or something similar.
            sleep(1);
         }
      }
      

      //
      // All processing complete.  Loop indefinitly until exit handler called.
      //
      sendMsg(config.outfd, ID_WAIT_TERMINATION, INFO);
      while (1) {
         sendMsg(config.outfd, ID_CHECK_TERMINATION, INFO);
         bool dead = dh->proc->isTerminated();
         sendMsg(config.outfd, ID_CHECK_TERMINATION, INFO, ID_PASS);
         if (dead) break;

         sendMsg(config.outfd, ID_WAIT_STATUS_CHANGE, DEBUG);
         if (!dh->bpatch->waitForStatusChange())
            sendMsg(config.outfd, ID_WAIT_STATUS_CHANGE, DEBUG, ID_FAIL);
         else
            sendMsg(config.outfd, ID_WAIT_STATUS_CHANGE, DEBUG, ID_PASS);
      }
      sendMsg(config.outfd, ID_WAIT_TERMINATION, INFO, ID_PASS);
   }
   return 0;
}
Ejemplo n.º 3
0
int main(int argc, const char *argv[]) {

  // Use BPatch_* classes to initialize
  BPatch bpatch;
  BPatch_binaryEdit* app = bpatch.openBinary("mutatee/c");
  BPatch_image* image = app->getImage();


  BPatch_Vector<BPatch_function *> found_funcs;
  app->loadLibrary("mutatee/liblib.so");

  found_funcs.clear();
  image->findFunction("foo3", found_funcs);
  Function* foo3_func = found_funcs[0]->getParseAPIFunc();

  // Here we go, create PatchAPI objects!
  vector<AddressSpace*> addrSpaces;
  app->getAS(addrSpaces);

  mapped_object* obj = addrSpaces[0]->getAOut();
  DynAddrSpacePtr as = DynAddrSpace::create(obj);
  PatchMgrPtr mgr = PatchMgr::create(as);

  mapped_object* lib_obj = addrSpaces[1]->getAOut();
  as->loadLibrary(lib_obj);

  // Find Points
  PatchFunction* foo3 = lib_obj->getFunc(foo3_func);
  const vector<PatchBlock*>& blks = foo3->getCallBlocks();
  for (int i = 0; i < blks.size(); i++) {
    vector<Point*> func_points;
    mgr->findPoints(blks[i], Point::PreInsn|Point::PostInsn, inserter(func_points, func_points.begin()));
    cerr << std::hex << blks[i]->start() << "--" << func_points.size() << " points found\n";
  }
  /*
  vector<Point*> pts;
  mgr->findPoints(foo3, Point::FuncExit, inserter(pts, pts.begin()));
  cerr << pts.size() << " exit points found\n";
  const vector<PatchBlock*>& blks2 = foo3->getExitBlocks();
  cerr << blks2.size() << " exit blocks\n";

  // Insert snippets
  BPatch_variableExpr *intCounter = app->malloc(*image->findType("int"));
  BPatch_arithExpr addOne(BPatch_assign, *intCounter,
                          BPatch_arithExpr(BPatch_plus, *intCounter, BPatch_constExpr(1)));
  BPatch_arithExpr addTwo(BPatch_assign, *intCounter,
                          BPatch_arithExpr(BPatch_plus, *intCounter, BPatch_constExpr(2)));
  BPatch_arithExpr addThree(BPatch_assign, *intCounter,
                            BPatch_arithExpr(BPatch_plus, *intCounter, BPatch_constExpr(3)));
  BPatch_arithExpr addFour(BPatch_assign, *intCounter,
                           BPatch_arithExpr(BPatch_plus, *intCounter, BPatch_constExpr(4)));

  SnippetRep<AstNodePtr> one(addOne.ast_wrapper);
  SnippetRep<AstNodePtr> two(addTwo.ast_wrapper);
  SnippetRep<AstNodePtr> three(addThree.ast_wrapper);
  SnippetRep<AstNodePtr> four(addFour.ast_wrapper);
  SnippetPtr snippet = Snippet::create(&one);
  SnippetPtr snippet1 = Snippet::create(&two);
  SnippetPtr snippet2 = Snippet::create(&three);
  SnippetPtr snippet3 = Snippet::create(&four);

  vector<InstancePtr> errorInstances;

  mgr->batchStart();
  func_points[0]->push_back(snippet);
  mgr->batchFinish(errorInstances);
  */
}
Ejemplo n.º 4
0
int main(int argc, char **argv) {
  if (argc < 3 || strncmp(argv[1], "-h", 2) == 0 || strncmp(argv[1], "--h", 3) == 0) {
    cout << "Usage: " << argv[0] << USAGE;
    return false;
  }

  if (!parseOptions(argc, argv)) {
    return EXIT_FAILURE;
  }
  
  if (do_bb == true) {
    if (DYNINST_MAJOR_VERSION < 9 || (DYNINST_MAJOR_VERSION == 9 && DYNINST_MINOR_VERSION < 3) || (DYNINST_MAJOR_VERSION == 9 && DYNINST_MINOR_VERSION == 3 && DYNINST_PATCH_VERSION <= 2)) {
      if (dynfix == false)
        fprintf(stderr, "Warning: your dyninst version does not include a critical fix, you should use the -f option!\n");
    } else {
      if (dynfix == true)
        fprintf(stderr, "Notice: your dyninst version is fixed, the -f option should not be necessary.\n");
    }
  }

  BPatch bpatch;
  BPatch_binaryEdit *appBin = bpatch.openBinary(originalBinary, instrumentLibraries.size() != 1);

  if (appBin == NULL) {
    cerr << "Failed to open binary" << endl;
    return EXIT_FAILURE;
  }

  BPatch_image *appImage = appBin->getImage();

  //get and iterate over all modules, instrumenting only the default and manually specified ones
  vector < BPatch_module * >*modules = appImage->getModules();
  vector < BPatch_module * >::iterator moduleIter;
  vector < BPatch_function * >*funcsInModule;
  BPatch_module *defaultModule = NULL;
  string defaultModuleName;

  // look for _init
  if (defaultModuleName.empty()) {
    for (moduleIter = modules->begin(); moduleIter != modules->end(); ++moduleIter) {
      funcsInModule = (*moduleIter)->getProcedures();
      vector < BPatch_function * >::iterator funcsIterator;
      for (funcsIterator = funcsInModule->begin(); funcsIterator != funcsInModule->end(); ++funcsIterator) {
        char funcName[1024];

        (*funcsIterator)->getName(funcName, 1024);
        if (string(funcName) == string("_init")) {
          char moduleName[1024];

          (*moduleIter)->getName(moduleName, 1024);
          defaultModuleName = string(moduleName);
          if (verbose) {
            cout << "Found _init in " << moduleName << endl;
          }
          break;
        }
      }
      if (!defaultModuleName.empty())
        break;
    }
  }
  // last resort, by name of the binary
  if (defaultModuleName.empty())
    defaultModuleName = string(originalBinary).substr(string(originalBinary).find_last_of("\\/") + 1);

  if (!appBin->loadLibrary(instLibrary)) {
    cerr << "Failed to open instrumentation library " << instLibrary << endl;
    cerr << "It needs to be located in the current working directory." << endl;
    return EXIT_FAILURE;
  }

  appImage = appBin->getImage();

  /* Find code coverage functions in the instrumentation library */
  BPatch_function *initAflForkServer;
  save_rdi = findFuncByName(appImage, (char *) "save_rdi");
  restore_rdi = findFuncByName(appImage, (char *) "restore_rdi");
  BPatch_function *bbCallback = findFuncByName(appImage, (char *) "bbCallback");
  BPatch_function *forceCleanExit = findFuncByName(appImage, (char *) "forceCleanExit");

  if (do_bb == true)
    initAflForkServer = findFuncByName(appImage, (char *) "initAflForkServer");
  else
    initAflForkServer = findFuncByName(appImage, (char *) "initOnlyAflForkServer");

  if (!initAflForkServer || !bbCallback || !save_rdi || !restore_rdi || !forceCleanExit) {
    cerr << "Instrumentation library lacks callbacks!" << endl;
    return EXIT_FAILURE;
  }

  int bbIndex = 0;

  // instrument all shared libraries:
  for (moduleIter = modules->begin(); moduleIter != modules->end(); ++moduleIter) {
    char moduleName[1024];

    (*moduleIter)->getName(moduleName, 1024);

    if ((*moduleIter)->isSharedLib()) {
      if (instrumentLibraries.find(moduleName) == instrumentLibraries.end()) {
        cout << "Skipping library: " << moduleName << endl;
        continue;
      }
    }

    if (string(moduleName).find(defaultModuleName) != string::npos) {
      defaultModule = (*moduleIter);
      if (skipMainModule)
        continue;
    }
    
    if (do_bb == true) {
      cout << "Instrumenting module: " << moduleName << endl;
      vector < BPatch_function * >*allFunctions = (*moduleIter)->getProcedures();
      vector < BPatch_function * >::iterator funcIter;

      // iterate over all functions in the module
      for (funcIter = allFunctions->begin(); funcIter != allFunctions->end(); ++funcIter) {
        BPatch_function *curFunc = *funcIter;
        char funcName[1024];
        int do_patch = 1;

        curFunc->getName(funcName, 1024);
        if (string(funcName) == string("_start"))
          continue;               // here's a bug on hlt // XXX: check what happens if removed
        
        if (!skipAddresses.empty()) {
          set < string >::iterator saiter;
          for (saiter = skipAddresses.begin(); saiter != skipAddresses.end() && do_patch == 1; saiter++)
            if (*saiter == string(funcName))
              do_patch = 0;
          if (do_patch == 0) {
            cout << "Skipping instrumenting function " << funcName << endl;
            continue;
          }
        }
        insertBBCallback(appBin, curFunc, funcName, bbCallback, &bbIndex);
      }
    }
  }

  // if an entrypoint was set then find function, else find _init
  BPatch_function *funcToPatch = NULL;

  if (!entryPoint) {
    BPatch_Vector < BPatch_function * >funcs;
    defaultModule->findFunction("_init", funcs);
    if (!funcs.size()) {
      cerr << "Couldn't locate _init, specify entry point manually with -e 0xaddr" << endl;
      return EXIT_FAILURE;
    }
    // there should really be only one
    funcToPatch = funcs[0];
  } else {
    funcToPatch = defaultModule->findFunctionByEntry(entryPoint);
  }
  if (!funcToPatch) {
    cerr << "Couldn't locate function at given entry point. " << endl;
    return EXIT_FAILURE;
  }
  if (!insertCallToInit(appBin, initAflForkServer, defaultModule, funcToPatch)) {
    cerr << "Could not insert init callback at given entry point." << endl;
    return EXIT_FAILURE;
  }

  if (!exitAddresses.empty()) {
    cout << "Instrumenting forced exit addresses." << endl;
    set < unsigned long >::iterator uliter;
    for (uliter = exitAddresses.begin(); uliter != exitAddresses.end(); uliter++) {
      if (*uliter > 0 && (signed long)*uliter != -1) {
        funcToPatch = defaultModule->findFunctionByEntry(*uliter);
        if (!funcToPatch) {
          cerr << "Could not find enty point 0x" << hex << *uliter << " (continuing)" << endl;
        } else {
          if (!insertCallToInit(appBin, forceCleanExit, defaultModule, funcToPatch))
            cerr << "Could not insert force clean exit callback at 0x" << hex << *uliter << " (continuing)" << endl;
        }
      }
    }
  }

  cout << "Saving the instrumented binary to " << instrumentedBinary << " ..." << endl;
  // Output the instrumented binary
  if (!appBin->writeFile(instrumentedBinary)) {
    cerr << "Failed to write output file: " << instrumentedBinary << endl;
    return EXIT_FAILURE;
  }

  if (!runtimeLibraries.empty()) {
    cout << "Instrumenting runtime libraries." << endl;
    set < string >::iterator rtLibIter;
    for (rtLibIter = runtimeLibraries.begin(); rtLibIter != runtimeLibraries.end(); rtLibIter++) {
      BPatch_binaryEdit *libBin = bpatch.openBinary((*rtLibIter).c_str(), false);

      if (libBin == NULL) {
        cerr << "Failed to open binary " << *rtLibIter << endl;
        return EXIT_FAILURE;
      }
      BPatch_image *libImg = libBin->getImage();

      vector < BPatch_module * >*modules = libImg->getModules();
      moduleIter = modules->begin();

      for (; moduleIter != modules->end(); ++moduleIter) {
        char moduleName[1024];

        (*moduleIter)->getName(moduleName, 1024);
        cout << "Instrumenting module: " << moduleName << endl;
        vector < BPatch_function * >*allFunctions = (*moduleIter)->getProcedures();
        vector < BPatch_function * >::iterator funcIter;
        // iterate over all functions in the module
        for (funcIter = allFunctions->begin(); funcIter != allFunctions->end(); ++funcIter) {
          BPatch_function *curFunc = *funcIter;
          char funcName[1024];

          curFunc->getName(funcName, 1024);
          if (string(funcName) == string("_start"))
            continue;
          insertBBCallback(libBin, curFunc, funcName, bbCallback, &bbIndex);
        }
      }
      if (!libBin->writeFile((*rtLibIter + ".ins").c_str())) {
        cerr << "Failed to write output file: " << (*rtLibIter + ".ins").c_str() << endl;
        return EXIT_FAILURE;
      } else {
        cout << "Saved the instrumented library to " << (*rtLibIter + ".ins").c_str() << "." << endl;
      }
    }
  }

  cout << "All done! Happy fuzzing!" << endl;
  return EXIT_SUCCESS;
}