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

# 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