7ab94ba276a404de805c95accce08d0f98b7ba40
[sfa.git] / sfa / dummy / dummyslices.py
1 import time
2 from types import StringTypes
3 from collections import defaultdict
4
5 from sfa.util.sfatime import utcparse, datetime_to_epoch
6 from sfa.util.sfalogging import logger
7 from sfa.util.xrn import Xrn, get_leaf, get_authority, urn_to_hrn
8
9 from sfa.rspecs.rspec import RSpec
10 from sfa.storage.model import SliverAllocation
11
12 from sfa.dummy.dummyxrn import DummyXrn, hrn_to_dummy_slicename
13
14 MAXINT =  2L**31-1
15
16 class DummySlices:
17
18     def __init__(self, driver):
19         self.driver = driver
20
21     def get_slivers(self, xrn, node=None):
22         hrn, type = urn_to_hrn(xrn)
23          
24         slice_name = hrn_to_dummy_slicename(hrn)
25         
26         slices = self.driver.shell.GetSlices({'slice_name': slice_name})
27         slice = slices[0]
28         # Build up list of users and slice attributes
29         user_ids = slice['user_ids']
30         # Get user information
31         all_users_list = self.driver.shell.GetUsers({'user_id':user_ids})
32         all_users = {}
33         for user in all_users_list:
34             all_users[user['user_id']] = user        
35
36         # Build up list of keys
37         all_keys = set()
38         for user in all_users_list:
39             all_keys.extend(user['keys'])
40
41         slivers = []
42         for slice in slices:
43             keys = all_keys
44             # XXX Sanity check; though technically this should be a system invariant
45             # checked with an assertion
46             if slice['expires'] > MAXINT:  slice['expires']= MAXINT
47             
48             slivers.append({
49                 'hrn': hrn,
50                 'name': slice['name'],
51                 'slice_id': slice['slice_id'],
52                 'expires': slice['expires'],
53                 'keys': keys,
54             })
55
56         return slivers
57  
58
59     def verify_slice_nodes(self, slice_urn, slice, rspec_nodes):
60
61         slivers = {}
62         for node in rspec_nodes:
63             hostname = node.get('component_name')
64             client_id = node.get('client_id')
65             component_id = node.get('component_id').strip()
66             if hostname:
67                 hostname = hostname.strip()
68             elif component_id:
69                 hostname = xrn_to_hostname(component_id)
70             if hostname:
71                 slivers[hostname] = {'client_id': client_id, 'component_id': component_id}
72         all_nodes = self.driver.shell.GetNodes()
73         requested_slivers = []
74         for node in all_nodes:
75             if node['hostname'] in slivers.keys():
76                 requested_slivers.append(node['node_id'])
77
78         if 'node_ids' not in slice.keys():
79             slice['node_ids']=[] 
80         nodes = self.driver.shell.GetNodes({'node_ids': slice['node_ids']})
81         current_slivers = [node['node_id'] for node in nodes]
82
83         # remove nodes not in rspec
84         deleted_nodes = list(set(current_slivers).difference(requested_slivers))
85
86         # add nodes from rspec
87         added_nodes = list(set(requested_slivers).difference(current_slivers))        
88
89         try:
90             self.driver.shell.AddSliceToNodes({'slice_id': slice['slice_id'], 'node_ids': added_nodes})
91             self.driver.shell.DeleteSliceFromNodes({'slice_id': slice['slice_id'], 'node_ids': deleted_nodes})
92
93         except: 
94             logger.log_exc('Failed to add/remove slice from nodes')
95
96         slices = self.driver.shell.GetSlices({'slice_name': slice['slice_name']})
97         resulting_nodes = self.driver.shell.GetNodes({'node_ids': slices[0]['node_ids']})
98
99         # update sliver allocations
100         for node in resulting_nodes:
101             client_id = slivers[node['hostname']]['client_id']
102             component_id = slivers[node['hostname']]['component_id']
103             sliver_hrn = '%s.%s-%s' % (self.driver.hrn, slice['slice_id'], node['node_id'])
104             sliver_id = Xrn(sliver_hrn, type='sliver').urn
105             record = SliverAllocation(sliver_id=sliver_id, client_id=client_id,
106                                       component_id=component_id,
107                                       slice_urn = slice_urn,
108                                       allocation_state='geni_allocated')
109             record.sync(self.driver.api.dbsession())
110         return resulting_nodes
111         
112
113     def verify_slice(self, slice_hrn, slice_record, expiration, options=None):
114         if options is None: options={}
115         slicename = hrn_to_dummy_slicename(slice_hrn)
116         parts = slicename.split("_")
117         login_base = parts[0]
118         slices = self.driver.shell.GetSlices({'slice_name': slicename}) 
119         if not slices:
120             slice = {'slice_name': slicename}
121             # add the slice                          
122             slice['slice_id'] = self.driver.shell.AddSlice(slice)
123             slice['node_ids'] = []
124             slice['user_ids'] = []
125         else:
126             slice = slices[0]
127             if slice_record and slice_record.get('expires'):
128                 requested_expires = int(datetime_to_epoch(utcparse(slice_record['expires'])))
129                 if requested_expires and slice['expires'] != requested_expires:
130                     self.driver.shell.UpdateSlice( {'slice_id': slice['slice_id'], 'fields':{'expires' : expiration}})
131        
132         return slice
133
134     def verify_users(self, slice_hrn, slice_record, users, options=None):
135         if options is None: options={}
136         slice_name = hrn_to_dummy_slicename(slice_hrn)
137         users_by_email = {}
138         for user in users:
139             user['urn'] = user['urn'].lower()
140             hrn, type = urn_to_hrn(user['urn'])
141             username = get_leaf(hrn)
142             user['username'] = username
143
144             if 'email' in user:
145                 user['email'] = user['email'].lower() 
146                 users_by_email[user['email']] = user
147         
148         # start building a list of existing users
149         existing_users_by_email = {}
150         existing_slice_users_by_email = {}
151         existing_users = self.driver.shell.GetUsers()
152         existing_slice_users_ids = self.driver.shell.GetSlices({'slice_name': slice_name})[0]['user_ids']
153         for user in existing_users:
154             existing_users_by_email[user['email']] = user  
155             if user['user_id'] in existing_slice_users_ids:
156                 existing_slice_users_by_email[user['email']] = user
157                 
158         add_users_by_email = set(users_by_email).difference(existing_slice_user_by_email)
159         delete_users_by_email = set(existing_slice_user_by_email).difference(users_by_email)
160         try:
161             for user in add_users_by_email: 
162                 self.driver.shell.AddUser()
163         except: 
164             pass
165             
166
167     def verify_keys(self, old_users, new_users, options=None):
168         if options is None: options={}
169         # existing keys 
170         existing_keys = []
171         for user in old_users:
172              existing_keys.append(user['keys'])
173         userdict = {}
174         for user in old_users:
175             userdict[user['email']] = user    
176     
177         # add new keys
178         requested_keys = []
179         updated_users = []
180         for user in new_users:
181             user_keys = user.get('keys', [])
182             updated_users.append(user)
183             for key_string in user_keys:
184                 requested_keys.append(key_string)
185                 if key_string not in existing_keys:
186                     key = key_string
187                     try:
188                         self.driver.shell.AddUserKey({'user_id': user['user_id'], 'key':key})
189                             
190                     except:
191                         pass        
192         # remove old keys (only if we are not appending)
193         append = options.get('append', True)
194         if append == False: 
195             removed_keys = set(existing_keys).difference(requested_keys)
196             for key in removed_keys:
197                  try:
198                      self.driver.shell.DeleteKey({'key': key})
199                  except:
200                      pass   
201