## ------------------------------------------------------------------------
## iSiLK - A python-based graphical user interface for SiLK
## ------------------------------------------------------------------------
## Copyright (c) 2007-2008 Carnegie Mellon University. All Rights Reserved.
## ------------------------------------------------------------------------
## Authors: Andrew Kompanek (ajk@cert.org)
## ------------------------------------------------------------------------
## GNU General Public License (GPL) Rights pursuant to Version 2, June 1991
## Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013
## ------------------------------------------------------------------------
## See the LICENSE file included in this distribution
## ------------------------------------------------------------------------

import pdb
import xml.dom.minidom as dom

class xmlPickler:
    
    def __init__(self, classDict=None):
        self._instance_stack = []
        self._classDict = classDict
        
    def _removeWhiteSpace(self, e):
        for c in e.childNodes:
            if c.nodeName == "#text":
                #print "Removing text node"
                e.removeChild(c)
            else:
                self._removeWhiteSpace(c)
        
    def parseString(self, xmlstr):
        x = dom.parseString(xmlstr)
        self._removeWhiteSpace(x.documentElement)
        return self.fromxml(x.documentElement)
        
    def toxml(self, obj, tag=None):
         #pdb.set_trace()
         if isinstance(obj, dict):
             return self.dicttoxml(obj, tag)
         elif type(obj) == list:
             return self.listtoxml(obj, tag)
         elif isinstance(obj, object):
             return self.instancetoxml(obj)
         raise "Can't convert to xml"

    def fromxml(self, e):
         if e.nodeName == "dict":
             return self.dictfromxml(e)
         elif e.nodeName == "list":
             return self.listfromxml(e)
         else:
             return self.instancefromxml(e)
             
    def _escape_string(self, val):
        return val
    
    # Fixme assumes attrs are always strings right now
    def instancetoxml(self, obj):
        attrs = obj._xml_children_()
        e = dom.Element("instance")
        e.setAttribute("type", obj.__class__.__name__)
        # Build a dictionary of the attrs to serialize
        d = {}
        for a in attrs:
            val = obj.__dict__[a]
            if type(val) == list or type(val) == dict:
                e.appendChild(self.toxml(val, a))
            elif type(val) in (str, int, bool, unicode):
                c = dom.Element(a)
                c.setAttribute("type", type(val).__name__)
                val_elem = dom.Element (type(val).__name__)
                c.appendChild(val_elem)
                val_elem.setAttribute("val", str(val))
                e.appendChild(c)
            elif isinstance(val, dict):
                e.appendChild(self.toxml(val, a))
            else:
                c = dom.Element(a)
                c.setAttribute("type", "instance")
                val_elem = self.instancetoxml(val)
                c.appendChild(val_elem)
                e.appendChild(c)               
                #if not val == None:
                #    c.appendChild(self.instancetoxml(val))
        return e
    
    # Initialize obj from xml element
    def instancefromxml(self, e):
        cls = self._classDict[e.getAttribute("type")]
        if len(self._instance_stack) > 0: 
            parent = self._instance_stack[len(self._instance_stack)-1]
        else:
            parent = None
        
        obj = cls(parent)
        self._instance_stack.append(obj)
        for c in e.childNodes:
            n = c.nodeName
            if n == "#text":
                continue
            t = c.getAttribute("type")
            if t == "list":
                obj.__dict__[n] = self.fromxml(c.childNodes[0])                
            elif t == "dict":
                val = self.fromxml(c.childNodes[0])
                if c.getAttribute("class") == "dict":
                    obj.__dict__[n] = val
                else:
                    subcls = self._classDict[c.getAttribute("class")]
                    obj.__dict__[n] = subcls(val)
            elif t=="bool":
                val_elem = c.childNodes[0]
                v = val_elem.getAttribute("val")
                obj.__dict__[n] = (v == "True")
            elif t=="str":
                val_elem = c.childNodes[0]
                v = val_elem.getAttribute("val")
                obj.__dict__[n] = v #eval("%s(\"%s\")" % (t,v))
            elif t=="unicode":
                val_elem = c.childNodes[0]
                v = val_elem.getAttribute("val")
                obj.__dict__[n] = v #eval("%s(\"%s\")" % (t,v))
            elif t=="int":
                val_elem = c.childNodes[0]
                v = val_elem.getAttribute("val")
                obj.__dict__[n] = eval("%s(\"%s\")" % (t,v))
            elif t=="instance":
                print c.toprettyxml()
                obj.__dict__[n] = self.fromxml(c.childNodes[0])
            else:
                raise "Unknown type referenced in xml: %s" % str(t)
        self._instance_stack.pop()
        obj._xml_postprocess_()
        return obj

    # assmes all values primitive types
    def listtoxml(self, alist, tag=None):
        if tag:
            root = dom.Element(tag)
            root.setAttribute("type", "list")
            info = dom.Element("list")
            root.appendChild(info)
        else:
            root = info = dom.Element("list")
        for val in alist:
            c = dom.Element("item")
            c.setAttribute("type", type(val).__name__)      
            if type(val) == list or type(val) == dict:
                c.appendChild(self.toxml(val, a))
            elif type(val) == str or type(val) == int or type(val) == unicode:
                c.setAttribute("val", str(val))
            else:
                c.appendChild(self.instancetoxml(val))
            info.appendChild(c)
        return root
    
    def listfromxml(self, e):
        l = []
        for e in e.childNodes:
           if e.nodeName == "#text":
               continue
           if not e.nodeName == "item":
              raise "Unexpected element in dictionary: expected <item>"
           n = str(e.getAttribute("name"))
           t = e.getAttribute("type")
           if t == "list" or t == "dict":
               for child in e.childNodes:
                   if child.nodeName != "#text":
                       break
               l.append(child)
