OLD | NEW |
| (Empty) |
1 # Copyright 2011 Google Inc. All Rights Reserved. | |
2 | |
3 """Multi-credential file store with lock support. | |
4 | |
5 This module implements a JSON credential store where multiple | |
6 credentials can be stored in one file. That file supports locking | |
7 both in a single process and across processes. | |
8 | |
9 The credential themselves are keyed off of: | |
10 * client_id | |
11 * user_agent | |
12 * scope | |
13 | |
14 The format of the stored data is like so: | |
15 { | |
16 'file_version': 1, | |
17 'data': [ | |
18 { | |
19 'key': { | |
20 'clientId': '<client id>', | |
21 'userAgent': '<user agent>', | |
22 'scope': '<scope>' | |
23 }, | |
24 'credential': { | |
25 # JSON serialized Credentials. | |
26 } | |
27 } | |
28 ] | |
29 } | |
30 """ | |
31 | |
32 __author__ = 'jbeda@google.com (Joe Beda)' | |
33 | |
34 import base64 | |
35 import errno | |
36 import logging | |
37 import os | |
38 import threading | |
39 | |
40 from anyjson import simplejson | |
41 from oauth2client.client import Storage as BaseStorage | |
42 from oauth2client.client import Credentials | |
43 from oauth2client import util | |
44 from locked_file import LockedFile | |
45 | |
46 logger = logging.getLogger(__name__) | |
47 | |
48 # A dict from 'filename'->_MultiStore instances | |
49 _multistores = {} | |
50 _multistores_lock = threading.Lock() | |
51 | |
52 | |
53 class Error(Exception): | |
54 """Base error for this module.""" | |
55 pass | |
56 | |
57 | |
58 class NewerCredentialStoreError(Error): | |
59 """The credential store is a newer version that supported.""" | |
60 pass | |
61 | |
62 | |
63 @util.positional(4) | |
64 def get_credential_storage(filename, client_id, user_agent, scope, | |
65 warn_on_readonly=True): | |
66 """Get a Storage instance for a credential. | |
67 | |
68 Args: | |
69 filename: The JSON file storing a set of credentials | |
70 client_id: The client_id for the credential | |
71 user_agent: The user agent for the credential | |
72 scope: string or iterable of strings, Scope(s) being requested | |
73 warn_on_readonly: if True, log a warning if the store is readonly | |
74 | |
75 Returns: | |
76 An object derived from client.Storage for getting/setting the | |
77 credential. | |
78 """ | |
79 # Recreate the legacy key with these specific parameters | |
80 key = {'clientId': client_id, 'userAgent': user_agent, | |
81 'scope': util.scopes_to_string(scope)} | |
82 return get_credential_storage_custom_key( | |
83 filename, key, warn_on_readonly=warn_on_readonly) | |
84 | |
85 | |
86 @util.positional(2) | |
87 def get_credential_storage_custom_string_key( | |
88 filename, key_string, warn_on_readonly=True): | |
89 """Get a Storage instance for a credential using a single string as a key. | |
90 | |
91 Allows you to provide a string as a custom key that will be used for | |
92 credential storage and retrieval. | |
93 | |
94 Args: | |
95 filename: The JSON file storing a set of credentials | |
96 key_string: A string to use as the key for storing this credential. | |
97 warn_on_readonly: if True, log a warning if the store is readonly | |
98 | |
99 Returns: | |
100 An object derived from client.Storage for getting/setting the | |
101 credential. | |
102 """ | |
103 # Create a key dictionary that can be used | |
104 key_dict = {'key': key_string} | |
105 return get_credential_storage_custom_key( | |
106 filename, key_dict, warn_on_readonly=warn_on_readonly) | |
107 | |
108 | |
109 @util.positional(2) | |
110 def get_credential_storage_custom_key( | |
111 filename, key_dict, warn_on_readonly=True): | |
112 """Get a Storage instance for a credential using a dictionary as a key. | |
113 | |
114 Allows you to provide a dictionary as a custom key that will be used for | |
115 credential storage and retrieval. | |
116 | |
117 Args: | |
118 filename: The JSON file storing a set of credentials | |
119 key_dict: A dictionary to use as the key for storing this credential. There | |
120 is no ordering of the keys in the dictionary. Logically equivalent | |
121 dictionaries will produce equivalent storage keys. | |
122 warn_on_readonly: if True, log a warning if the store is readonly | |
123 | |
124 Returns: | |
125 An object derived from client.Storage for getting/setting the | |
126 credential. | |
127 """ | |
128 filename = os.path.expanduser(filename) | |
129 _multistores_lock.acquire() | |
130 try: | |
131 multistore = _multistores.setdefault( | |
132 filename, _MultiStore(filename, warn_on_readonly=warn_on_readonly)) | |
133 finally: | |
134 _multistores_lock.release() | |
135 key = util.dict_to_tuple_key(key_dict) | |
136 return multistore._get_storage(key) | |
137 | |
138 | |
139 class _MultiStore(object): | |
140 """A file backed store for multiple credentials.""" | |
141 | |
142 @util.positional(2) | |
143 def __init__(self, filename, warn_on_readonly=True): | |
144 """Initialize the class. | |
145 | |
146 This will create the file if necessary. | |
147 """ | |
148 self._file = LockedFile(filename, 'r+b', 'rb') | |
149 self._thread_lock = threading.Lock() | |
150 self._read_only = False | |
151 self._warn_on_readonly = warn_on_readonly | |
152 | |
153 self._create_file_if_needed() | |
154 | |
155 # Cache of deserialized store. This is only valid after the | |
156 # _MultiStore is locked or _refresh_data_cache is called. This is | |
157 # of the form of: | |
158 # | |
159 # ((key, value), (key, value)...) -> OAuth2Credential | |
160 # | |
161 # If this is None, then the store hasn't been read yet. | |
162 self._data = None | |
163 | |
164 class _Storage(BaseStorage): | |
165 """A Storage object that knows how to read/write a single credential.""" | |
166 | |
167 def __init__(self, multistore, key): | |
168 self._multistore = multistore | |
169 self._key = key | |
170 | |
171 def acquire_lock(self): | |
172 """Acquires any lock necessary to access this Storage. | |
173 | |
174 This lock is not reentrant. | |
175 """ | |
176 self._multistore._lock() | |
177 | |
178 def release_lock(self): | |
179 """Release the Storage lock. | |
180 | |
181 Trying to release a lock that isn't held will result in a | |
182 RuntimeError. | |
183 """ | |
184 self._multistore._unlock() | |
185 | |
186 def locked_get(self): | |
187 """Retrieve credential. | |
188 | |
189 The Storage lock must be held when this is called. | |
190 | |
191 Returns: | |
192 oauth2client.client.Credentials | |
193 """ | |
194 credential = self._multistore._get_credential(self._key) | |
195 if credential: | |
196 credential.set_store(self) | |
197 return credential | |
198 | |
199 def locked_put(self, credentials): | |
200 """Write a credential. | |
201 | |
202 The Storage lock must be held when this is called. | |
203 | |
204 Args: | |
205 credentials: Credentials, the credentials to store. | |
206 """ | |
207 self._multistore._update_credential(self._key, credentials) | |
208 | |
209 def locked_delete(self): | |
210 """Delete a credential. | |
211 | |
212 The Storage lock must be held when this is called. | |
213 | |
214 Args: | |
215 credentials: Credentials, the credentials to store. | |
216 """ | |
217 self._multistore._delete_credential(self._key) | |
218 | |
219 def _create_file_if_needed(self): | |
220 """Create an empty file if necessary. | |
221 | |
222 This method will not initialize the file. Instead it implements a | |
223 simple version of "touch" to ensure the file has been created. | |
224 """ | |
225 if not os.path.exists(self._file.filename()): | |
226 old_umask = os.umask(0177) | |
227 try: | |
228 open(self._file.filename(), 'a+b').close() | |
229 finally: | |
230 os.umask(old_umask) | |
231 | |
232 def _lock(self): | |
233 """Lock the entire multistore.""" | |
234 self._thread_lock.acquire() | |
235 self._file.open_and_lock() | |
236 if not self._file.is_locked(): | |
237 self._read_only = True | |
238 if self._warn_on_readonly: | |
239 logger.warn('The credentials file (%s) is not writable. Opening in ' | |
240 'read-only mode. Any refreshed credentials will only be ' | |
241 'valid for this run.' % self._file.filename()) | |
242 if os.path.getsize(self._file.filename()) == 0: | |
243 logger.debug('Initializing empty multistore file') | |
244 # The multistore is empty so write out an empty file. | |
245 self._data = {} | |
246 self._write() | |
247 elif not self._read_only or self._data is None: | |
248 # Only refresh the data if we are read/write or we haven't | |
249 # cached the data yet. If we are readonly, we assume is isn't | |
250 # changing out from under us and that we only have to read it | |
251 # once. This prevents us from whacking any new access keys that | |
252 # we have cached in memory but were unable to write out. | |
253 self._refresh_data_cache() | |
254 | |
255 def _unlock(self): | |
256 """Release the lock on the multistore.""" | |
257 self._file.unlock_and_close() | |
258 self._thread_lock.release() | |
259 | |
260 def _locked_json_read(self): | |
261 """Get the raw content of the multistore file. | |
262 | |
263 The multistore must be locked when this is called. | |
264 | |
265 Returns: | |
266 The contents of the multistore decoded as JSON. | |
267 """ | |
268 assert self._thread_lock.locked() | |
269 self._file.file_handle().seek(0) | |
270 return simplejson.load(self._file.file_handle()) | |
271 | |
272 def _locked_json_write(self, data): | |
273 """Write a JSON serializable data structure to the multistore. | |
274 | |
275 The multistore must be locked when this is called. | |
276 | |
277 Args: | |
278 data: The data to be serialized and written. | |
279 """ | |
280 assert self._thread_lock.locked() | |
281 if self._read_only: | |
282 return | |
283 self._file.file_handle().seek(0) | |
284 simplejson.dump(data, self._file.file_handle(), sort_keys=True, indent=2) | |
285 self._file.file_handle().truncate() | |
286 | |
287 def _refresh_data_cache(self): | |
288 """Refresh the contents of the multistore. | |
289 | |
290 The multistore must be locked when this is called. | |
291 | |
292 Raises: | |
293 NewerCredentialStoreError: Raised when a newer client has written the | |
294 store. | |
295 """ | |
296 self._data = {} | |
297 try: | |
298 raw_data = self._locked_json_read() | |
299 except Exception: | |
300 logger.warn('Credential data store could not be loaded. ' | |
301 'Will ignore and overwrite.') | |
302 return | |
303 | |
304 version = 0 | |
305 try: | |
306 version = raw_data['file_version'] | |
307 except Exception: | |
308 logger.warn('Missing version for credential data store. It may be ' | |
309 'corrupt or an old version. Overwriting.') | |
310 if version > 1: | |
311 raise NewerCredentialStoreError( | |
312 'Credential file has file_version of %d. ' | |
313 'Only file_version of 1 is supported.' % version) | |
314 | |
315 credentials = [] | |
316 try: | |
317 credentials = raw_data['data'] | |
318 except (TypeError, KeyError): | |
319 pass | |
320 | |
321 for cred_entry in credentials: | |
322 try: | |
323 (key, credential) = self._decode_credential_from_json(cred_entry) | |
324 self._data[key] = credential | |
325 except: | |
326 # If something goes wrong loading a credential, just ignore it | |
327 logger.info('Error decoding credential, skipping', exc_info=True) | |
328 | |
329 def _decode_credential_from_json(self, cred_entry): | |
330 """Load a credential from our JSON serialization. | |
331 | |
332 Args: | |
333 cred_entry: A dict entry from the data member of our format | |
334 | |
335 Returns: | |
336 (key, cred) where the key is the key tuple and the cred is the | |
337 OAuth2Credential object. | |
338 """ | |
339 raw_key = cred_entry['key'] | |
340 key = util.dict_to_tuple_key(raw_key) | |
341 credential = None | |
342 credential = Credentials.new_from_json(simplejson.dumps(cred_entry['credenti
al'])) | |
343 return (key, credential) | |
344 | |
345 def _write(self): | |
346 """Write the cached data back out. | |
347 | |
348 The multistore must be locked. | |
349 """ | |
350 raw_data = {'file_version': 1} | |
351 raw_creds = [] | |
352 raw_data['data'] = raw_creds | |
353 for (cred_key, cred) in self._data.items(): | |
354 raw_key = dict(cred_key) | |
355 raw_cred = simplejson.loads(cred.to_json()) | |
356 raw_creds.append({'key': raw_key, 'credential': raw_cred}) | |
357 self._locked_json_write(raw_data) | |
358 | |
359 def _get_credential(self, key): | |
360 """Get a credential from the multistore. | |
361 | |
362 The multistore must be locked. | |
363 | |
364 Args: | |
365 key: The key used to retrieve the credential | |
366 | |
367 Returns: | |
368 The credential specified or None if not present | |
369 """ | |
370 return self._data.get(key, None) | |
371 | |
372 def _update_credential(self, key, cred): | |
373 """Update a credential and write the multistore. | |
374 | |
375 This must be called when the multistore is locked. | |
376 | |
377 Args: | |
378 key: The key used to retrieve the credential | |
379 cred: The OAuth2Credential to update/set | |
380 """ | |
381 self._data[key] = cred | |
382 self._write() | |
383 | |
384 def _delete_credential(self, key): | |
385 """Delete a credential and write the multistore. | |
386 | |
387 This must be called when the multistore is locked. | |
388 | |
389 Args: | |
390 key: The key used to retrieve the credential | |
391 """ | |
392 try: | |
393 del self._data[key] | |
394 except KeyError: | |
395 pass | |
396 self._write() | |
397 | |
398 def _get_storage(self, key): | |
399 """Get a Storage object to get/set a credential. | |
400 | |
401 This Storage is a 'view' into the multistore. | |
402 | |
403 Args: | |
404 key: The key used to retrieve the credential | |
405 | |
406 Returns: | |
407 A Storage object that can be used to get/set this cred | |
408 """ | |
409 return self._Storage(self, key) | |
OLD | NEW |