# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# Mobius Forensic Toolkit
# Copyright (C) 2008,2009,2010,2011,2012,2013,2014 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 os
import os.path
import shutil
import zipfile
import sys
import libxml2
import mobius

BASEDIR = os.path.join (os.path.expanduser ("~"), '.mobiusft', 'extensions')

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Extension metadata
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Metadata (object):

  def __init__ (self):
    self.id = ''
    self.name = ''
    self.author = ''
    self.version = ''
    self.description = ''

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Get extension path
# @param extension_id Extension ID
# @param *args sub-paths
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_path (extension_id, *args):
  extension_path = os.path.normpath (os.path.join (BASEDIR, extension_id))

  if not extension_path.startswith (BASEDIR):
    raise Exception, 'Warning: bad crafted extension ID <%s>' % extension_id

  path = os.path.join (extension_path, *args)

  if not path.startswith (extension_path):
    raise Exception, 'Warning: bad crafted path. Extension ID <%s>, Path <%s>' % (extension_id, '/'.join (args))

  return path

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Get extension metadata
# @param extension_id Extension ID
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_metadata (extension_id):
  path = get_path (extension_id, "metadata.xml")

  if os.path.exists (path):
    doc = libxml2.parseFile (path)
    node = doc.getRootElement ()

    metadata = Metadata ()
    metadata.id = (node.prop ('id') or '').decode ('utf-8')
    metadata.name = (node.prop ('name') or '').decode ('utf-8')
    metadata.author = (node.prop ('author') or '').decode ('utf-8')
    metadata.version = (node.prop ('version') or '').decode ('utf-8')
    metadata.description = (node.prop ('description') or '').decode ('utf-8')
    doc.freeDoc ()

  else:
    metadata = None

  return metadata

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Get metadata from extension file
# @param path Extension path
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_metadata_from_file (path):

  if os.path.exists (path):
    zfile = zipfile.ZipFile (path)
    metadata = zfile.read ('metadata.xml')
    zfile.close ()

    doc = libxml2.parseDoc (metadata)
    node = doc.getRootElement ()

    metadata = Metadata ()
    metadata.id = (node.prop ('id') or '').decode ('utf-8')
    metadata.name = (node.prop ('name') or '').decode ('utf-8')
    metadata.author = (node.prop ('author') or '').decode ('utf-8')
    metadata.version = (node.prop ('version') or '').decode ('utf-8')
    metadata.description = (node.prop ('description') or '').decode ('utf-8')
    doc.freeDoc ()

  else:
    metadata = None

  return metadata

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Get icon data from extension file
# @param path Extension path
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def get_icon_data_from_file (path):
  icon_data = None

  if os.path.exists (path):
    zfile = zipfile.ZipFile (path)
    icon_data = zfile.read ('icon.png')
    zfile.close ()

  return icon_data

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Install extension
# @param path Extension path
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def install (path):
  zfile = zipfile.ZipFile (path)

  # read metadata
  metadata = zfile.read ('metadata.xml')

  doc = libxml2.parseDoc (metadata)
  node = doc.getRootElement ()
  extension_id = (node.prop ('id') or '').decode ('utf-8')
  doc.freeDoc ()

  # set directory names
  extension_dir = get_path (extension_id)	# extension directory
  install_dir = extension_dir + '.install'	# installation directory
  old_dir = extension_dir + '.old'		# old version directory

  # remove previous .install and .old directory, if any
  if os.path.exists (install_dir):
    shutil.rmtree (install_dir)

  if os.path.exists (old_dir):
    shutil.rmtree (old_dir)
  
  # extract files to .install directory
  zfile.extractall (install_dir)
  zfile.close ()

  # run setup.py if it exists
  setup_py = os.path.join (install_dir, 'build', 'setup.py')

  if os.path.exists (setup_py):
    oldpwd = os.getcwd ()
    os.chdir (install_dir)

    fp = open (setup_py)
    exec fp in {}
    fp.close ()

    os.chdir (oldpwd)

  # rename extension dir to .old
  if os.path.exists (extension_dir):
    shutil.move (extension_dir, old_dir)

  # rename .install directory to extension dir
  shutil.move (install_dir, extension_dir)

  # remove old version, if any
  if os.path.exists (old_dir):
    shutil.rmtree (old_dir)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Uninstall extension
# @param extension_id Extension ID
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def uninstall (extension_id):
  path = get_path (extension_id)
  shutil.rmtree (path)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief List installed extensions
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def installed ():

  for path in os.listdir (BASEDIR):
    fullpath = os.path.join (BASEDIR, path)
    ext = os.path.splitext (fullpath)[1]

    if os.path.isdir (fullpath) and ext != '.install':
      yield path

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Data
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
g_extensions = {}	# running extensions

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.get-metadata
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_get_metadata (extension_id):
  return get_metadata (extension_id)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.get-metadata-from-file
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_get_metadata_from_file (path):
  return get_metadata_from_file (path)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.get-icon-data-from-file
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_get_icon_data_from_file (path):
  return get_icon_data_from_file (path)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.get-icon-path
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_get_icon_path (extension_id):
  return get_path (extension_id, 'icon.png')

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.get-resource-path
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_get_resource_path (extension_id, *args):
  return get_path (extension_id, 'resource', *args)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.list
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_list ():
  return installed ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.is-running
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_is_running (extension_id):
  return extension_id in g_extensions



# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Global variable container
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class GlobalData (object):
  pass

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.start-all
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_start_all ():

  # start API
  for extension_id in installed ():
    try:
      mobius.mediator.call ('extension.start-api', extension_id)
    except Exception, e:
      print 'Error starting extension "%s" API: %s' % (extension_id, e)

  # start extensions
  for extension_id in g_extensions.iterkeys ():
    try:
      mobius.mediator.call ('extension.start', extension_id)
    except Exception, e:
      print 'Error starting extension "%s": %s' % (extension_id, e)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.stop-all
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_stop_all ():

  # stop extensions
  for extension_id in g_extensions.keys ():
    mobius.mediator.call ('extension.stop', extension_id)

  # stop extensions API
  for extension_id in g_extensions.keys ():
    mobius.mediator.call ('extension.stop-api', extension_id)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.start-api
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_start_api (extension_id):

  # instance data
  gdata = GlobalData ()
  gdata.mediator = mobius.mediator.copy ()
  context = {'gdata' : gdata}

  # execute code
  extension_dir = get_path (extension_id)
  sys.path.insert (0, extension_dir)

  try:
    pyfile = os.path.join (extension_dir, "build", "main.py")

    fp = open (pyfile)
    sourcecode = fp.read ()
    fp.close ()

    bytecode = compile (sourcecode, '%s:main' % extension_id, 'exec')
    exec bytecode in context

    start_api = context.get ('pvt_start_api')
    start_api ()
    g_extensions[extension_id] = context

  finally:

    # clean up local imports
    del sys.path[0]

    for name in sys.modules.keys ():
      if name.startswith ('extension'):
        del sys.modules[name]

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.stop-api
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_stop_api (extension_id):
  context = g_extensions.pop (extension_id)
  mediator = context.get ('gdata').mediator
  mediator.clear ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.start
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_start (extension_id):
  context = g_extensions.get (extension_id)
  pvt_start = context.get ('pvt_start')

  if pvt_start:
    pvt_start ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.stop
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_stop (extension_id):
  context = g_extensions.get (extension_id)
  pvt_stop = context.get ('pvt_stop')

  if pvt_stop:
    pvt_stop ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.install
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_install (path):
  install (path)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.uninstall
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_uninstall (extension_id):
  uninstall (extension_id)
