text: # Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
  
  import os
  import io
  import ntpath
  import hashlib
  import fnmatch
  import speakeasy.winenv.defs.windows.windows as windefs
  from speakeasy.errors import FileSystemEmuError
  
  
  def normalize_response_path(path):
      def _get_speakeasy_root():
          return os.path.join(os.path.dirname(__file__), os.pardir)
  
      root_var = '$ROOT$'
      if root_var in path:
          root = _get_speakeasy_root()
          return path.replace(root_var, root)
      return path
  
  
  class MapView(object):
      """
      Represents a shared memory view
      """
      def __init__(self, base, offset, size, protect, process=None):
          self.base = base
          self.offset = offset
          self.size = size
          self.protect = protect
          self.process = process
  
  
  class FileMap(object):
      """
      Represents a memory mapped file
      """
      curr_handle = 0x280
  
      def __init__(self, name, size, prot, backed_file=None):
          self.name = name
          self.backed_file = backed_file
          self.views = {}
          self.size = size
  
      def get_handle(self):
          hmap = FileMap.curr_handle
          FileMap.curr_handle += 4
          return hmap
  
      def get_name(self):
          return self.name
  
      def get_backed_file(self):
          return self.backed_file
  
      def add_view(self, base, offset, size, protect):
          view = MapView(base, offset, size, protect)
          self.views.update({base: view})
  
  
  class File(object):
      """
      Base class for an emulated file
      """
      curr_handle = 0x80
  
      def __init__(self, path, config={}, data=b''):
          self.path = path
          self.data = None
          self.bytes_written = 0
          if data:
              self.data = io.BytesIO(data)
          self.curr_offset = 0
          self.is_dir = False
          self.config = config
  
      def get_handle(self):
          hfile = File.curr_handle
          File.curr_handle += 4
          return hfile
  
      def get_path(self):
          return self.path
  
      def get_hash(self):
          h = hashlib.sha256()
          data = self.get_data()
          h.update(data)
          return h.hexdigest()
  
      def get_size(self):
          if not self.data and self.config:
              self.data = self.handle_file_data()
          if not self.data:
              return 0
          off = self.data.tell()
          self.data.seek(0, io.SEEK_SET)
          size = len(self.data.read())
          self.data.seek(off, io.SEEK_SET)
          return size
  
      def get_data(self, size=-1, reset_pointer=False):
          if not self.data and self.config:
              self.data = self.handle_file_data()
  
          if not self.data:
              return b''
  
          off = self.data.tell()
          if off == self.get_size():
              if reset_pointer:
                  # Reset the file pointer
                  self.data.seek(0)
              else:
                  return None
  
          return self.data.read(size)
  
      def seek(self, offset, whence):
          if whence not in [io.SEEK_CUR, io.SEEK_SET, io.SEEK_END]:
              return
          if self.data:
              self.data.seek(offset, whence)
  
      def add_data(self, data):
  
          if not self.data:
              self.data = io.BytesIO()
          off = self.data.tell()
          self.data.seek(0, io.SEEK_END)
          self.data.write(data)
          self.data.seek(off, io.SEEK_SET)
          self.bytes_written += len(data)
  
      def remove_data(self):
          self.data = io.BytesIO(b'')
  
      def is_directory(self):
          return self.is_dir
  
      def handle_file_data(self):
          """
          Based on the emulation config, determine what data
          to return from the read request
          """
  
          path = self.config.get('path')
          if path:
              path = normalize_response_path(path)
              with open(path, 'rb') as f:
                  return io.BytesIO(f.read())
          bf = self.config.get('byte_fill')
          if bf:
              byte = bf.get('byte')
              if byte.startswith('0x'):
                  byte = 0xFF & int(byte, 0)
              else:
                  byte = 0xFF & int(byte, 16)
              size = bf.get('size')
              b = (byte).to_bytes(1, 'little')
              return b * size
          return io.BytesIO(b'')
  
  
  class Pipe(File):
      """
      Emulated named pipe objects
      """
      curr_handle = 0x400
  
      def __init__(self, name, mode, num_instances, out_size, in_size, config={}):
          super(Pipe, self).__init__(path=name, config=config)
          self.name = name
          self.mode = mode
          self.num_instances = num_instances
          self.out_size = out_size
          self.in_size = in_size
  
      def get_handle(self):
          hpipe = Pipe.curr_handle
          Pipe.curr_handle += 4
          return hpipe
  
  
  class FileManager(object):
      """
      Manages file system activity during emulation
      """
      def __init__(self, config=None):
          super(FileManager, self).__init__()
          self.file_handles = {}
          self.pipe_handles = {}
          self.file_maps = {}
          self.config = config
          self.files = []
  
      def file_create_mapping(self, hfile, name, size, prot):
          if hfile not in (windefs.INVALID_HANDLE_VALUE, 0):
              f = self.get_file_from_handle(hfile)
              fm = FileMap(name, size, prot, f)
              hnd = fm.get_handle()
              self.file_maps.update({hnd: fm})
              return hnd
          else:
              fm = FileMap(name, size, prot, None)
              hnd = fm.get_handle()
              self.file_maps.update({hnd: fm})
              return hnd
  
      def walk_files(self):
          for f in self.config.get('files', []):
              path = f.get('emu_path')
              if not path:
                  continue
              yield path
  
      def get_dropped_files(self):
          return [f for f in self.files if f.bytes_written == f.get_size()]
  
      def get_mapping_from_handle(self, handle):
          return self.file_maps.get(handle)
  
      def get_mapping_from_addr(self, addr):
          for h, fmap in self.file_maps.items():
              for base, view in fmap.views.items():
                  if base == addr:
                      return fmap
  
      def get_file_from_handle(self, handle):
          return self.file_handles.get(handle)
  
      def get_pipe_from_handle(self, handle):
          return self.pipe_handles.get(handle)
  
      def get_file_from_path(self, path):
          for f in self.files:
              if f.get_path().lower() == path.lower():
                  return f
          return None
  
      def get_all_files(self):
          return self.files
  
      def handle_file_data(self, fconf):
  
          path = fconf.get('path')
          if path:
              path = normalize_response_path(path)
              with open(path, 'rb') as f:
                  return f.read()
          bf = fconf.get('byte_fill')
          if bf:
              byte = bf.get('byte')
              if byte.startswith('0x'):
                  byte = 0xFF & int(byte, 0)
              else:
                  byte = 0xFF & int(byte, 16)
              size = bf.get('size')
              b = (byte).to_bytes(1, 'little')
              return b * size
  
      def add_existing_file(self, path, data):
          """
          Register an existing file already included in the emulation space
          (with data)
          """
          f = File(path, data=data)
          self.files.append(f)
          return f
  
      def create_file(self, path):
          f = self.get_file_from_path(path)
          if f:
              self.files.remove(f)
          f = File(path)
          self.files.append(f)
          return f
  
      def delete_file(self, path):
          f = self.get_file_from_path(path)
          if f:
              self.files.remove(f)
              return True
          return False
  
      def get_emu_file(self, path):
          # Does this file exist in our emulation environment
          # See if we have a handler for this exact file
          for f in self.config.get('files', []):
              mode = f.get('mode')
              if mode == 'full_path':
                  if fnmatch.fnmatch(path.lower(), f.get('emu_path').lower()):
                      return f
  
          # If no full path handler exists, do we have an extension handler?
          for f in self.config.get('files', []):
              path_ext = ntpath.splitext(path)[-1:][0].strip('.')
              if path_ext:
                  mode = f.get('mode')
                  if mode == 'by_ext':
                      if path_ext.lower() == f.get('ext'):
                          return f
  
          # Finally, do we have a catch-all default handler?
          for f in self.config.get('files', []):
  
              mode = f.get('mode')
              if mode == 'default':
                  return f
          return None
  
      def pipe_open(self, path, mode, num_instances, out_size, in_size):
          hnd = None
          fconf = self.get_emu_file(path)
          if not fconf:
              return hnd
          p = Pipe(path, mode, num_instances, out_size, in_size, config=fconf)
          hnd = p.get_handle()
          self.pipe_handles.update({hnd: p})
          return hnd
  
      def does_file_exist(self, path):
          if self.get_file_from_path(path):
              return True
  
          if self.get_emu_file(path):
              return True
          return False
  
      def get_object_from_handle(self, handle):
          obj = self.file_maps.get(handle)
          if obj:
              return obj
          obj = self.pipe_handles.get(handle)
          if obj:
              return obj
          obj = self.file_handles.get(handle)
          if obj:
              return obj
  
      def file_open(self, path, create=False, truncate=False, is_dir=False):
          hnd = None
  
          if create:
              f = self.create_file(path)
              hnd = f.get_handle()
              self.file_handles.update({hnd: f})
          else:
  
              f = self.get_file_from_path(path)
              if f:
                  hnd = f.get_handle()
                  self.file_handles.update({hnd: f})
                  return hnd
  
              fconf = self.get_emu_file(path)
              if not fconf:
                  return hnd
  
              real_path = fconf.get('path', '')
              real_path = normalize_response_path(real_path)
              if not truncate:
                  if real_path and not os.path.exists(real_path):
                      raise FileSystemEmuError('File path not found: %s' % (real_path))
                  f = File(path, config=fconf)
                  self.files.append(f)
              else:
                  if real_path and not os.path.exists(real_path):
                      raise FileSystemEmuError('File path not found: %s' % (real_path))
                  f = File(path, config=fconf)
                  self.files.append(f)
              hnd = f.get_handle()
              self.file_handles.update({hnd: f})
  
          return hnd
  