Package oauth2client :: Module multistore_file
[hide private]
[frames] | no frames]

Source Code for Module oauth2client.multistore_file

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