def reload_statistics(self):
   if self.state_file is not None:
     with open(self.state_file, "rb") as f:
       self.stats = json.loads(f.read())
   line = "Reloaded statistics: Min %d, Max %d, Avg %f"
   line = line % (self.stats["min"], self.stats["max"], self.stats["avg"])
   log(line)

  def debug_server(self, shared_queue):
    self.read_configuration()

    uid = int(self.server_uid)
    if os.getuid() != uid:
      os.setresuid(uid, uid, uid)

    gid = int(self.server_gid)
    if os.getgid() != gid:
      os.setresgid(gid, gid, gid)

    for key in self.env:
      debug("Setting environment variable %s=%s" % (key, self.env[key]))
      os.putenv(key, self.env[key])

    if self.pre_command is not None:
      os.system(self.pre_command)

    crash = None
    for i in range(0,3):
      try:
        crash = self.launch_debugger(self.timeout, self.command, "")
        break
      except:
        log("Exception: %s" % sys.exc_info()[1])
        continue

    if self.post_command is not None:
      os.system(self.post_command)

    if crash is not None:
      self.crash_info = crash
      shared_queue.put(crash)
      return True
    return False

  def __init__(self, tube_prefix):
    if tube_prefix.endswith("-samples"):
      log("Notice: Removing '-samples' suffix from the queue name")
      tube_prefix = tube_prefix.replace("-samples", "")

    self.tube_prefix = tube_prefix
    self.q = get_queue(watch=False, name="%s-samples" % tube_prefix)

  def run(self, timeout=60):
    def target():
      debug('Thread started')
      self.process = subprocess.Popen("exec %s" % self.cmd, shell=True)
      self.process.communicate()
      debug('Thread finished')

    thread = threading.Thread(target=target)
    thread.start()

    thread.join(timeout)
    if thread.is_alive():
      log('Terminating process after timeout (%d)' % timeout)
      try:
        self.process.terminate()
        self.process.terminate()
        self.process.kill()
        self.process.wait()
      except:
        log("Error killing process: %s" % str(sys.exc_info()[1]))

      thread.join()
    self.process.wait()
    ret = self.process.returncode

    # A negative return code means a signal was received and the return
    # code is -1 * SIGNAL. Return the expected Unix return code.
    if ret is not None and ret < 0:
      ret = abs(ret) + 128
    return ret

 def resolve_windbg_path(self):
   try:
     reg = ConnectRegistry(None,HKEY_LOCAL_MACHINE)
     key = OpenKey(reg, r"SOFTWARE\Microsoft\Microsoft SDKs\Windows")
     if key:
       for i in range(QueryInfoKey(key)[0]):
         value = EnumKey(key, i)
         if value:
           full_key = r"SOFTWARE\Microsoft\Microsoft SDKs\Windows\\" + value
           key2 = OpenKey(reg, full_key)
           if key2:
             name = QueryValueEx(key2, "ProductName")
             name = name[0]
             if name and name.startswith("Microsoft Windows SDK for Windows"):
               vals = QueryValueEx(key2, "InstallationFolder")
               val = vals[0]
               if val is not None:
                 log("Found installation path at %s" % val)
                 self.windbg_path = val
                 break
             CloseKey(key2)
     CloseKey(key)
   except WindowsError:
     print "Cannot resolve Windows SDKs path:", sys.exc_info()[1]
     print "Did you install Windows SDKs for Windows?"
   except:
     print "Cannot resolve Windows SDKs path:", sys.exc_info()[1]

  def check_cpu(self):
    while True:
      try:
        if self.pid is None:
          time.sleep(0.2)
          continue

        proc = psutil.Process(self.pid)
        if proc is None:
          break

        cpu = 0
        l = []
        for x in xrange(20):
          tmp = int(proc.cpu_percent(interval=0.1))
          cpu += tmp
          l.append(tmp)

        if cpu is not None and (cpu <= 100 or l.count(0) > 10):
          log("CPU at 0%, killing")
          self.do_stop = True
          pykd.breakin()
          break
        else:
          time.sleep(0.5)
      except psutil.NoSuchProcess:
        self.do_stop = True
        break

