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

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief registry constants
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
STARTTIME = None
DEBUG = False

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief print log message
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def LOG (*args):
  if not DEBUG:
    return

  now = datetime.datetime.utcnow ()

  global STARTTIME
  if not STARTTIME:
    STARTTIME = now

  d = now - STARTTIME

  print >> sys.stderr, 'mobius.forensics.registry.main', now, '%d.%06d' % (d.seconds, d.microseconds), ' '.join (str (a) for a in args)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @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:
    LOG ('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:
    LOG ('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:
    LOG ('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:
    LOG ('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):

    # if item has datasource, retrieves data
    if self.__item.datasource:

      # try to retrieve cached data
      if not self.__data:
        path = self.__item.get_data_path ('registry/data.xml')

        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
    ret = []

    for data in self.__data.reg_list:
      registry = mobius.forensics.registry.registry ()

      for hivefile in data.hivefiles:
        localpath = self.__item.get_data_path (os.path.join ('registry', '%04d' %   data.idx, hivefile.filename))
        registry.add_file_by_path (hivefile.role, hivefile.path, localpath)

      ret.append (registry)
            
    return ret

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief retrieve artifacts
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def retrieve_data (self):
    self.__data = None

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

    # retrieve artifacts
    next_idx = 1
    registry = None
    hashes = set ()
    reg_list = []
    user_hivefiles = []

    # scan VFS root entries
    for root in vfs:
            
      # system entries
      system_hivefiles = self.__get_system_hivefiles (root)
      
      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
      user_hivefiles = self.__get_user_hivefiles (root)

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

    # retrieve hivefiles
    for registry in reg_list:
      for hivefile in registry.hivefiles:
        hash_sha2_512 = self.__retrieve_file (hivefile.entry, registry.idx, hivefile.filename)
        hashes.add (hash_sha2_512)
        del hivefile.entry

    # build data model
    self.__data = dataholder ()
    self.__data.mtime = datetime.datetime.utcnow ()
    self.__data.reg_list = reg_list

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

    # save KFF
    if hashes:
      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 ()

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

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

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

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

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

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

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

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

      LOG ('__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)

    return hivefiles

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

    LOG ('__get_user_hivefiles (2)', users_dir)

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

      for child in get_children (users_dir):
        LOG ('__get_user_hivefiles (2.1.1)', child.name)
        if is_folder (child):
          LOG ('__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')
    LOG ('__get_user_hivefiles (3)', config_dir)

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

    # system profile
    config_dir = get_child_by_path (root, 'windows/system32/config/systemprofile')
    LOG ('__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
    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)

    return hivefiles

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief retrieve file
  # @param entry entry
  # @param path fullpath
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __retrieve_file (self, entry, idx, filename):
    rpath = os.path.join ('registry', '%04d' % idx, filename)
    path = self.__item.create_data_path (rpath)
    reader = entry.new_reader ()

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

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

    fp.close ()
    
    return h.get_hex_digest ()
