# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# Mobius Forensic Toolkit
# Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019 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.xml
import pymobius.vfs
import os.path
import hashlib
import datetime
import shutil
import sys

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief generic dataholder
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class dataholder (object):
  pass

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @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, e:
    pymobius.LOG ('pymobius.registry.get_child_by_name', 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, e:
    pymobius.LOG ('pymobius.registry.get_child_by_path', 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, e:
    pymobius.LOG ('pymobius.registry.get_children', 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, e:
    pymobius.LOG ('pymobius.registry.is_folder', 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:
        path = self.__item.get_data_path ('registry/data.xml')

        if os.path.exists (path):
          pickle = pymobius.xml.Pickle ()
          data = pickle.load (path)

      # 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):
    pymobius.LOG ('pymobius.registry.get_data (1)')

    # if item has datasource, retrieves data
    if self.__item.datasource:
      pymobius.LOG ('pymobius.registry.get_data (2)')

      # try to retrieve cached data
      if not self.__data:
        pymobius.LOG ('pymobius.registry.get_data (3)')
        path = self.__item.get_data_path ('registry/data.xml')
        pymobius.LOG ('pymobius.registry.get_data (4)', path)

        if os.path.exists (path):
          pickle = pymobius.xml.Pickle ()
          self.__data = pickle.load (path)

      # 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
    pymobius.LOG ('pymobius.registry.get_data (5)')
    ret = []

    for data in self.__data.reg_list:
      pymobius.LOG ('pymobius.registry.get_data (6)')
      registry = mobius.registry.registry ()

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

      ret.append (registry)
            
    pymobius.LOG ('pymobius.registry.get_data (8)')
    return ret

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief retrieve artifacts
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def retrieve_data (self):
    pymobius.LOG ('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
    pymobius.LOG ('pymobius.registry.retrieve_data (2)')
    next_idx = 1
    registry = None
    hashes = set ()
    reg_list = []
    user_hivefiles = []

    # scan VFS root entries
    pymobius.LOG ('pymobius.registry.retrieve_data (3)')
    for root in vfs:
      pymobius.LOG ('pymobius.registry.retrieve_data (4)')

      # system entries
      system_hivefiles = self.__get_system_hivefiles (root)
      pymobius.LOG ('pymobius.registry.retrieve_data (5)', len (system_hivefiles))
      
      if system_hivefiles:
        registry = dataholder ()
        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
      pymobius.LOG ('pymobius.registry.retrieve_data (6)')
      user_hivefiles = self.__get_user_hivefiles (root)
      pymobius.LOG ('pymobius.registry.retrieve_data (7)', len (user_hivefiles))

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

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

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

    # build data model
    pymobius.LOG ('pymobius.registry.retrieve_data (9)')
    self.__data = dataholder ()
    self.__data.mtime = datetime.datetime.utcnow ()
    self.__data.reg_list = reg_list

    # save data model
    pymobius.LOG ('pymobius.registry.retrieve_data (10)')
    path = self.__item.create_data_path ('registry/data.xml')
    pickle = pymobius.xml.Pickle ()
    pickle.save (path, self.__data)

    # save KFF
    pymobius.LOG ('pymobius.registry.retrieve_data (11)')
    if hashes:
      pymobius.LOG ('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 ()

      pymobius.LOG ('pymobius.registry.retrieve_data (13)')

    pymobius.LOG ('pymobius.registry.retrieve_data (14)')

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief remove registry data for an item
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __remove_data (self):

    # 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):
    pymobius.LOG ('pymobius.registry.__get_system_hivefiles (1)', root.path)
    entries = []

    # %windir%/system32/config    
    config_dir = get_child_by_path (root, 'windows/system32/config')
    pymobius.LOG ('pymobius.registry.__get_system_hivefiles (2)', config_dir)

    if config_dir:
      components = get_child_by_name (config_dir, 'components')
      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (3)', components)
      if components:
        entries.append (components)

      sam = get_child_by_name (config_dir, 'sam')
      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (4)', sam)
      if sam:
        entries.append (sam)

      security = get_child_by_name (config_dir, 'security')
      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (5)', security)
      if security:
        entries.append (security)

      software = get_child_by_name (config_dir, 'software')
      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (6)', software)
      if software:
        entries.append (software)

      system = get_child_by_name (config_dir, 'system')
      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (7)', system)
      if system:
        entries.append (system)

      default = get_child_by_name (config_dir, 'default')
      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (8)', default)
      if default:
        entries.append (default)

      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (9)')

    # create hivefiles
    hivefiles = []

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

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

      pymobius.LOG ('pymobius.registry.__get_system_hivefiles (10)', hivefile.filename, hivefile.role, hivefile.path)

    pymobius.LOG ('pymobius.registry.__get_system_hivefiles (11)')
    return hivefiles

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief get user hivefiles
  # @param filesystem's root entry
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __get_user_hivefiles (self, root):
    pymobius.LOG ('pymobius.registry.__get_user_hivefiles (1)', 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')

    pymobius.LOG ('pymobius.registry.__get_user_hivefiles (2)', users_dir)

    if users_dir:
      pymobius.LOG ('pymobius.registry.__get_user_hivefiles (2.1)', users_dir.name)

      for child in get_children (users_dir):
        pymobius.LOG ('pymobius.registry.__get_user_hivefiles (2.1.1)', child.name)
        if is_folder (child):
          pymobius.LOG ('pymobius.registry.__get_user_hivefiles (2.2)', child.name)
          ntuser = get_child_by_name (child, 'ntuser.dat')
          if ntuser:
            entries.append (ntuser)

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

    if config_dir:
      localservice = get_child_by_path (config_dir, 'localservice/ntuser.dat')
      pymobius.LOG ('pymobius.registry.__get_user_hivefiles (4)', localservice)
      if localservice:
        entries.append (localservice)
                  
      networkservice = get_child_by_path (config_dir, 'networkservice/ntuser.dat')
      pymobius.LOG ('pymobius.registry.__get_user_hivefiles (5)', networkservice)
      if networkservice:
        entries.append (networkservice)

    # system profile
    config_dir = get_child_by_path (root, 'windows/system32/config/systemprofile')
    pymobius.LOG ('pymobius.registry.__get_user_hivefiles (6)', config_dir)
    if config_dir:
      f = get_child_by_name (config_dir, "ntuser.dat")
      if f:
        entries.append (f)

    # create hivefiles
    pymobius.LOG ('pymobius.registry.__get_user_hivefiles (7)')
    hivefiles = []

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

      hivefile = dataholder ()
      hivefile.path = entry.path.replace ('/', '\\')
      hivefile.role = 'NTUSER'
      hivefile.filename = 'NTUSER-%s.dat' % parent.name.upper ()
      hivefile.entry = entry
      hivefiles.append (hivefile)

    pymobius.LOG ('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, 'w')
    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 ()
