added add_element() method
[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.config import Config  
8 from sfa.util.faults import SfaNotImplemented, InvalidRSpec
9
10 class RSpec:
11     header = '<?xml version="1.0"?>\n'
12     template = """<RSpec></RSpec>"""
13     namespaces = {}
14     config = Config()
15     xml = None
16   
17     def __init__(self, rspec="", namespaces={}):
18         if rspec:
19             self.parse_rspec(rspec, namespaces)
20         else:
21             self.create()
22
23     def create(self, type="advertisement"):
24         """
25         Create root element
26         """
27         # eg. 2011-03-23T19:53:28Z 
28         date_format = '%Y-%m-%dT%H:%M:%SZ'
29         now = datetime.utcnow()
30         generated_ts = now.strftime(date_format)
31         expires_ts = (now + timedelta(hours=1)).strftime(date_format) 
32         self.parse_rspec(self.template, self.namespaces)
33         self.xml.set('valid_until', expires_ts)
34         self.xml.set('generated', generated_ts)
35     
36     def parse_rspec(self, rspec, namespaces={}):
37         """
38         parse rspec into etree
39         """
40         parser = etree.XMLParser(remove_blank_text=True)
41         try:
42             tree = etree.parse(rspec, parser)
43         except IOError:
44             # 'rspec' file doesnt exist. 'rspec' is proably an xml string
45             try:
46                 tree = etree.parse(StringIO(rspec), parser)
47             except:
48                 raise InvalidRSpec('Must specify a xml file or xml string. Received: ' + rspec )
49
50         self.xml = tree.getroot()  
51         if namespaces:
52            self.namespaces = namespaces
53
54     def add_attribute(self, elem, name, value):
55         """
56         Add attribute to specified etree element    
57         """
58         opt = etree.SubElement(elem, name)
59         opt.text = value
60
61     def add_element(self, name, attrs={}, parent=None, text=""):
62         """
63         Generic wrapper around etree.SubElement(). Adds an element to 
64         specified parent node. Adds element to root node is parent is 
65         not specified. 
66         """
67         if parent == None:
68             parent = self.xml
69         element = etree.SubElement(parent, name)
70         if text:
71             element.text = text
72         if isinstance(attrs, dict):
73             for attr in attrs:
74                 element.set(attr, attrs[attr])  
75         return element
76
77     def remove_attribute(self, elem, name, value):
78         """
79         Removes an attribute from an element
80         """
81         if elem is not None:
82             opts = elem.iterfind(name)
83             if opts is not None:
84                 for opt in opts:
85                     if opt.text == value:
86                         elem.remove(opt)
87
88
89     def validate(self, schema):
90         """
91         Validate against rng schema
92         """
93         
94         relaxng_doc = etree.parse(schema)
95         relaxng = etree.RelaxNG(relaxng_doc)
96         if not relaxng(self.xml):
97             error = relaxng.error_log.last_error
98             message = "%s (line %s)" % (error.message, error.line)
99             raise InvalidRSpec(message)
100         return True
101         
102
103     def __str__(self):
104         return self.toxml()
105
106     def toxml(self):
107         return self.header + etree.tostring(self.xml, pretty_print=True)  
108         
109     def save(self, filename):
110         f = open(filename, 'w')
111         f.write(self.toxml())
112         f.close()
113  
114 if __name__ == '__main__':
115     rspec = RSpec()
116     print rspec