Source Code for Module oauth2client.multistore_file

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