__author__ = 'jbeda@google.com (Joe Beda)'
import base64
import errno
import logging
import os
import threading
from anyjson import simplejson
from oauth2client.client import Storage as BaseStorage
from oauth2client.client import Credentials
from oauth2client import util
from locked_file import LockedFile
logger = logging.getLogger(__name__)
_multistores = {}
_multistores_lock = threading.Lock()
class Error(Exception):
pass
class NewerCredentialStoreError(Error):
pass
@util.positional(4)
def get_credential_storage(filename, client_id, user_agent, scope,
warn_on_readonly=True):
key = {'clientId': client_id, 'userAgent': user_agent,
'scope': util.scopes_to_string(scope)}
return get_credential_storage_custom_key(
filename, key, warn_on_readonly=warn_on_readonly)
@util.positional(2)
def get_credential_storage_custom_string_key(
filename, key_string, warn_on_readonly=True):
key_dict = {'key': key_string}
return get_credential_storage_custom_key(
filename, key_dict, warn_on_readonly=warn_on_readonly)
@util.positional(2)
def get_credential_storage_custom_key(
filename, key_dict, warn_on_readonly=True):
multistore = _get_multistore(filename, warn_on_readonly=warn_on_readonly)
key = util.dict_to_tuple_key(key_dict)
return multistore._get_storage(key)
@util.positional(1)
def get_all_credential_keys(filename, warn_on_readonly=True):
multistore = _get_multistore(filename, warn_on_readonly=warn_on_readonly)
multistore._lock()
try:
return multistore._get_all_credential_keys()
finally:
multistore._unlock()
@util.positional(1)
def _get_multistore(filename, warn_on_readonly=True):
filename = os.path.expanduser(filename)
_multistores_lock.acquire()
try:
multistore = _multistores.setdefault(
filename, _MultiStore(filename, warn_on_readonly=warn_on_readonly))
finally:
_multistores_lock.release()
return multistore
class _MultiStore(object):
@util.positional(2)
def __init__(self, filename, warn_on_readonly=True):
self._file = LockedFile(filename, 'r+b', 'rb')
self._thread_lock = threading.Lock()
self._read_only = False
self._warn_on_readonly = warn_on_readonly
self._create_file_if_needed()
self._data = None
class _Storage(BaseStorage):
def __init__(self, multistore, key):
self._multistore = multistore
self._key = key
def acquire_lock(self):
self._multistore._lock()
def release_lock(self):
self._multistore._unlock()
def locked_get(self):
credential = self._multistore._get_credential(self._key)
if credential:
credential.set_store(self)
return credential
def locked_put(self, credentials):
self._multistore._update_credential(self._key, credentials)
def locked_delete(self):
self._multistore._delete_credential(self._key)
def _create_file_if_needed(self):
if not os.path.exists(self._file.filename()):
old_umask = os.umask(0177)
try:
open(self._file.filename(), 'a+b').close()
finally:
os.umask(old_umask)
def _lock(self):
self._thread_lock.acquire()
self._file.open_and_lock()
if not self._file.is_locked():
self._read_only = True
if self._warn_on_readonly:
logger.warn('The credentials file (%s) is not writable. Opening in '
'read-only mode. Any refreshed credentials will only be '
'valid for this run.' % self._file.filename())
if os.path.getsize(self._file.filename()) == 0:
logger.debug('Initializing empty multistore file')
self._data = {}
self._write()
elif not self._read_only or self._data is None:
self._refresh_data_cache()
def _unlock(self):
self._file.unlock_and_close()
self._thread_lock.release()
def _locked_json_read(self):
assert self._thread_lock.locked()
self._file.file_handle().seek(0)
return simplejson.load(self._file.file_handle())
def _locked_json_write(self, data):
assert self._thread_lock.locked()
if self._read_only:
return
self._file.file_handle().seek(0)
simplejson.dump(data, self._file.file_handle(), sort_keys=True, indent=2)
self._file.file_handle().truncate()
def _refresh_data_cache(self):
self._data = {}
try:
raw_data = self._locked_json_read()
except Exception:
logger.warn('Credential data store could not be loaded. '
'Will ignore and overwrite.')
return
version = 0
try:
version = raw_data['file_version']
except Exception:
logger.warn('Missing version for credential data store. It may be '
'corrupt or an old version. Overwriting.')
if version > 1:
raise NewerCredentialStoreError(
'Credential file has file_version of %d. '
'Only file_version of 1 is supported.' % version)
credentials = []
try:
credentials = raw_data['data']
except (TypeError, KeyError):
pass
for cred_entry in credentials:
try:
(key, credential) = self._decode_credential_from_json(cred_entry)
self._data[key] = credential
except:
logger.info('Error decoding credential, skipping', exc_info=True)
def _decode_credential_from_json(self, cred_entry):
raw_key = cred_entry['key']
key = util.dict_to_tuple_key(raw_key)
credential = None
credential = Credentials.new_from_json(simplejson.dumps(cred_entry['credential']))
return (key, credential)
def _write(self):
raw_data = {'file_version': 1}
raw_creds = []
raw_data['data'] = raw_creds
for (cred_key, cred) in self._data.items():
raw_key = dict(cred_key)
raw_cred = simplejson.loads(cred.to_json())
raw_creds.append({'key': raw_key, 'credential': raw_cred})
self._locked_json_write(raw_data)
def _get_all_credential_keys(self):
return [dict(key) for key in self._data.keys()]
def _get_credential(self, key):
return self._data.get(key, None)
def _update_credential(self, key, cred):
self._data[key] = cred
self._write()
def _delete_credential(self, key):
try:
del self._data[key]
except KeyError:
pass
self._write()
def _get_storage(self, key):
return self._Storage(self, key)