#               l.append(self.fromxml(e.childNodes[0]))
           elif t=="str" or t=="int" or t=="unicode":
               v = e.getAttribute("val")
               if t=="str" or t=="unicode":
                   l.append(v)
               else:
                   l.append(eval("%s(\"%s\")" % (t,v)))           
           else:
               # Instantiate an empty instance... (better way?)
               l.append(self.fromxml(e.childNodes[0]))

        return l

    # fixme: assumes values are all lists or primitive types
    def dicttoxml(self, adict, tag=None):
        if tag:
            root = dom.Element(tag)
            root.setAttribute("type", "dict")
            if type(adict) != dict: # subclass
                root.setAttribute("class", type(adict).__name__)
            else:
                root.setAttribute("class", "dict")
            info = dom.Element("dict")
            root.appendChild(info)
        else:
            root = info = dom.Element("dict")
        for i in adict.keys():
            e = dom.Element("item")
            e.setAttribute("name", str(i))
            e.setAttribute("type", type(adict[i]).__name__)
            if type(adict[i]) == list or type(adict[i]) == dict:
                c = self.toxml(adict[i])
                e.appendChild(c)
            else:
                e.setAttribute("val",  str(adict[i]))
            info.appendChild(e)
        return root
    
    def dictfromxml(self, info):
        d = {}
        for e in info.childNodes:
            if e.nodeName == "#text":
                continue
            if not e.nodeName == "item":
                raise "Unexpected element in dictionary: expected <item>"
            n = str(e.getAttribute("name"))
            t = e.getAttribute("type")
            if t == "list" or t == "dict":
                # Skip text nodes
                for child in e.childNodes:
                    if child.nodeName != "#text":
                        break
                d[n] = self.fromxml(child)
            else:
                v = e.getAttribute("val")
                if t == "bool": 
                    d[n] = eval(v)
                else:
                    d[n] = eval("%s(\"%s\")" % (t,v))
                #print "*** Read dict item n=%s, t=%s, v=%s" % (n,t,v)
                #print "    %s" % d[n]
        return d
    
 