#!/usr/bin/env python
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# Mobius Forensic Toolkit
# Copyright (C) 2008,2009,2010 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 optparse
import sys
import unittest

XML_ENCODING='utf-8'

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Placeholder class, just for refactoring
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class DummyRefactoringClass (object):
  pass

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @deprecated version=0.4.7
# @brief Persistence layer for case
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
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 application data
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load (self, path):
    doc = libxml2.parseFile (path)
    node = doc.getRootElement ()

    case = self.load_case (node)
    case.icon_data = ICON_CASE_DATA
    case.filename = path
    case.is_new = False
    case.is_modified = False
    case.uid = 1

    model = DummyRefactoringClass ()
    model.next_uid = int (case.attributes.pop ('next_uid')) + 1
    model.objects = case.objects
    model.objects[1] = case

    for obj in model.objects.itervalues ():
      obj.model = model
      
      for rel, b in obj.rel_to:
        obj_b = model.objects.get (b)
        obj_b.rel_from.append ((rel, obj.uid))
 
    return model

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load case from node
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_case (self, node):
    case = DummyRefactoringClass ()
    case.attributes = {}
    case.attributes['next_uid'] = self.__get_prop (node, 'next_uid') or '1'
    case.attributes['rootdir'] = self.__get_prop (node, 'rootdir')
    case.attributes['category'] = 'case'
    case.rel_to = []
    case.rel_from = []
    case.objects = {}

    # load children
    node = node.children

    while node:
      if node.type == 'element' and node.name == 'item':
        child = self.load_item (node, case)
        case.objects[child.uid] = child
        case.rel_to.append (('composite', child.uid))

      elif node.type == 'element' and node.name == 'attribute':
        name, value = self.load_attribute (node)
        case.attributes[name] = value

      node = node.next

    return case

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load item from node
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_item (self, node, case):
    item = DummyRefactoringClass ()
    item.uid = int (self.__get_prop (node, 'uid')) + 1
    item.attributes = {'category' : self.__get_prop (node, 'category')}
    item.rel_to = []
    item.rel_from = []

    # load children
    node = node.children

    while node:
      if node.type == 'element' and node.name == 'item':
        child = self.load_item (node, case)
        case.objects[child.uid] = child
        item.rel_to.append (('composite', child.uid))

      elif node.type == 'element' and node.name == 'attribute':
        id, value = self.load_attribute (node)
        item.attributes[id] = value

      node = node.next

    return item

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load attribute from node
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_attribute (self, node):
    name = self.__get_prop (node, 'name')
    value = self.__get_prop (node, 'value')
    return name, value

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @deprecated version=0.5.3
# @brief Persistence layer for model
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Pickle2 (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 Load model from XML file
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load (self, path):
    doc = libxml2.parseFile (path)
    node = doc.getRootElement ()
    model = self.load_model (node)

    return model

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <model>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_model (self, node):
    model = Model ()
    model.next_uid = int (self.__get_prop (node, 'next_uid') or '1')

    # load children
    node = node.children

    while node:
      if node.type == 'element' and node.name == 'object':
        child = self.load_object (node)
        child.model = model
        model.objects[child.uid] = child

      elif node.type == 'element' and node.name == 'relation':
        rel, a, b = self.load_relation (node)
        obj_a = model.objects.get (a)
        obj_a.rel_to.append ((rel, b))

        obj_b = model.objects.get (b)
        obj_b.rel_from.append ((rel, a))

      node = node.next

    return model

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <object>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_object (self, node):
    obj = Object ()
    obj.uid = int (self.__get_prop (node, 'uid'))
    obj.rel_to = []
    obj.rel_from = []

    # load children
    node = node.children

    while node:
      if node.type == 'element' and node.name == 'attribute':
        name, value = self.load_attribute (node)
        obj.attributes[name] = value
        
      node = node.next

    return obj

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <relation>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_relation (self, node):
    rel = self.__get_prop (node, 'id')
    a = int (self.__get_prop (node, 'a'))
    b = int (self.__get_prop (node, 'b'))
    return rel, a, b

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <attribute>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_attribute (self, node):
    name = self.__get_prop (node, 'name')
    value = self.__get_prop (node, 'value')
    return name, value

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Mobius model
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Model (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self):
    self.next_uid = 1
    self.objects = {}


# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Mobius object
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Object (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Initialize object
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def __init__ (self):
    self.uid = -1
    self.attributes = {}
    self.rel_to = []
    self.rel_from = []

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Generic object's instance
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class Instance (object):
  pass

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Persistence layer
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class XMLPickle (object):

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load XML file
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load (self, path):
    self.objects = {}

    doc = libxml2.parseFile (path)
    node = doc.getRootElement ()
    name, value = self.load_item (node)

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load any item
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_item (self, node):
    name = node.prop ('name')
    value = None

    # basic types
    if node.type == 'element' and node.name == 'attribute':
      value, datatype = self.load_attribute (node)

    elif node.type == 'element' and node.name == 'string':
      value = self.load_string (node)

    elif node.type == 'element' and node.name == 'unicode':
      value = self.load_unicode (node)

    # container types
    elif node.type == 'element' and node.name == 'reference':
      value = self.load_reference (node)

    elif node.type == 'element' and node.name == 'tuple':
      value = self.load_tuple (node)

    elif node.type == 'element' and node.name == 'list':
      value = self.load_list (node)

    elif node.type == 'element' and node.name == 'set':
      value = self.load_set (node)

    elif node.type == 'element' and node.name == 'dict':
      value = self.load_dict (node)

    elif node.type == 'element' and node.name == 'object':
      value = self.load_object (node)

    return name, value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <attribute>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_attribute (self, node):
    value = node.prop ('value')
    datatype = node.prop ('datatype')

    if datatype == 'int':
      value = int (value)

    elif datatype == 'bool':
      value = (value == "true")

    elif datatype == 'float':
      value = float (value)

    elif datatype == 'unicode':
      value = value.decode (XML_ENCODING)

    elif datatype == 'NoneType':
      value = None

    return value, datatype

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <string>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_string (self, node):
    value = node.getContent ()

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <unicode>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_unicode (self, node):
    value = unicode (node.getContent (), XML_ENCODING)

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <reference>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_reference (self, node):
    object_id = node.prop ('id')
    return self.objects.get (object_id)

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <tuple>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_tuple (self, node):
    object_id = node.prop ('id')
    value = []
    self.objects[object_id] = value

    node = node.children

    while node:
      if node.type == 'element':
        item_name, item_value = self.load_item (node)
        value.append (item_value)
      node = node.next

    value = tuple (value)
    self.objects[object_id] = value

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <list>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_list (self, node):
    object_id = node.prop ('id')
    value = []
    self.objects[object_id] = value

    node = node.children

    while node:
      if node.type == 'element':
        item_name, item_value = self.load_item (node)
        value.append (item_value)
      node = node.next

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <set>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_set (self, node):
    object_id = node.prop ('id')
    value = set ()
    self.objects[object_id] = value

    node = node.children

    while node:
      if node.type == 'element':
        item_name, item_value = self.load_item (node)
        value.add (item_value)
      node = node.next

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <dict>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_dict (self, node):
    object_id = node.prop ('id')
    value = {}
    self.objects[object_id] = value

    node = node.children

    while node:
      if node.type == 'element':
        # @begin-deprecated version=0.5.3
        name = node.prop ('name')
        if name:
          item_key, item_value = self.load_item (node)
        else:
        # @end-deprecated
          item_key, item_value = self.load_tuple (node)

        value[item_key] = item_value
      node = node.next

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Load <object>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def load_object (self, node):
    object_id = node.prop ('id')
    value = Instance ()
    self.objects[object_id] = value

    node = node.children

    while node:
      item_name, item_value = self.load_item (node)
      if item_name:
        setattr (value, item_name, item_value)
      node = node.next

    return value

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save XML file
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save (self, path, value):
    self.objects = {}

    doc = libxml2.newDoc ('1.0')
    node = self.save_item (value)
    doc.addChild (node)
    doc.saveFormatFileEnc (path, XML_ENCODING, 1)
    doc.freeDoc ()

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save any item
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_item (self, value):
    datatype = type (value).__name__

    # basic types 
    if datatype == 'int':
      node = self.save_attribute (value, datatype)

    elif datatype == 'bool':
      node = self.save_attribute (value, datatype)

    elif datatype == 'float':
      node = self.save_attribute (value, datatype)

    elif datatype == 'NoneType':
      node = self.save_attribute (value, datatype)

    elif datatype == 'str':
      node = self.save_string (value)

    elif datatype == 'unicode':
      node = self.save_unicode (value)

    # container types
    else:
      node = self.save_container (value)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <attribute>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_attribute (self, value, datatype):
    node = libxml2.newNode ('attribute')

    # datatype property
    if datatype != 'str':
      node.setProp ('datatype', datatype)

    # value property
    if datatype == 'unicode':
      node.setProp ('value', value.encode (XML_ENCODING))

    elif datatype == 'bool':
      if value:
        node.setProp ('value', 'true')

    elif datatype == 'str':
      node.setProp ('value', value)

    elif datatype != 'NoneType':
      node.setProp ('value', str (value))

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <string>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_string (self, value):
    if '\n' in value:
      node = libxml2.newNode ('string')
      node.addContent (value)
    else:
      node = self.save_attribute (value, 'str')

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <unicode>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_unicode (self, value):
    if '\n' in value:
      node = libxml2.newNode ('unicode')
      node.addContent (value.encode (XML_ENCODING))
    else:
      node = self.save_attribute (value, 'unicode')

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save container objects (dict, set, list, tuple, object)
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_container (self, value):
    object_id = id (value)

    if object_id in self.objects:
      node = libxml2.newNode ('reference')
      node.setProp ('id', str (object_id))

    else:
      self.objects[object_id] = value
      datatype = type (value).__name__

      if datatype == 'tuple':
        node = self.save_tuple (value)

      elif datatype == 'list':
        node = self.save_list (value)

      elif datatype == 'set':
        node = self.save_set (value)

      elif datatype == 'dict':
        node = self.save_dict (value)

      else:
        node = self.save_object (value)

      node.setProp ('id', str (object_id))

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <tuple>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_tuple (self, value):
    node = libxml2.newNode ('tuple')

    for item in value:
      child = self.save_item (item)
      node.addChild (child)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <list>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_list (self, value):
    node = libxml2.newNode ('list')

    for item in value:
      child = self.save_item (item)
      node.addChild (child)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <set>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_set (self, value):
    node = libxml2.newNode ('set')

    for item in value:
      child = self.save_item (item)
      node.addChild (child)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <dict>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_dict (self, value):
    node = libxml2.newNode ('dict')

    for key, value in value.iteritems ():
      child = self.save_tuple ((key, value))
      node.addChild (child)

    return node

  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  # @brief Save <object>
  # =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  def save_object (self, value):
    node = libxml2.newNode ('object')

    for name, value in vars (value).iteritems ():
      child = self.save_item (value)
      child.setProp ('name', name)
      node.addChild (child)

    return node

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# @brief Unittests for XML.Pickle
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
class TestXMLPickle (unittest.TestCase):

  def setUp (self):
    class Struct (object): pass

    model = Struct ()
    model.name = 'testing'
    model.uname = u'my testing \u0123'
    model.utext = u'''abc \u0123
test123 \u0122'''
    model.owner = Struct ()
    model.owner.name = 'John Doe'
    model.owner.data = [1, 5, 3, {'abc':5}]

    model.int_a = 23
    model.int_b = 0
    model.int_c = -5
    model.bool_true = True
    model.bool_false = False
    model.float_a = 0.5
    model.float_b = 0.0
    model.float_c = 10000000.5
    model.none = None
    model.tuple_a = (1,3,5)
    model.tuple_b = ('yellow', 5, 'blue', 3, 0.5)
    model.tuple_c = ()
    model.tuple_d = (1,)
    model.list_a = [1, -2, 'abc']
    model.list_b = []
    model.set_a = set ()
    model.set_b = set ([55, -1, 0.5, 'abc'])
    model.dict_a = {'a' : 5, 'b' : 3, 'c' : -1}
    model.dict_b = {1 : None, 'a' : 18, u'test\u0121' : 0.5}
    model.dict_c = {}
    model.composite_a = [(1,5), set ([3,'abc']), {1 : None, 'a' : 0.5}, []]
    model.composite_b = ([0, 0, 'j'], set (['a', 'b', None]), {}, {})
    model.composite_c = set ([0, -1, 1, (1, 5), ('a', 'b', 'c'), 'abc', u'abc'])
    model.composite_d = {1: {'a' : 5}, 2: [1, 5], 3: ('ab', 'cd'), 4: set ()}

    # circular reference
    model.owner.model = model
    model.reference_a = {'l' : []}
    model.reference_a['l'].append (model.reference_a)
    model.reference_b = ([1, 5], 5, 'abc', 'd')
    model.reference_b[0].append (model.reference_b)

    # cross reference
    model.obj_a = Struct ()
    model.obj_a.name = 'sibling 1'
    model.obj_b = Struct ()
    model.obj_b.name = 'sibling 2'
    model.obj_a.sibling = model.obj_b
    model.obj_b.sibling = model.obj_a

    xmlpickle = XMLPickle ()
    xmlpickle.save ('unittest.xml', model)

  def testnone (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.none).__name__, 'NoneType')
    self.assertEqual (model.none, None)

  def testint (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.int_a).__name__, 'int')
    self.assertEqual (type (model.int_b).__name__, 'int')
    self.assertEqual (type (model.int_c).__name__, 'int')
    self.assertEqual (model.int_a, 23)
    self.assertEqual (model.int_b, 0)
    self.assertEqual (model.int_c, -5)

  def testbool (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.bool_true).__name__, 'bool')
    self.assertEqual (type (model.bool_false).__name__, 'bool')
    self.assertEqual (model.bool_true, True)
    self.assertFalse (model.bool_false)

  def testfloat (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.float_a).__name__, 'float')
    self.assertEqual (type (model.float_b).__name__, 'float')
    self.assertEqual (type (model.float_c).__name__, 'float')
    self.assertAlmostEqual (model.float_a, 0.5)
    self.assertAlmostEqual (model.float_b, 0.0)
    self.assertAlmostEqual (model.float_c, 10000000.5)

  def teststring (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.name).__name__, 'str')
    self.assertEqual (model.name, 'testing')

  def testunicode (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.uname).__name__, 'unicode')
    self.assertEqual (type (model.utext).__name__, 'unicode')
    self.assertEqual (model.uname, u'my testing \u0123')
    self.assertEqual (model.utext, u'''abc \u0123
test123 \u0122''')

  def testobject (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (model.owner.name, 'John Doe')
    self.assertEqual (model.owner.data, [1, 5, 3, {'abc':5}])

  def testtuple (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.tuple_a).__name__, 'tuple')
    self.assertEqual (type (model.tuple_b).__name__, 'tuple')
    self.assertEqual (type (model.tuple_c).__name__, 'tuple')
    self.assertEqual (type (model.tuple_d).__name__, 'tuple')
    self.assertEqual (len (model.tuple_a), 3)
    self.assertEqual (len (model.tuple_b), 5)
    self.assertEqual (len (model.tuple_c), 0)
    self.assertEqual (len (model.tuple_d), 1)
    self.assertEqual (model.tuple_a, (1,3,5))
    self.assertEqual (model.tuple_b, ('yellow', 5, 'blue', 3, 0.5))
    self.assertEqual (model.tuple_c, ())
    self.assertEqual (model.tuple_d, (1,))

  def testlist (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.list_a).__name__, 'list')
    self.assertEqual (type (model.list_b).__name__, 'list')
    self.assertEqual (len (model.list_a), 3)
    self.assertEqual (len (model.list_b), 0)
    self.assertEqual (model.list_a, [1, -2, 'abc'])
    self.assertEqual (model.list_b, [])

  def testset (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.set_a).__name__, 'set')
    self.assertEqual (type (model.set_b).__name__, 'set')
    self.assertEqual (len (model.set_a), 0)
    self.assertEqual (len (model.set_b), 4)
    self.assertEqual (model.set_a, set ())
    self.assertEqual (model.set_b, set ([55, -1, 0.5, 'abc']))

  def testdict (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.dict_a).__name__, 'dict')
    self.assertEqual (type (model.dict_b).__name__, 'dict')
    self.assertEqual (type (model.dict_c).__name__, 'dict')
    self.assertEqual (len (model.dict_a), 3)
    self.assertEqual (len (model.dict_b), 3)
    self.assertEqual (len (model.dict_c), 0)
    self.assertEqual (model.dict_a, {'a' : 5, 'b' : 3, 'c' : -1})
    self.assertEqual (model.dict_b, {1 : None, 'a' : 18, u'test\u0121' : 0.5})
    self.assertEqual (model.dict_c, {})

  def testcomposite (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (type (model.composite_a).__name__, 'list')
    self.assertEqual (type (model.composite_b).__name__, 'tuple')
    self.assertEqual (type (model.composite_c).__name__, 'set')
    self.assertEqual (type (model.composite_d).__name__, 'dict')
    self.assertEqual (len (model.composite_a), 4)
    self.assertEqual (len (model.composite_b), 4)
    self.assertEqual (len (model.composite_c), 6)
    self.assertEqual (len (model.composite_d), 4)
    self.assertEqual (model.composite_a, [(1,5), set ([3,'abc']), {1 : None, 'a' : 0.5}, []])
    self.assertEqual (model.composite_b, ([0, 0, 'j'], set (['a', 'b', None]), {}, {}))
    self.assertEqual (model.composite_c, set ([0, -1, 1, (1, 5), ('a', 'b', 'c'), 'abc', u'abc']))
    self.assertEqual (model.composite_d, {1: {'a' : 5}, 2: [1, 5], 3: ('ab', 'cd'), 4: set ()})

  def testreference (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (model.owner.model.name, 'testing')

    self.assertEqual (type (model.reference_a).__name__, 'dict')
    self.assertEqual (len (model.reference_a), 1)
    self.assertEqual (id (model.reference_a), id (model.reference_a['l'][0]))

    self.assertEqual (type (model.reference_b).__name__, 'tuple')
    self.assertEqual (len (model.reference_b), 4)
    self.assertEqual (len (model.reference_b[0]), 3)

    # tuples cannot have elements referencing the tuple
    # self.assertEqual (type (model.reference_b[0][2]).__name__, 'tuple')
    # self.assertEqual (id (model.reference_b), id (model.reference_b[0][2]))

  def testcrossreference (self):
    xmlpickle = XMLPickle ()
    model = xmlpickle.load ('unittest.xml')

    self.assertEqual (model.obj_a.name, 'sibling 1')
    self.assertEqual (model.obj_b.name, 'sibling 2')
    self.assertEqual (model.obj_a.sibling.name, 'sibling 2')
    self.assertEqual (model.obj_b.sibling.name, 'sibling 1')


# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# show message
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
print 'Mobius Forensic Toolkit - Case update v1.0'
print 'by Eduardo Aguiar'
print

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# read command line
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
parser = optparse.OptionParser ()
parser.add_option ("-u", "--unittest", action="store_true", dest="unittest", default=False, help="Run unittest and exits")
parser.add_option ("-o", "--output", dest="output", help="Output filename")
(option, args) = parser.parse_args ()

if option.output and len (args) > 1:
  print '"-o" cannot be used with more than one file'
  sys.exit (2)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# run unittest
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
if option.unittest:
  suite1 = unittest.makeSuite (TestXMLPickle)
  suite = unittest.TestSuite ((suite1,))
  unittest.TextTestRunner (verbosity=2).run (suite)
  sys.exit (0)

# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
# update files
# =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
xmlpickle = XMLPickle ()

for path in args:
  data = open (path).read ()

  if '<case' in data: # version < 0.4.7
    pickle = Pickle ()
    print '>> "%s". version prior to 0.4.7 detected' % path

  elif '<model' in data: # version < 0.5.3
    pickle = Pickle2 ()
    print '>> "%s". version prior to 0.5.3 detected' % path

  else:
    print '>> "%s" is already up to date. Skipped' % path
    continue

  model = pickle.load (path)
  output = option.output or path
  print '>> writing "%s"' % output
  xmlpickle.save (output, model)
