no need to import sfa.util.config.Config here
[sfa.git] / sfa / rspecs / rspec.py
1 #!/usr/bin/python 
2 from lxml import etree
3 from StringIO import StringIO
4 from datetime import datetime, timedelta
5 from sfa.util.xrn import *
6 from sfa.util.plxrn import hostname_to_urn
7 from sfa.util.faults import SfaNotImplemented, InvalidRSpec
8
9 class RSpec:
10     header = '<?xml version="1.0"?>\n'
11     template = """<RSpec></RSpec>"""
12     xml = None
13     type = None
14     version = None
15     namespaces = None
16     user_options = {}
17   
18     def __init__(self, rspec="", namespaces={}, type=None, user_options={}):
19         self.type = type
20         self.user_options = user_options
21         if rspec:
22             self.parse_rspec(rspec, namespaces)
23         else:
24             self.create()
25
26     def create(self):
27         """
28         Create root element
29         """
30         # eg. 2011-03-23T19:53:28Z 
31         date_format = '%Y-%m-%dT%H:%M:%SZ'
32         now = datetime.utcnow()
33         generated_ts = now.strftime(date_format)
34         expires_ts = (now + timedelta(hours=1)).strftime(date_format) 
35         self.parse_rspec(self.template, self.namespaces)
36         self.xml.set('expires', expires_ts)
37         self.xml.set('generated', generated_ts)
38     
39     def parse_rspec(self, rspec, namespaces={}):
40         """
41         parse rspec into etree
42         """
43         parser = etree.XMLParser(remove_blank_text=True)
44         try:
45             tree = etree.parse(rspec, parser)
46         except IOError:
47             # 'rspec' file doesnt exist. 'rspec' is proably an xml string
48             try:
49                 tree = etree.parse(StringIO(rspec), parser)
50             except:
51                 raise InvalidRSpec('Must specify a xml file or xml string. Received: ' + rspec )
52         self.xml = tree.getroot()  
53         if namespaces:
54            self.namespaces = namespaces
55
56     def xpath(self, xpath):
57         return this.xml.xpath(xpath, namespaces=self.namespaces)
58
59     def add_attribute(self, elem, name, value):
60         """
61         Add attribute to specified etree element    
62         """
63         opt = etree.SubElement(elem, name)
64         opt.text = value
65
66     def add_element(self, name, attrs={}, parent=None, text=""):
67         """
68         Generic wrapper around etree.SubElement(). Adds an element to 
69         specified parent node. Adds element to root node is parent is 
70         not specified. 
71         """
72         if parent == None:
73             parent = self.xml
74         element = etree.SubElement(parent, name)
75         if text:
76             element.text = text
77         if isinstance(attrs, dict):
78             for attr in attrs:
79                 element.set(attr, attrs[attr])  
80         return element
81
82     def remove_attribute(self, elem, name, value):
83         """
84         Removes an attribute from an element
85         """
86         if elem is not None:
87             opts = elem.iterfind(name)
88             if opts is not None:
89                 for opt in opts:
90                     if opt.text == value:
91                         elem.remove(opt)
92
93     def remove_element(self, element_name, root_node = None):
94         """
95         Removes all occurences of an element from the tree. Start at 
96         specified root_node if specified, otherwise start at tree's root.   
97         """
98         if not root_node:
99             root_node = self.xml
100
101         if not element_name.startswith('//'):
102             element_name = '//' + element_name
103
104         elements = root_node.xpath('%s ' % element_name, namespaces=self.namespaces)
105         for element in elements:
106             parent = element.getparent()
107             parent.remove(element)
108          
109
110     def merge(self, in_rspec):
111         pass
112
113     def validate(self, schema):
114         """
115         Validate against rng schema
116         """
117         
118         relaxng_doc = etree.parse(schema)
119         relaxng = etree.RelaxNG(relaxng_doc)
120         if not relaxng(self.xml):
121             error = relaxng.error_log.last_error
122             message = "%s (line %s)" % (error.message, error.line)
123             raise InvalidRSpec(message)
124         return True
125         
126     def cleanup(self):
127         """
128         Optional method which inheriting classes can choose to implent. 
129         """
130         pass 
131
132     def __process_slivers(self, slivers):
133         """
134         Creates a dict of sliver details for each sliver host
135         
136         @param slivers a single hostname, list of hostanmes or list of dicts keys on hostname,
137         Returns a list of dicts 
138         """
139         if not isinstance(slivers, list):
140             slivers = [slivers]
141         dicts = []
142         for sliver in slivers:
143             if isinstance(sliver, dict):
144                 dicts.append(sliver)
145             elif isinstance(sliver, basestring):
146                 dicts.append({'hostname': sliver}) 
147         return dicts
148
149     def __str__(self):
150         return self.toxml()
151
152     def toxml(self, cleanup=False):
153         if cleanup:
154             self.cleanup()
155         return self.header + etree.tostring(self.xml, pretty_print=True)  
156         
157     def save(self, filename):
158         f = open(filename, 'w')
159         f.write(self.toxml())
160         f.close()
161  
162 if __name__ == '__main__':
163     rspec = RSpec()
164     print rspec