text: ############################################
  # Copyright (C) 2019 FireEye, Inc.
  #
  # Author: Moritz Raabe
  #
  # ironstrings is an IDAPython script that uses flare-emu (which combines Unicorn and
  # IDA Pro) to recover stackstrings from malware using code emulation.
  #
  # Dependencies:
  # https://github.com/fireeye/flare-emu
  # https://github.com/unicorn-engine/unicorn
  ############################################
  
  from __future__ import print_function
  
  import time
  import logging
  from collections import namedtuple
  
  import idc
  import idaapi
  import idautils
  import unicorn
  
  import strings
  try:
      import flare_emu
  except ImportError as e:
      print("Could not import flare_emu: {}\nExiting.".format(e.message))
      raise
  
  
  #
  # USER OPTIONS
  #
  
  # only analyze the currently selected function; otherwise analyze all functions
  ANALYZE_SINGLE_FUNC = False
  
  # jump to currently analyzed function
  JUMP_TO_FUNC = False
  
  # print listing of unique stackstrings after running
  PRINT_PLAIN_SUMMARY = True
  
  # add comments for recovered stackstrings
  COMMENT_STACKSTRINGS = True
  
  # use a repeatable comment on the global memory address
  COMMENT_STACKSTRING_GLOBAL_REPEATABLE = False
  
  # prefix used for stackstring comment
  COMMENT_STACKSTRING_PREFIX = "stackstring: '"
  
  # suffix used for stackstring comment
  COMMENT_STACKSTRING_SUFFIX = "'"
  
  #
  # STACKSTRINGS RECOVERY OPTIONS
  #
  
  # maximum number of code paths to find through each function
  MAX_CODE_PATHS = 10
  
  # minimum number of memory writes in a basic block to trigger stackstring recovery
  MIN_MEM_WRITE_COUNT = 5
  
  # maximum number an instruction is repeated
  MAX_INSTR_REP = 0x100
  
  # maximum size of current stack frame
  # IMPROVEMENT fine tune this value
  MAX_STACK_SIZE = 0x10000
  
  
  StackString = namedtuple("StackString",
                           [
                               # function address from which the stackstring was extracted
                               "fva",
                               # program counter at which the stackstring existed
                               "pc",
                               # string contents
                               "s",
                               # for local stackstring: offset from the function frame at which the stackstring
                               # existed, for globalstring: offset in global memory
                               "offset",
                               # address where first character of stackstring was written
                               "written_at",
                           ])
  
  
  def main():
      eh = flare_emu.EmuHelper()
  
      # dictionary that stores data used across emulation runs, function emulation specific data is set below
      userData = {
          # found stackstrings in stack memory
          "stackstrings": [],
          # found stackstrings in global memory (globalstrings)
          "globalstrings": []
      }
  
      cnt_functions = 0
      cnt_analyzed = 0
      cnt_found_ss = 0
      cnt_commented_ss = 0
      errors = []
  
      start = time.time()
      print("Started ironstrings stackstring deobfuscation")
      print_header()
  
      if ANALYZE_SINGLE_FUNC:
          fvas = [idc.get_func_attr(idc.here(), idc.FUNCATTR_START)]
      else:
          fvas = idautils.Functions()
  
      for fva in fvas:
          logging.debug("running on 0x%X", fva)
          if JUMP_TO_FUNC:
              idc.jumpto(fva)
  
          if fva == idaapi.BADADDR:
              logging.debug("skipping invalid function address")
              continue
          if idc.get_func_flags(fva) & (idc.FUNC_LIB | idc.FUNC_THUNK):
              logging.debug("skipping library or thunk function 0x%X", fva)
              continue
  
          # function start address
          userData["funcStart"] = fva
  
          # list of addresses of last instruction for all basic blocks in function
          userData["bb_ends"] = get_bb_ends(fva)
  
          # memory writes in current function
          userData["mem_writes"] = {}
  
          # start and end addresses of all memory writes in function
          userData["writelog"] = set()
  
          # memory write count in current basic block
          userData["mem_write_count"] = 0
  
          # cache previous address to count instructions that are executed multiple times, e.g. rep prefixed
          userData["prevAddress"] = 0
  
          # number same instruction has been executed in a row
          userData["repCount"] = 0
  
          cnt_functions += 1
          try:
              # emulate various paths through function via flare-emu, use hooks to reconstruct strings
              eh.iterateAllPaths(fva, noop, hookData=userData, callHook=call_hook, instructionHook=instr_hook,
                                 memAccessHook=hook_mem_write, hookApis=False, maxPaths=MAX_CODE_PATHS)
          except unicorn.UcError as e:
              errors.append("Error analyzing function 0x{:X}: {}".format(fva, str(e)))
          else:
              cnt_analyzed += 1
  
              # print stackstrings found in this function
              f_ss = filter(lambda s: s.fva == fva, userData["stackstrings"])
              cnt_found_ss += len(f_ss)
              for ss in sorted(f_ss, key=lambda s: s.written_at):
                  print_string(ss.fva, ss.written_at, ss.offset, ss.s)
                  # IMPROVEMENT adjust stack frame member size in IDA view
  
              # print globalstrings found in this function
              f_gs = filter(lambda g: g.fva == fva, userData["globalstrings"])
              cnt_found_ss += len(f_gs)
              for gs in sorted(f_gs, key=lambda g: g.written_at):
                  print_string(gs.fva, gs.written_at, gs.offset, gs.s)
  
              if COMMENT_STACKSTRINGS:
                  for ss in f_ss:
                      if not ss.written_at:
                          errors.append("Can't get location where '{}' was written in 0x{:X}.".format(ss.s, ss.fva))
                          continue
                      ss_cmt = format_comment(ss.s)
                      if append_comment(ss.written_at, ss_cmt):
                          cnt_commented_ss += 1
                      else:
                          errors.append("Failed to set comment at 0x{:X}: {}".format(ss.written_at, ss_cmt))
  
                  for gs in f_gs:
                      if COMMENT_STACKSTRING_GLOBAL_REPEATABLE:
                          repeatable = True
                          cmt_va = gs.offset
                      else:
                          repeatable = False
                          cmt_va = gs.written_at
  
                      if not cmt_va:
                          errors.append("Can't get location where '{}' was written in 0x{:X}.".format(gs.s, gs.fva))
                          continue
                      gs_cmt = format_comment(gs.s)
                      if append_comment(cmt_va, gs_cmt, repeatable):
                          cnt_commented_ss += 1
                      else:
                          errors.append("Failed to set comment at 0x{:X}: {}".format(cmt_va, gs_cmt))
  
          # update IDA view
          idc.refresh_idaview_anyway()
  
          # clean up memory after each function
          eh.resetEmulatorHeapAndStack()
  
      print_summary(cnt_functions, cnt_analyzed, cnt_found_ss, cnt_commented_ss, errors)
  
      if PRINT_PLAIN_SUMMARY:
          print_plain_summary(userData["stackstrings"] + userData["globalstrings"])
  
      print("\nFinished ironstrings stackstring deobfuscation after {:.2f} seconds".format(time.time() - start))
  
  
  def instr_hook(uc, address, size, userData):
      """
      Hook that runs before every instruction. Performs emulation sanity checks.
      Extract stackstrings at end of a basic block (bb) if bb contains enough memory writes.
      No success with UC_HOOK_BLOCK as bbs may differ between IDA and Unicorn
      :param uc: Unicorn object
      :param address: current instruction address
      :param size: current instruction size
      :param userData: user-provided data
      :return: None
      """
      update_init_sp(address, userData)
  
      # IMPROVEMENT generalize to count maxhits for all instructions
      if address == userData["prevAddress"]:
          userData["repCount"] += 1
      if userData["repCount"] > MAX_INSTR_REP:
          logging.debug("address hit more than %d times, skipping", MAX_INSTR_REP)
          userData["repCount"] = 0
          eh = userData["EmuHelper"]
          eh.skipInstruction(userData)
      userData["prevAddress"] = address
  
      # trigger stackstring extraction based on number of memory writes in current bb
      if address in userData["bb_ends"]:
          if userData["mem_write_count"] >= MIN_MEM_WRITE_COUNT:
              logging.debug("end of basic block with %d memory writes", userData["mem_write_count"])
              update_stackstrings(address, userData)
              update_globalstrings(address, userData)
          # reset counter at end of basic block
          userData["mem_write_count"] = 0
  
  
  def call_hook(address, argv, funcName, userData):
      """
      Hook all function calls. Extract stackstrings.
      :param address: address of call
      :param argv: call arguments
      :param funcName: name of called function
      :param userData: user-provided data
      :return: None
      """
      # IMPROVEMENT record funcName
      update_stackstrings(address, userData)
      update_globalstrings(address, userData)
      # reset memory write counter, all strings up to call have been captured
      userData["mem_write_count"] = 0
  
  
  def hook_mem_write(uc, access, address, size, value, userData):
      """
      Hook all memory read and write events. Count memory writes and record written address locations.
      :param uc: Unicorn object
      :param access: memory access type
      :param address: memory address
      :param size: memory address size
      :param value: memory value
      :param userData: user-provided data
      :return: None
      """
      eh = userData["EmuHelper"]
      pc = eh.getRegVal("pc")
  
      # only record and count memory writes, create write log
      if access == unicorn.UC_MEM_WRITE:
          # IMPROVEMENT store multiple locations where address was written
          userData["mem_writes"][address] = pc
          userData["mem_write_count"] += 1
          userData["writelog"].add((address, address + size))
  
  
  def noop(*args):
      """
      No operation callback
      :param args:
      :return: None
      """
  
  
  def get_bb_ends(address):
      """
      Get end addresses of all bbs in function containing address.
      :param address: address in function
      :return: list of bb end addresses
      """
      function = idaapi.get_func(address)
      flowchart = idaapi.FlowChart(function)
      return [idc.prev_head(bb.end_ea) for bb in flowchart]
  
  
  def update_init_sp(address, userData):
      """
      Hack to get stack pointer at beginning of emulation.
      :param address: current instruction address
      :param userData: user-provided data
      :return: None
      """
      eh = userData["EmuHelper"]
      if address == userData["funcStart"]:
          userData["init_sp"] = eh.getRegVal("sp")
  
  
  def update_stackstrings(address, userData):
      """
      Extract stackstrings and update userData value.
      :param address: current instruction address
      :param userData: user-provided data
      :return: None
      """
      eh = userData["EmuHelper"]
      stack_top = eh.getRegVal("sp")
      stack_start, stack_end = eh.getEmuMemRegion(eh.stack)
      if stack_top < stack_start:
          stack_top = stack_start
  
      if "init_sp" not in userData:
          logging.debug("could not get initial SP, using base of stack: 0x%X", stack_start)
          userData["init_sp"] = stack_start
      stack_bottom = userData["init_sp"]  # SP at start of function, most of the times equal to BP
      stackstrings = userData["stackstrings"]
      mem_writes = userData["mem_writes"]
  
      # Extract only the bytes on the stack between stack pointer at function entry and current stack pointer
      current_stack_size = stack_bottom - stack_top
      if current_stack_size < 0:
          logging.debug("can't read negative stack size: %d", current_stack_size)
          # IMPROVEMENT return false to indicate failure
          return
      if stack_top + current_stack_size > stack_end:
          logging.debug("can't read past stack end")
          return
      if current_stack_size > MAX_STACK_SIZE:
          logging.debug("stack size too big: 0x%X", current_stack_size)
          return
  
      stack_buf = eh.getEmuBytes(stack_top, current_stack_size)
      logging.debug("extracting stackstrings at checkpoint: 0x%X, stacksize: 0x%X", address, current_stack_size)
      for s in extract_all_strings(stack_buf):
          logging.debug("found stackstring '%s'", s.s)
          frame_offset = current_stack_size - s.offset - eh.size_pointer
          offset, written_at = get_offset_written_at(stack_top + s.offset, mem_writes)
          ss = StackString(userData["funcStart"], address, s.s[offset:], frame_offset, written_at)
          if does_contain(stackstrings, ss) or does_contain_substr(stackstrings, ss):
              # already recovered
              continue
          if not extend_existing(ss, stackstrings):
              stackstrings.append(ss)
  
  
  def update_globalstrings(address, userData):
      """
      Extract stackstrings in global memory and update userData value.
      :param address: current instruction address
      :param userData: user-provided data
      :return: None
      """
      globalstrings = userData["globalstrings"]
      mem_writes = userData["mem_writes"]
      # consolidate memory write ranges
      userData["writelog"] = set(consolidate(userData["writelog"]))
  
      eh = userData["EmuHelper"]
      stack_start, stack_end = eh.getEmuMemRegion(eh.stack)
  
      # ignore stack memory
      # IMPROVEMENT use this mechanism for all stackstrings
      userData["writelog"] = set(filter(lambda m: m[0] < stack_start or m[0] > stack_end, userData["writelog"]))
  
      logging.debug("extracting global stackstrings at checkpoint: 0x%X", address)
      for mem_start, mem_end in userData["writelog"]:
          bytez = eh.getEmuBytes(mem_start, mem_end - mem_start)
          for s in extract_all_strings(bytez):
              # ignore artifacts in null page
              if mem_start == 0:
                  continue
  
              logging.debug("found globalstring '%s' at 0x%X", s.s, mem_start + s.offset)
              _, written_at = get_offset_written_at(mem_start + s.offset, mem_writes)
              gs = StackString(userData["funcStart"], address, s.s, mem_start + s.offset, written_at)
              if does_contain(globalstrings, gs) or does_contain_substr(globalstrings, gs):
                  # already recovered
                  continue
              if not extend_existing(gs, globalstrings):
                  globalstrings.append(gs)
  
  
  def extract_all_strings(stack_buf):
      """
      Extract ASCII and UTF-16 strings from buffer.
      :param stack_buf: memory buffer
      :return: Extracted String namedtuples
      """
      for s in strings.extract_ascii_strings(stack_buf):
          yield s
      for s in strings.extract_unicode_strings(stack_buf):
          yield s
  
  
  def get_offset_written_at(ss_va, mem_writes):
      """
      Get instruction address where string was initially written.
      Removes false positives, e.g. when addresses got pushed that contain a part in ASCII range
      :param ss_va: identified string offset
      :param mem_writes: record of all memory writes
      :return: tuple: string offset, address where string was written
      """
      offset = 0
      written_at = None
      if ss_va in mem_writes:
          logging.debug("string in mem_writes at 0x%X", mem_writes[ss_va])
          written_at = mem_writes[ss_va]
      # one off ASCII
      elif (ss_va + 1) in mem_writes:
          logging.debug("string+1 in mem_writes at 0x%X", mem_writes[ss_va + 1])
          offset = 1
          written_at = mem_writes[ss_va + 1]
      # one off UTF-16LE
      elif (ss_va + 2) in mem_writes:
          logging.debug("string+2 in mem_writes at 0x%X", mem_writes[ss_va + 2])
          offset = 1  # still just one character off
          written_at = mem_writes[ss_va + 2]
      return offset, written_at
  
  
  def consolidate(intervals):
      """
      Consolidate interval ranges via https://codereview.stackexchange.com/a/69249.
      :param intervals: set of tuples with start and end address
      :return: list of consolidated intervals
      """
      sorted_by_lower_bound = sorted(intervals, key=lambda tup: tup[0])
      merged = []
      for higher in sorted_by_lower_bound:
          if not merged:
              merged.append(higher)
          else:
              lower = merged[-1]
              # test for intersection between lower and higher:
              # we know via sorting that lower[0] <= higher[0]
              if higher[0] <= lower[1]:
                  upper_bound = max(lower[1], higher[1])
                  merged[-1] = (lower[0], upper_bound)  # replace by merged interval
              else:
                  merged.append(higher)
      return merged
  
  
  def does_contain(stackstrings, ss):
      """
      Check existence of stackstring in list.
      :param stackstrings: list of all recovered stackstrings
      :param ss: new stackstring candidate
      :return: True if candidate already in stackstring list, False otherwise
      """
      hashable_ss = (ss.fva, ss.s, ss.written_at)
      for s in stackstrings:
          hashable = (s.fva, s.s, s.written_at)
          if hashable == hashable_ss:
              return True
      return False
  
  
  def does_contain_substr(stackstrings, new_ss):
      """
      Check if new string is a substring of an existing stackstring.
      :param stackstrings: list of all recovered stackstrings
      :param new_ss: new stackstring candidate
      :return: True if candidate substring of an existing stackstring, False otherwise
      """
      # IMPROVEMENT incorporate written_at in check
      for ss in filter(lambda es: es.fva == new_ss.fva, stackstrings):
          if new_ss.s in ss.s:
              return True
      return False
  
  
  def extend_existing(new_string, existing_strings):
      """
      Update an existing string if a new string extends it.
      :param new_string: new recovered string
      :param existing_strings: list of already recovered strings
      :return: True if extended existing string, False otherwise
      """
      for i, existing in enumerate(existing_strings):
          if existing.s in new_string.s:
              logging.debug("updating string %s to %s", existing.s, new_string.s)
              existing_strings[i] = new_string
              return True
      return False
  
  
  #
  # print and comment functions
  #
  def print_header():
      """
      Print formatted header
      :return: None
      """
      print("{:16}   {:16}   {:16}   {:16}".format("Function", "Written at", "Offset", "String"))
      print("{s:-<16}   {s:-<16}   {s:-<16}   {s:-<16}".format(s=""))
  
  
  def print_string(fva, written_at, offset, s):
      """
      Print formatted stackstring
      :param fva: function address containing stackstring
      :param written_at: instruction address where stackstring was written
      :param offset: memory offset (stack relative or global) where string was written
      :param s: recovered stackstring
      :return: None
      """
      if written_at:
          # IMPROVEMENT indicate if stack offset
          print("0x{:<16X} 0x{:<16X} 0x{:<16X} {:<}".format(fva, written_at, offset, s))
      else:
          print("0x{:<16X}   {:16} 0x{:<16X} {:<}".format(fva, "", offset, s))
  
  
  def print_summary(cnt_functions, cnt_analyzed, cnt_found, cnt_commented, errors):
      """
      Print analysis summary.
      :param cnt_functions: number of identified functions
      :param cnt_analyzed: number of successfully analyzed functions
      :param cnt_found: number of recovered stackstrings
      :param cnt_commented: number of commented stackstrings
      :param errors: list of encountered error message strings
      :return: None
      """
      print("\nironstrings summary\n{:-<24}".format(""))
      print("Ran successfully on {:d}/{:d} functions".format(cnt_analyzed, cnt_functions))
      print("Found {:d} stackstrings".format(cnt_found))
      print("Commented {:d} stackstrings".format(cnt_commented))
  
      if errors:
          print("Encountered {:d} errors".format(len(errors)))
          for err in errors:
              print(" - {}".format(err))
  
  
  def print_plain_summary(stackstrings):
      unique = get_unique_strings(stackstrings)
      print("\nRecovered {:d} unique stackstrings\n{:-<24}".format(len(unique), ""))
      for s in unique:
          print("{}".format(s))
  
  
  def get_unique_strings(stackstrings):
      """
      Get unique recovered stackstrings.
      :param stackstrings: list of all stackstrings
      :return: list of unique stackstrings
      """
      unique = set()
      for ss in stackstrings:
          unique.add(ss.s)
      return unique
  
  
  def format_comment(comment):
      """
      Wrap comment with prefix and suffix.
      :param comment: comment string
      :return: formatted comment string
      """
      return "{prefix}{cmt}{suffix}".format(prefix=COMMENT_STACKSTRING_PREFIX, cmt=comment,
                                            suffix=COMMENT_STACKSTRING_SUFFIX)
  
  
  def append_comment(va, new_cmt, repeatable=False):
      """
      Append a comment to an address in IDA Pro.
      :param va: comment address
      :param new_cmt: comment string
      :param repeatable: if True, append as repeatable comment
      :return: True if success
      """
      cmt = idc.get_cmt(va, repeatable)
      if not cmt:
          # no existing comment
          cmt = new_cmt
      else:
          if new_cmt in cmt:
              # comment already exists
              return True
          cmt = cmt + "\n" + new_cmt
      return idc.set_cmt(va, cmt, repeatable)
  
  
  if __name__ == "__main__":
      logging.basicConfig(level=logging.INFO)
      main()
  