# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# Mobius Forensic Toolkit
# Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019,2020,2021,2022 Eduardo Aguiar
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
import pymobius.vfs
import pymobius.registry
import pymobius.registry.main
import pymobius.registry.installed_programs
import mobius
import os.path
import traceback

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Scan installed operating systems
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def scan (item):

  # get item VFS
  vfs = pymobius.vfs.get_item_vfs (item)
  if not vfs:
    raise Exception ("Datasource is not available")

  # scan VFS
  return [os_win (vfs, item)]

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Safe file wrapper (no exceptions)
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class file_wrapper (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self, f, item):
    self.__f = f
    self.__item = item

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get attributes from file
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __getattr__ (self, name):
    return self.__dict__.get (name, getattr (self.__f, name))

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Set file handled
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def set_handled (self):
    if not self.__f:
      return

    reader = self.__f.new_reader ()
    if not reader:
      return

    # calculate hash sha2-512
    h = mobius.crypt.hash ('sha2_512')
    data = reader.read (65536)

    while data:
      h.update (data)
      data = reader.read (65536)

    # add to kff file
    case = self.__item.case
    path = case.create_path (os.path.join ('hashset', '%04d-handled.ignore' % self.__item.uid))

    fp = open (path, 'a')
    fp.write ('%s\n' % h.get_hex_digest ())
    fp.close ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Safe folder wrapper (no exceptions)
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class folder_wrapper (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self, folder, item, case_sensitive=True):
    self.__folder = folder
    self.__item = item
    self.__case_sensitive = case_sensitive

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get attributes from folder
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __getattr__ (self, name):
    return self.__dict__.get (name, getattr (self.__folder, name))

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get children
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_children (self):
    children = []

    try:
      for child in self.__folder.get_children ():
        if child.is_folder ():
          f = folder_wrapper (child, self.__item, self.__case_sensitive)
        else:
          f = file_wrapper (child, self.__item)
        children.append (f)

    except Exception as e:
      mobius.core.logf ('WRN (path=%s): %s' % (self.__folder.path, e))

    return children

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get child by name
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_child_by_name (self, name, case_sensitive=False):
    child = None

    try:
      child = self.__folder.get_child_by_name (name, self.__case_sensitive)

      if not child:
        pass
      elif child.is_folder ():
        child = folder_wrapper (child, self.__item, self.__case_sensitive)
      else:
        child = file_wrapper (child, self.__item)

    except Exception as e:
      mobius.core.logf ('WRN (name=%s): %s' % (self.__folder.path, e))

    return child

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get child by path
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_child_by_path (self, path, case_sensitive=False):
    child = None

    try:
      child = self.__folder.get_child_by_path (path, self.__case_sensitive)

      if not child:
        pass

      elif child.is_folder ():
        child = folder_wrapper (child, self.__item, self.__case_sensitive)

      else:
        child = file_wrapper (child, self.__item)

    except Exception as e:
      mobius.core.logf ('WRN %s %s' % (str (e), traceback.format_exc ()))

    return child

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief OS Windows object
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class os_win (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self, vfs, item):
    self.__vfs = vfs
    self.__installed_programs = []
    self.__installed_programs_loaded = False
    self.__item = item

    # retrieve registry object
    ant = pymobius.registry.main.Ant (item)
    data = ant.get_data ()

    if not data:
      raise Exception ('OS.Win has no registry')

    elif len (data) == 1:
      self.__registry = data[0]

    else:
      mobius.core.logf ('INF Item with %d registries' % len (data))
      self.__registry = data[0]

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get item
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_item (self):
    return self.__item

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get registry object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_registry (self):
    return self.__registry

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get system folder path
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_folder (self, name):
    return pymobius.registry.get_system_folder (self.__registry, name)

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get System root folders
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_root_folders (self):
    return [ folder_wrapper (root, self.__item, False) for root in self.__vfs ]

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get installed programs
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_installed_programs (self):
    if not self.__installed_programs_loaded:
      self.__installed_programs = pymobius.registry.installed_programs.get (self.__registry)
      self.__installed_programs_loaded = True

    return self.__installed_programs

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get system profiles
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_profiles (self):

    # get profiles path
    profiles_path = self.get_folder ('PROFILES') or 'C:\\Users'
    profiles_path = profiles_path[3:].replace ('\\', '/')

    # retrieve profiles
    for root in self.__vfs:
      root_wrapper = folder_wrapper (root, self.__item, False)
      profile_folder = root_wrapper.get_child_by_path (profiles_path)

      if profile_folder:
        for child in profile_folder.get_children ():
          if not child.is_reallocated () and child.is_folder ():
            yield os_win_profile (self, child)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief OS Windows profile object
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class os_win_profile (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self, operating_system, folder):
    self.__folder = folder
    self.__operating_system = operating_system
    self.__ntuser_dat = None
    self.folder = folder
    self.path = folder.path.replace ('/', '\\')
    self.username = folder.path.rsplit ('/', 1)[-1]
    self.is_deleted = folder.is_deleted ()

    self.__control_folders = {
        'appdata' : 'AppData/Roaming',
        'localappdata' : 'AppData/Local',
        'cookies' : 'AppData/Roaming/Microsoft/Windows/Cookies',
        'desktop' : 'Desktop',
        'downloads' : 'Downloads',
        'favorites' : 'Favorites',
    }

    # get NTUSER.DAT root key
    nt_path = folder.path + '/NTUSER.DAT'
    nt_path = nt_path.replace ('/', '\\').lower ()

    registry = operating_system.get_registry ()

    for f in registry.files:
      if f.role == 'NTUSER' and f.path.lower () == nt_path:
        self.__ntuser_dat = f.root_key

    # get user shell folders
    if self.__ntuser_dat:
      for value in self.__ntuser_dat.get_value_by_mask ('Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\User Shell Folders\\*'):
        path = value.data.get_data_as_string ('utf-16')
        if path.startswith ('%USERPROFILE%\\'):
          folder_id = value.name.lower ().replace (' ', '')
          self.__control_folders[folder_id] = path[14:].replace ('\\', '/')

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get filesystem entry by path
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_entry_by_path (self, path):
    rpath = path

    if path.startswith ('%'):
      pos = path.find ('%', 1)

      if pos != -1:
        folder_id = path[1:pos]
        folder_path = self.__control_folders.get (folder_id)

        if folder_path:
          rpath = folder_path + path[pos+1:]

    try:
      f = self.__folder.get_child_by_path (rpath)
    except Exception as e:
      mobius.core.logf ('WRN ' + str (e) + ' ' + traceback.format_exc ())
      f = None

    return f

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get AppData folder
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_appdata_folder (self, rpath):
    appdata = self.__control_folders.get ('appdata', 'AppData/Roaming')

    if rpath:
      appdata = appdata + '/' + rpath

    folder = None

    try:
      folder = self.__folder.get_child_by_path (appdata)
    except Exception as e:
      mobius.core.logf ('WRN ' + str (e) + ' ' + traceback.format_exc ())
      folder = None

    return folder

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get AppData entry
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_appdata_entry (self, rpath):
    appdata = self.__control_folders.get ('appdata', 'AppData/Roaming')

    if rpath:
      appdata = appdata + '/' + rpath

    entry = None

    try:
      entry = self.__folder.get_child_by_path (appdata)
    except Exception as e:
      mobius.core.logf ('WRN ' + str (e) + ' ' + traceback.format_exc ())

    return entry

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get AppDataLocal folder
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_appdata_local_folder (self, rpath):
    appdata = self.__control_folders.get ('local appdata', 'AppData/Local')

    if rpath:
      appdata = appdata + '/' + rpath

    folder = None

    try:
      folder = self.__folder.get_child_by_path (appdata)
    except Exception as e:
      mobius.core.logf ('WRN ' + str (e) + ' ' + traceback.format_exc ())
      folder = None

    return folder

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get control folder
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_control_folder (self, folder_id, folder_default):
    rpath = self.__control_folders.get (folder_id, folder_default)
    folder = None

    try:
      folder = self.__folder.get_child_by_path (rpath)
    except Exception as e:
      mobius.core.logf ('WRN ' + str (e) + ' ' + traceback.format_exc ())
      folder = None

    return folder