def do_fuzz(cfg, section):
  try:
    fuzzer = CGenericFuzzer(cfg, section)
    fuzzer.fuzz()
  except KeyboardInterrupt:
    log("Aborted")
  except:
    log("Error: %s" % str(sys.exc_info()[1]))

 def timeout_func(self):
   log("Timeout (%d seconds), killing the target..." % self.timeout)
   self.do_stop = True
   try:
     pykd.breakin()
   except:
     # A race condition might happen in the timeout function and in 
     # such cases we must ignore the error.
     pass

 def launch_debugger(self, timeout, command, filename):
   self.iface.timeout = int(timeout)
   
   if command.find("@@") > -1:
     cmd = [command.replace("@@", filename), ]
   else:
     cmd = [command, filename]
   
   log("Launching debugger with command %s" % " ".join(cmd))
   crash = self.iface.main(cmd)
   return crash

  def get_html_buffer(self, lines, skip_tags):
    log("Rewriting HTML...")
    rewriter = CHTMLRewriter()
    ret = rewriter.rewrite("\n".join(lines), skip_tags)
    if not ret:
      return None, None

    lines = ret
    log("Total line(s) %d" % len(lines))

    return "".join(lines), lines

  def run(self, timeout=60, get_output=False):
    def target():
      debug('Thread started')
      if os.name == "nt":
        line = self.cmd
        shell = False
      else: # Unix based
        line = "exec %s" % self.cmd
        shell = True
      
      if get_output:
        self.process = subprocess.Popen(line, stdout=subprocess.PIPE,\
                                      stderr=subprocess.PIPE, shell=shell)
        self.pid = self.process.pid
        out, err = self.process.communicate()
        self.stdout = out[:8192]
        self.stderr = err[:8192]
      else:
        self.process = subprocess.Popen(line, shell=shell)
        self.pid = self.process.pid
        self.process.communicate()

      debug('Thread finished')

    thread = threading.Thread(target=target)
    thread.start()

    if str(timeout).lower() == "auto":
      self.thread = threading.Thread(target=self.check_cpu)
      self.thread.start()
      thread.join(self.default_timeout)
    else:
      thread.join(timeout)

    if thread.is_alive():
      log('Terminating process after timeout (%s)' % str(timeout))
      try:
        self.do_kill()
      except:
        log("Error killing process: %s" % str(sys.exc_info()[1]))

      thread.join()

    self.process.wait()
    ret = self.process.returncode

    # A negative return code means a signal was received and the return
    # code is -1 * SIGNAL. Return the expected Unix return code.
    if ret is not None and ret < 0:
      if os.name == "nt":
        ret = ret & 0xFFFFFFFF
      else:
        ret = abs(ret) + 128
    return ret

  def recalculate_statistics(self, old_stats, bbs):
    self.stats["max"] = bbs
    self.stats["min"] = old_stats["max"]
    self.stats["avg"] = (self.stats["max"] + self.stats["min"]) / 2.
    #self.stats = self.mgr.dict(self.stats)
    line = "New statistics: Min %d, Max %d, Avg %f"
    line = line % (self.stats["min"], self.stats["max"], self.stats["avg"])
    log(line)

    if self.state_file is not None:
      with open(self.state_file, "wb") as f:
        f.write(json.dumps(dict(self.stats)))

  def minimize(self, template, outdir):
    self.read_template(template)

    log("Performing line-level test case minimization")
    start_at = os.getenv("NIGHTMARE_ITERATION")
    if start_at is not None:
      start_at = int(start_at)
      log("Starting from iteration %d\n" % start_at)
    else:
      start_at = 0

    self.do_try(outdir, start_at)

  def apply_bytes(self, offset, size, buf):
    debug("Acquiring lock...")
    self.lock.acquire()
    try:
      debug("Saving old generation (%s)" % sha1(self.template).hexdigest())
      if len(self.generations) >= self.max_generations:
        del self.generations[0]
      self.generations.append([bytearray(self.template), dict(self.stats), self.generation_value])
      
      if self.save_generations:
        file_hash = sha1(buf).hexdigest()
        ext = os.path.splitext(self.input_file)[1]
        filename = "generation_%s%s" % (file_hash, ext)
        filename = os.path.join(self.output, filename)
        log("Writing discovered generation file %s (%s)" % (file_hash, filename))

        with open(filename, "wb") as f:
          f.write(buf)

      if not self.radamsa:
        debug("Applying patch at offset %d of size %d" % (offset, size))
      else:
        debug("Replacing old buffer")

      self.template = buf
      """
      if self.skip_bytes > 0:
        header = self.template[0:self.skip_bytes]

      if len(buf) > len(self.template):
        self.template = bytearray(buf)
      else:
        for i in range(size):
          self.template[offset+i] = buf[i]

      if self.skip_bytes > 0:
        self.template[0:self.skip_bytes] = header
      """

      if self.current_state is not None:
        ext = os.path.splitext(self.input_file)[1]
        filename = "%s%s" % (self.current_state, ext)
        filename = os.path.join(self.output, filename)
        file_hash = sha1(self.template).hexdigest()

        debug("Creating or updating current state file %s (%s)" % (filename, file_hash))
        with open(filename, "wb") as f:
          f.write(self.template)

    finally:
      debug("Releasing lock...")
      self.lock.release()

  def disasm_around(self):
    try:
      lines = pykd.dbgCommand("u %s-c L12" % self.pc_register)
      for line in lines.split("\n"):
        tmp = re.findall("([a-f0-9]{1,}) ([a-f0-9]{2,}) (.*)", line)
        if len(tmp) > 0:
          line = tmp[0]

          addr = line[0]
          dis = line[2]
          self.crash_data.add_data("disassembly", int(addr, 16), dis)
    except:
      log("Error in disasm_around: %s" % str(sys.exc_info()[1]))

  def fuzz(self):
    log("Launching fuzzer, listening in tube %s" % self.tube_name)
    while 1:
      value = self.q.stats_tube(self.tube_name)["current-jobs-ready"]
      debug("Total of %d job(s) in queue" % value)
      job = self.q.reserve()
      buf, temp_file = json.loads(job.body)
      buf = base64.b64decode(buf)

      debug("Launching sample %s..." % os.path.basename(temp_file))
      if self.launch_sample(buf):
        log("We have a crash, moving to %s queue..." % self.crash_tube)
        crash = self.crash_info
        d = {temp_file:self.crash_info}
        self.crash_q.put(json.dumps(d))
        self.crash_info = None

        log("$PC 0x%08x Signal %s Exploitable %s " % (crash["pc"], crash["signal"], crash["exploitable"]))
        if crash["disasm"] is not None:
          log("%08x: %s" % (crash["disasm"][0], crash["disasm"][1]))
      else:
        file_delete = os.path.basename(temp_file)
        self.delete_q.put(str(file_delete))
      
      if self.cleanup is not None:
        debug("Running clean-up command %s" % self.cleanup)
        os.system(self.cleanup)
        debug("Done")
      job.delete()
      
      if self.iface == gdb_iface:
        break

 def launch_debugger(self, timeout, command, filename):
   self.iface.timeout = int(timeout)
   
   if command.find("@@") > -1:
     cmd = [command.replace("@@", filename), ]
   else:
     cmd = [command, filename]
   
   log("Launching debugger with command %s" % " ".join(cmd))
   if self.iface != pykd_iface:
     crash = self.iface.main(" ".join(cmd))
   else:
     reload(pykd_iface)
     crash = pykd_iface.main(cmd, mode=self.mode, windbg_path=self.windbg_path, exploitable_path=self.exploitable_path)
   return crash

  def launch_target(self, temp_file, lines, outdir):
    try:
      crashed = False

      for key in self.env:
        os.putenv(key, self.env[key])

      self.remove_crash_path()

      if self.pre_command is not None:
        log("Running pre-command %s" % self.pre_command)
        os.system(self.pre_command)

      if self.command.find("@@") == -1:
        cmd = "%s %s" % (self.command, temp_file)
      else:
        cmd = self.command.replace("@@", temp_file)
      ret = self.execute_command(cmd, self.timeout)

      if self.post_command is not None:
        log("Running post-command %s" % self.post_command)
        os.system(self.post_command)

      if ret in RETURN_SIGNALS or (self.signal is not None and ret == self.signal) or \
         self.crash_file_exists():
          
        crashed = True
        self.template = lines
        log("Process crashed as expected...")
        buf = "\n".join(self.template)
        if not os.path.exists(outdir):
          log("Directory %s does not exists, creating it..." % outdir)
          os.mkdir(outdir)

        filename = os.path.join(outdir, "last_minimized%d%s" % (os.getpid(), self.extension))
        with open(filename, "wb") as f:
          f.write(buf)
        log("Last minimized test case %s written to disk." % filename)

        if self.should_notify_crash():
          # TODO: Write a temporary file and put an enqueue the crash
          self.put_new_crash(buf)

      self.remove_crash_path()
    finally:
      os.remove(temp_file)
    
    return crashed

  def put(self, filename):
    temp_file = tempfile.mktemp()
    
    try:
      buf = open(filename, "rb").read()
      
      with open(temp_file, "wb") as f:
        f.write(buf)

      json_buf = json.dumps([base64.b64encode(buf), temp_file])
      self.q.put(json_buf)
      l = "File '%s' put in queue %s as temporary file '%s'"
      log(l % (filename, self.tube_prefix, temp_file))
    except:
      raise
      os.remove(temp_file)

  def __init__(self, arch, cfg, section, metrics=10):
    if int(arch) not in [32, 64]:
      raise Exception("Invalid architecture %s" % str(arch))

    self.arch = arch
    self.cfg = cfg
    self.section = section
    self.metrics = 10
    self.bininst_tool=None

    self.mgr = Manager()
    self.stats = self.mgr.dict()

    self.read_configuration()

    cpus = os.getenv("NIGHTMARE_PROCESSES")
    if cpus is not None:
      cpus = int(cpus)
    else:
      cpus = cpu_count()
    self.procs = cpus

    self.discard_data = self.mgr.list()

    if self.procs > self.metrics:
      log("The number of processes is bigger than the number of metrics, adjusting it to %d" % self.procs)
      self.metrics = self.procs

    # Default output directory is current path
    self.output = "."
    self.input_file = None
    self.is_dir = False

    # Maximum number of bytes to mutate per each try
    self.max_size = random.randint(1, 8)
    log("Selected a maximum size of %d change(s) to apply" % self.max_size)
    
    # Only for the iterative mutator
    self.stats["iteration"] = 0
    self.stats["iteration_char"] = 0
    
    self.generations = self.mgr.list()
    self.generation_value = 0
    self.max_generations = 10
    
    self.bugs = 0