# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# Mobius Forensic Toolkit
# Copyright (C) 2008, 2009 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 libxml2
import mobius
XML_ENCODING='utf-8'

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

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize extension
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self):
    self.id = ''
    self.name = ''
    self.author = ''
    self.version = ''
    self.description = ''
    self.code = ''
    self.mediator = None
    self.icon_data = None
    self.services = []

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Service
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Service (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize service
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self):
    self.id = u''
    self.description = u''
    self.code = u''
    self.is_deprecated = False
    self.is_autostarted = True
    self.args = []
    self.use_varargs = False
    self.use_kwds = False

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Service argument
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Arg (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize arg
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self):
    self.id = u''
    self.description = u''
    self.defaultvalue = u''

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Persistence layer for extension
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Pickle (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Get node property with correct encoding
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __get_prop (self, node, name):
    value = node.prop (name)
    if value:
      value = value.decode (XML_ENCODING)
    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Set node property with correct encoding
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __set_prop (self, node, name, value):
    if value != None:
      node.setProp (name, value.encode (XML_ENCODING))

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load extension
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load (self, path):
    doc = libxml2.parseFile (path)
    node = doc.getRootElement ()

    extension = self.load_extension (node)
    return extension

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <extension>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_extension (self, node):
    extension = Extension ()
    extension.id = self.__get_prop (node, 'id')
    extension.name = self.__get_prop (node, 'name')
    extension.author = self.__get_prop (node, 'author')
    extension.version = self.__get_prop (node, 'version')
    extension.description = self.__get_prop (node, 'description')
    extension.services = []

    # load children
    node = node.children

    while node:
      if node.type == 'element' and node.name == 'service':
        extension.services.append (self.load_service (node))

      elif node.type == 'element' and node.name == 'icon':
        extension.icon_data = self.load_icon (node)

      elif node.type == 'element' and node.name == 'code':
        extension.code = self.load_code (node)

      node = node.next

    return extension

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <service>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_service (self, node):
    service = Service ()
    service.id = self.__get_prop (node, 'id')
    service.description = self.__get_prop (node, 'description')
    service.is_deprecated = self.__get_prop (node, 'is_deprecated') == 'true'
    service.is_autostarted = self.__get_prop (node, 'is_autostarted') == 'true'
    service.args = []
    service.use_varargs = self.__get_prop (node, 'use_varargs') == 'true'
    service.use_kwds = self.__get_prop (node, 'use_kwds') == 'true'

    # load children
    node = node.children

    while node:
      if node.type == 'element' and node.name == 'arg':
        service.args.append (self.load_arg (node))

      elif node.type == 'element' and node.name == 'code':
        service.code = self.load_code (node)

      node = node.next

    return service

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <arg>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_arg (self, node):
    arg = Arg ()
    arg.id = self.__get_prop (node, 'id')
    arg.defaultvalue = self.__get_prop (node, 'defaultvalue')
    arg.description = self.__get_prop (node, 'description')

    return arg

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <icon>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_icon (self, node):
    icon_data = node.getContent ()
    return icon_data

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <code>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_code (self, node):
    code = node.getContent ().strip ()
    return code

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save extension
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save (self, extension, path):
    doc = libxml2.newDoc ('1.0')
    node = self.save_extension (extension)
    doc.addChild (node)
    doc.saveFormatFileEnc (path, XML_ENCODING, 1)
    doc.freeDoc ()

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <extension>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_extension (self, extension):
    node = libxml2.newNode ('extension')
    self.__set_prop (node, 'id', extension.id)
    self.__set_prop (node, 'name', extension.name)
    self.__set_prop (node, 'author', extension.author)
    self.__set_prop (node, 'version', extension.version)
    self.__set_prop (node, 'description', extension.description)

    # services
    for service in extension.services:
      child = self.save_service (service)
      node.addChild (child)

    # icon
    if extension.icon_data:
      child = self.save_icon (extension.icon_data)
      node.addChild (child)

    # code
    if extension.code:
      child = self.save_code (extension.code)
      node.addChild (child)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <service>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_service (self, service):
    node = libxml2.newNode ('service')
    self.__set_prop (node, 'id', service.id)
    self.__set_prop (node, 'description', service.description)
    self.__set_prop (node, 'is_deprecated', 'true' if service.is_deprecated else 'false')
    self.__set_prop (node, 'is_autostarted', 'true' if service.is_autostarted else 'false')

    # args
    for arg in service.args:
      child = self.save_arg (arg)
      node.addChild (child)

    self.__set_prop (node, 'use_varargs', 'true' if service.use_varargs else 'false')
    self.__set_prop (node, 'use_kwds', 'true' if service.use_kwds else 'false')

    # code
    if service.code:
      child = self.save_code (service.code)
      node.addChild (child)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <arg>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_arg (self, arg):
    node = libxml2.newNode ('arg')
    self.__set_prop (node, 'id', arg.id)
    self.__set_prop (node, 'description', arg.description)
    self.__set_prop (node, 'defaultvalue', arg.defaultvalue)
    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <icon>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_icon (self, icon_data):
    node = libxml2.newNode ('icon')
    node.addContent (icon_data)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <code>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_code (self, code):
    node = libxml2.newNode ('code')
    node.addContent (code.strip ())

    return node

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.new
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_new ():
  extension = Extension ()
  return extension

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.open
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_open (path):
  pickle = Pickle ()
  extension = pickle.load (path)
  return extension

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.save
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_save (extension, path):
  pickle = Pickle ()
  pickle.save (extension, path)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.start-api
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_start_api (extension):
  extension.mediator = mobius.mediator.copy ()
  extension.module = {'i_mediator' : extension.mediator,
                      'i_extension_id' : extension.id,
		      'i_extension_name' : extension.name,
		      'i_extension_version' : extension.version,
		      'i_extension_description' : extension.description,
		      'i_extension_icon_data' : extension.icon_data}

  bytecode = compile (extension.code + '\n', '%s:main' % extension.id, 'exec')
  exec bytecode in extension.module

  # @begin-deprecated
  if 'Extension' in extension.module:
    extension_class = extension.module['Extension']()
    extension_class.id = extension.id
    extension_class.name = extension.name
    extension_class.version = extension.version
    extension_class.description = extension.description
    extension_class.icon_data = extension.icon_data
    extension_class.mediator = mobius.mediator.copy ()
    extension_class.path = extension.path
    extension.old_extension_class = extension_class

    start_api = getattr (extension_class, 'start_api', None)

    if start_api:
      start_api ()
  # @end-deprecated

  # register services
  prefix = 'svc_' + extension.id.replace ('-', '_') + '_'

  for service in extension.services:
    funcname = prefix + service.id.replace ('.', '_').replace ('-', '_')
    modulename = '%s:%s' % (extension.id, service.id)

    # compose code
    args = []
    for a in service.args:
      value = a.id
      if a.defaultvalue:
        value += '=%s' % a.defaultvalue
      args.append (value)

    if service.use_varargs:
      args.append ('*args')

    if service.use_kwds:
      args.append ('**kwds')
 
    code = 'def %s (%s):\n  ' % (funcname, ', '.join (args)) + service.code.replace ('\n', '\n  ') + '\n'

    # execute code
    bytecode = compile (code, modulename, 'exec')
    exec bytecode in extension.module

    if service.is_autostarted:
      extension.mediator.advertise (service.id, extension.module[funcname])

  return extension

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.stop-api
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_stop_api (extension):
  extension.mediator.clear ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.start
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_start (extension):
  extension_class = getattr (extension, 'old_extension_class', None)

  if extension_class:
    start = getattr (extension_class, 'start', None)
    if start:
      start ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.stop
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_stop (extension):
  extension_class = getattr (extension, 'old_extension_class', None)

  if extension_class:
    stop = getattr (extension_class, 'stop', None)
    if stop:
      stop ()

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.new-service
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_new_service ():
  service = Service ()
  return service

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief service: extension.new-service-arg
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
def svc_extension_new_service_arg ():
  arg = Arg ()
  return arg
