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

Source Code for Module oauth2client.multistore_file

  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['credential'])) 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)
410