# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# 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 mobius
import pymobius
import pymobius.vfs
import os.path
import datetime
import shutil
import traceback

ATTR_NAME = 'data.registry'

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief get folder's child by name
# @param folder folder
# @param name child's name
# @return child or None
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_child_by_name (folder, name):
  child = None

  try:
    child = folder.get_child_by_name (name, False)
  except Exception as e:
    mobius.core.logf ('DBG ' + str (e))

  return child

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief get folder's child by path
# @param folder folder
# @param path child's path
# @return child or None
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_child_by_path (folder, path):
  child = None

  try:
    child = folder.get_child_by_path (path, False)
  except Exception as e:
    mobius.core.logf ('DBG ' + str (e))

  return child

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief get folder's children
# @param folder folder
# @return children
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_children (folder):
  children = []

  try:
    children = folder.get_children ()
  except Exception as e:
    mobius.core.logf ('DBG ' + str (e))

  return children

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief check if entry is folder
# @param entry entry
# @return true/false
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def is_folder (entry):
  rc = False

  try:
    rc = entry.is_folder ()
  except Exception as e:
    mobius.core.logf ('DBG ' + str (e))

  return rc

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief forensics: win registry
# @author Eduardo Aguiar
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Ant (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self, item):
    self.__item = item
    self.__data = None

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief check if data is cached
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def is_data_cached (self):
    is_cached = False

    if self.__item.datasource:
      data = self.__data

      # get cached data, if any
      if not data:
        data = self.__item.get_attribute (ATTR_NAME)

      # check if data is up to date
      if data and data.mtime >= self.__item.datasource.mtime:
        is_cached = True

    return is_cached

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get registry data
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def get_data (self):
    mobius.core.logf ('DBG pymobius.registry.get_data (1)')

    # if item has datasource, retrieves data
    if self.__item.datasource:
      mobius.core.logf ('DBG pymobius.registry.get_data (2)')

      # try to retrieve cached data
      if not self.__data:
        mobius.core.logf ('DBG pymobius.registry.get_data (3)')
        self.__data = self.__item.get_attribute (ATTR_NAME)

      # no data yet, retrieves data from disk
      if not self.__data or self.__item.datasource.mtime > self.__data.mtime:
        self.retrieve_data ()

    # no datasource, removes old cached data
    else:
      self.remove_data ()
      self.__data = None
      return []

    # builds return list
    mobius.core.logf ('DBG pymobius.registry.get_data (5)')
    ret = []

    for data in self.__data.reg_list:
      mobius.core.logf ('DBG pymobius.registry.get_data (6)')
      registry = mobius.os.win.registry.registry ()

      for hivefile in data.hivefiles:
        try:
          localpath = self.__item.get_data_path (os.path.join ('registry', '%04d' % data.idx, hivefile.filename))
          mobius.core.logf ('DBG pymobius.registry.get_data (7) %s %s' % (hivefile.filename, localpath))
          registry.add_file_by_path (hivefile.role, hivefile.path, localpath)
        except Exception as e:
          mobius.core.logf ('WRN ' + str (e) + ' ' + traceback.format_exc ())

      ret.append (registry)

    mobius.core.logf ('DBG pymobius.registry.get_data (8)')
    return ret

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief retrieve artifacts
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def retrieve_data (self):
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (1)')
    self.__data = None

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

    # retrieve artifacts
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (2)')
    next_idx = 1
    registry = None
    hashes = set ()
    reg_list = []
    user_hivefiles = []

    # scan VFS root entries
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (3)')
    for root in vfs:
      mobius.core.logf ('DBG pymobius.registry.retrieve_data (4)')

      # system entries
      system_hivefiles = self.__get_system_hivefiles (root)
      mobius.core.logf ('DBG pymobius.registry.retrieve_data (5) %d' % len (system_hivefiles))

      if system_hivefiles:
        registry = pymobius.Data ()
        registry.idx = next_idx
        registry.hivefiles = system_hivefiles
        reg_list.append (registry)
        next_idx += 1

        if user_hivefiles:      # previously found NTUSER hivefiles
          registry.hivefiles.extend (user_hivefiles)
          user_hivefiles = []

      # user hivefiles
      mobius.core.logf ('DBG pymobius.registry.retrieve_data (6)')
      user_hivefiles = self.__get_user_hivefiles (root)
      mobius.core.logf ('DBG pymobius.registry.retrieve_data (7) %d' % len (user_hivefiles))

      if user_hivefiles and registry:
        registry.hivefiles.extend (user_hivefiles)
        user_hivefiles = []

    # retrieve hivefiles
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (8)')
    for registry in reg_list:
      mobius.core.logf ('DBG pymobius.registry.retrieve_data (8.1)')
      for hivefile in registry.hivefiles:
        mobius.core.logf ('DBG pymobius.registry.retrieve_data (8.2) %s %s' % (hivefile.filename, hivefile.role))
        hash_sha2_512 = self.__retrieve_file (hivefile.entry, registry.idx, hivefile.filename)
        mobius.core.logf ('DBG pymobius.registry.retrieve_data (8.3) %s' % hash_sha2_512)

        if hash_sha2_512:
          hashes.add (hash_sha2_512)
        del hivefile.entry

    # build data model
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (9)')
    self.__data = pymobius.Data ()
    self.__data.mtime = datetime.datetime.utcnow ()
    self.__data.reg_list = reg_list

    # save data model
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (10)')
    self.__item.set_attribute (ATTR_NAME, self.__data)

    # save KFF
    mobius.core.logf ('DBG pymobius.registry.retrieve_data (11)')
    if hashes:
      mobius.core.logf ('DBG pymobius.registry.retrieve_data (12)')
      case = self.__item.case
      path = case.create_path (os.path.join ('hashset', '%04d-registry.ignore' % self.__item.uid))

      fp = open (path, 'w')
      for h in hashes:
        fp.write ('%s\n' % h)
      fp.close ()

      mobius.core.logf ('DBG pymobius.registry.retrieve_data (13)')

    mobius.core.logf ('DBG pymobius.registry.retrieve_data (14)')

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief remove registry data for an item
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def remove_data (self):
    self.__item.remove_attribute (ATTR_NAME)

    # data/registry
    path = self.__item.get_data_path ('registry')
    if os.path.exists (path):
      shutil.rmtree (path)

    # .ignore
    case = self.__item.case
    path = case.get_path (os.path.join ('hashset', '%04d-registry.ignore' % self.__item.uid))

    if os.path.exists (path):
      os.remove (path)

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get system hivefiles
  # @param filesystem's root entry
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __get_system_hivefiles (self, root):
    mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (1) %s' %  root.path)
    entries = []

    # %windir%/system32/config
    config_dir = get_child_by_path (root, 'windows/system32/config')
    mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (2) %s' %  config_dir)

    if config_dir:
      components = get_child_by_name (config_dir, 'components')
      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (3) %s' % components)
      if components:
        entries.append (components)

      sam = get_child_by_name (config_dir, 'sam')
      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (4) %s' % sam)
      if sam:
        entries.append (sam)

      security = get_child_by_name (config_dir, 'security')
      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (5) %s' % security)
      if security:
        entries.append (security)

      software = get_child_by_name (config_dir, 'software')
      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (6) %s' % software)
      if software:
        entries.append (software)

      system = get_child_by_name (config_dir, 'system')
      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (7) %s' % system)
      if system:
        entries.append (system)

      default = get_child_by_name (config_dir, 'default')
      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (8) %s' % default)
      if default:
        entries.append (default)

      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (9)')

    # create hivefiles
    hivefiles = []

    for entry in entries:
      parent = entry.get_parent ()

      hivefile = pymobius.Data ()
      hivefile.path = entry.path.replace ('/', '\\')
      hivefile.role = entry.name.upper ()
      hivefile.filename = entry.name.upper ()
      hivefile.entry = entry
      hivefiles.append (hivefile)

      mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (10) %s %s %s' % (hivefile.filename, hivefile.role, hivefile.path))

    mobius.core.logf ('DBG pymobius.registry.__get_system_hivefiles (11)')
    return hivefiles

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get user hivefiles
  # @param filesystem's root entry
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __get_user_hivefiles (self, root):
    mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (1) %s' % root.path)
    entries = []

    # %profilesdir%/*/ntuser.dat
    users_dir = get_child_by_name (root, 'users')

    if not users_dir:
      users_dir = get_child_by_name (root, 'documents and settings')

    mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (2) %s' % users_dir)

    if users_dir:
      mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (2.1) %s' % users_dir.name)

      for child in get_children (users_dir):
        mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (2.1.1) %s' % child.name)
        if is_folder (child):
          mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (2.2) %s' % child.name)
          ntuser = get_child_by_name (child, 'ntuser.dat')
          if ntuser:
            entries.append ((child.name, ntuser))

    # %windir%/serviceprofiles/{localservice,networkservice}/ntuser.dat
    config_dir = get_child_by_path (root, 'windows/serviceprofiles')
    mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (3) %s' % config_dir)

    if config_dir:
      localservice = get_child_by_path (config_dir, 'localservice/ntuser.dat')
      mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (4) %s' % localservice)
      if localservice:
        entries.append ((config_dir.name, localservice))

      networkservice = get_child_by_path (config_dir, 'networkservice/ntuser.dat')
      mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (5) %s' % networkservice)
      if networkservice:
        entries.append ((config_dir.name, networkservice))

    # system profile
    config_dir = get_child_by_path (root, 'windows/system32/config/systemprofile')
    mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (6) %s' % config_dir)
    if config_dir:
      f = get_child_by_name (config_dir, "ntuser.dat")
      if f:
        entries.append ((config_dir.name, f))

    # create hivefiles
    mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (7)')
    hivefiles = []

    for username, entry in entries:
      hivefile = pymobius.Data ()
      hivefile.path = entry.path.replace ('/', '\\')
      hivefile.role = 'NTUSER'
      hivefile.filename = 'NTUSER-%s.dat' % username.upper ()
      hivefile.entry = entry
      hivefiles.append (hivefile)

    mobius.core.logf ('DBG pymobius.registry.__get_user_hivefiles (8)')
    return hivefiles

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief retrieve file
  # @param f file
  # @param path fullpath
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __retrieve_file (self, f, idx, filename):

    # create reader
    reader = f.new_reader ()

    if not reader:
      return None

    # create output file
    rpath = os.path.join ('registry', '%04d' % idx, filename)
    path = self.__item.create_data_path (rpath)

    fp = open (path, 'wb')
    h = mobius.crypt.hash ('sha2_512')
    BLOCK_SIZE = 65536

    # copy data
    data = reader.read (BLOCK_SIZE)
    while data:
      fp.write (data)
      h.update (data)
      data = reader.read (BLOCK_SIZE)

    fp.close ()

    # return hash
    return h.get_hex_digest ()
