OLD | NEW |
(Empty) | |
| 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 fcntl |
| 37 import logging |
| 38 import os |
| 39 import threading |
| 40 |
| 41 from anyjson import simplejson |
| 42 from client import Storage as BaseStorage |
| 43 from client import Credentials |
| 44 |
| 45 logger = logging.getLogger(__name__) |
| 46 |
| 47 # A dict from 'filename'->_MultiStore instances |
| 48 _multistores = {} |
| 49 _multistores_lock = threading.Lock() |
| 50 |
| 51 |
| 52 class Error(Exception): |
| 53 """Base error for this module.""" |
| 54 pass |
| 55 |
| 56 |
| 57 class NewerCredentialStoreError(Error): |
| 58 """The credential store is a newer version that supported.""" |
| 59 pass |
| 60 |
| 61 |
| 62 def get_credential_storage(filename, client_id, user_agent, scope, |
| 63 warn_on_readonly=True): |
| 64 """Get a Storage instance for a credential. |
| 65 |
| 66 Args: |
| 67 filename: The JSON file storing a set of credentials |
| 68 client_id: The client_id for the credential |
| 69 user_agent: The user agent for the credential |
| 70 scope: string or list of strings, Scope(s) being requested |
| 71 warn_on_readonly: if True, log a warning if the store is readonly |
| 72 |
| 73 Returns: |
| 74 An object derived from client.Storage for getting/setting the |
| 75 credential. |
| 76 """ |
| 77 filename = os.path.realpath(os.path.expanduser(filename)) |
| 78 _multistores_lock.acquire() |
| 79 try: |
| 80 multistore = _multistores.setdefault( |
| 81 filename, _MultiStore(filename, warn_on_readonly)) |
| 82 finally: |
| 83 _multistores_lock.release() |
| 84 if type(scope) is list: |
| 85 scope = ' '.join(scope) |
| 86 return multistore._get_storage(client_id, user_agent, scope) |
| 87 |
| 88 |
| 89 class _MultiStore(object): |
| 90 """A file backed store for multiple credentials.""" |
| 91 |
| 92 def __init__(self, filename, warn_on_readonly=True): |
| 93 """Initialize the class. |
| 94 |
| 95 This will create the file if necessary. |
| 96 """ |
| 97 self._filename = filename |
| 98 self._thread_lock = threading.Lock() |
| 99 self._file_handle = None |
| 100 self._read_only = False |
| 101 self._warn_on_readonly = warn_on_readonly |
| 102 |
| 103 self._create_file_if_needed() |
| 104 |
| 105 # Cache of deserialized store. This is only valid after the |
| 106 # _MultiStore is locked or _refresh_data_cache is called. This is |
| 107 # of the form of: |
| 108 # |
| 109 # (client_id, user_agent, scope) -> OAuth2Credential |
| 110 # |
| 111 # If this is None, then the store hasn't been read yet. |
| 112 self._data = None |
| 113 |
| 114 class _Storage(BaseStorage): |
| 115 """A Storage object that knows how to read/write a single credential.""" |
| 116 |
| 117 def __init__(self, multistore, client_id, user_agent, scope): |
| 118 self._multistore = multistore |
| 119 self._client_id = client_id |
| 120 self._user_agent = user_agent |
| 121 self._scope = scope |
| 122 |
| 123 def acquire_lock(self): |
| 124 """Acquires any lock necessary to access this Storage. |
| 125 |
| 126 This lock is not reentrant. |
| 127 """ |
| 128 self._multistore._lock() |
| 129 |
| 130 def release_lock(self): |
| 131 """Release the Storage lock. |
| 132 |
| 133 Trying to release a lock that isn't held will result in a |
| 134 RuntimeError. |
| 135 """ |
| 136 self._multistore._unlock() |
| 137 |
| 138 def locked_get(self): |
| 139 """Retrieve credential. |
| 140 |
| 141 The Storage lock must be held when this is called. |
| 142 |
| 143 Returns: |
| 144 oauth2client.client.Credentials |
| 145 """ |
| 146 credential = self._multistore._get_credential( |
| 147 self._client_id, self._user_agent, self._scope) |
| 148 if credential: |
| 149 credential.set_store(self) |
| 150 return credential |
| 151 |
| 152 def locked_put(self, credentials): |
| 153 """Write a credential. |
| 154 |
| 155 The Storage lock must be held when this is called. |
| 156 |
| 157 Args: |
| 158 credentials: Credentials, the credentials to store. |
| 159 """ |
| 160 self._multistore._update_credential(credentials, self._scope) |
| 161 |
| 162 def _create_file_if_needed(self): |
| 163 """Create an empty file if necessary. |
| 164 |
| 165 This method will not initialize the file. Instead it implements a |
| 166 simple version of "touch" to ensure the file has been created. |
| 167 """ |
| 168 if not os.path.exists(self._filename): |
| 169 old_umask = os.umask(0177) |
| 170 try: |
| 171 open(self._filename, 'a+b').close() |
| 172 finally: |
| 173 os.umask(old_umask) |
| 174 |
| 175 def _lock(self): |
| 176 """Lock the entire multistore.""" |
| 177 self._thread_lock.acquire() |
| 178 # Check to see if the file is writeable. |
| 179 try: |
| 180 self._file_handle = open(self._filename, 'r+b') |
| 181 fcntl.lockf(self._file_handle.fileno(), fcntl.LOCK_EX) |
| 182 except IOError, e: |
| 183 if e.errno != errno.EACCES: |
| 184 raise e |
| 185 self._file_handle = open(self._filename, 'rb') |
| 186 self._read_only = True |
| 187 if self._warn_on_readonly: |
| 188 logger.warn('The credentials file (%s) is not writable. Opening in ' |
| 189 'read-only mode. Any refreshed credentials will only be ' |
| 190 'valid for this run.' % self._filename) |
| 191 if os.path.getsize(self._filename) == 0: |
| 192 logger.debug('Initializing empty multistore file') |
| 193 # The multistore is empty so write out an empty file. |
| 194 self._data = {} |
| 195 self._write() |
| 196 elif not self._read_only or self._data is None: |
| 197 # Only refresh the data if we are read/write or we haven't |
| 198 # cached the data yet. If we are readonly, we assume is isn't |
| 199 # changing out from under us and that we only have to read it |
| 200 # once. This prevents us from whacking any new access keys that |
| 201 # we have cached in memory but were unable to write out. |
| 202 self._refresh_data_cache() |
| 203 |
| 204 def _unlock(self): |
| 205 """Release the lock on the multistore.""" |
| 206 if not self._read_only: |
| 207 fcntl.lockf(self._file_handle.fileno(), fcntl.LOCK_UN) |
| 208 self._file_handle.close() |
| 209 self._thread_lock.release() |
| 210 |
| 211 def _locked_json_read(self): |
| 212 """Get the raw content of the multistore file. |
| 213 |
| 214 The multistore must be locked when this is called. |
| 215 |
| 216 Returns: |
| 217 The contents of the multistore decoded as JSON. |
| 218 """ |
| 219 assert self._thread_lock.locked() |
| 220 self._file_handle.seek(0) |
| 221 return simplejson.load(self._file_handle) |
| 222 |
| 223 def _locked_json_write(self, data): |
| 224 """Write a JSON serializable data structure to the multistore. |
| 225 |
| 226 The multistore must be locked when this is called. |
| 227 |
| 228 Args: |
| 229 data: The data to be serialized and written. |
| 230 """ |
| 231 assert self._thread_lock.locked() |
| 232 if self._read_only: |
| 233 return |
| 234 self._file_handle.seek(0) |
| 235 simplejson.dump(data, self._file_handle, sort_keys=True, indent=2) |
| 236 self._file_handle.truncate() |
| 237 |
| 238 def _refresh_data_cache(self): |
| 239 """Refresh the contents of the multistore. |
| 240 |
| 241 The multistore must be locked when this is called. |
| 242 |
| 243 Raises: |
| 244 NewerCredentialStoreError: Raised when a newer client has written the |
| 245 store. |
| 246 """ |
| 247 self._data = {} |
| 248 try: |
| 249 raw_data = self._locked_json_read() |
| 250 except Exception: |
| 251 logger.warn('Credential data store could not be loaded. ' |
| 252 'Will ignore and overwrite.') |
| 253 return |
| 254 |
| 255 version = 0 |
| 256 try: |
| 257 version = raw_data['file_version'] |
| 258 except Exception: |
| 259 logger.warn('Missing version for credential data store. It may be ' |
| 260 'corrupt or an old version. Overwriting.') |
| 261 if version > 1: |
| 262 raise NewerCredentialStoreError( |
| 263 'Credential file has file_version of %d. ' |
| 264 'Only file_version of 1 is supported.' % version) |
| 265 |
| 266 credentials = [] |
| 267 try: |
| 268 credentials = raw_data['data'] |
| 269 except (TypeError, KeyError): |
| 270 pass |
| 271 |
| 272 for cred_entry in credentials: |
| 273 try: |
| 274 (key, credential) = self._decode_credential_from_json(cred_entry) |
| 275 self._data[key] = credential |
| 276 except: |
| 277 # If something goes wrong loading a credential, just ignore it |
| 278 logger.info('Error decoding credential, skipping', exc_info=True) |
| 279 |
| 280 def _decode_credential_from_json(self, cred_entry): |
| 281 """Load a credential from our JSON serialization. |
| 282 |
| 283 Args: |
| 284 cred_entry: A dict entry from the data member of our format |
| 285 |
| 286 Returns: |
| 287 (key, cred) where the key is the key tuple and the cred is the |
| 288 OAuth2Credential object. |
| 289 """ |
| 290 raw_key = cred_entry['key'] |
| 291 client_id = raw_key['clientId'] |
| 292 user_agent = raw_key['userAgent'] |
| 293 scope = raw_key['scope'] |
| 294 key = (client_id, user_agent, scope) |
| 295 credential = None |
| 296 credential = Credentials.new_from_json(simplejson.dumps(cred_entry['credenti
al'])) |
| 297 return (key, credential) |
| 298 |
| 299 def _write(self): |
| 300 """Write the cached data back out. |
| 301 |
| 302 The multistore must be locked. |
| 303 """ |
| 304 raw_data = {'file_version': 1} |
| 305 raw_creds = [] |
| 306 raw_data['data'] = raw_creds |
| 307 for (cred_key, cred) in self._data.items(): |
| 308 raw_key = { |
| 309 'clientId': cred_key[0], |
| 310 'userAgent': cred_key[1], |
| 311 'scope': cred_key[2] |
| 312 } |
| 313 raw_cred = simplejson.loads(cred.to_json()) |
| 314 raw_creds.append({'key': raw_key, 'credential': raw_cred}) |
| 315 self._locked_json_write(raw_data) |
| 316 |
| 317 def _get_credential(self, client_id, user_agent, scope): |
| 318 """Get a credential from the multistore. |
| 319 |
| 320 The multistore must be locked. |
| 321 |
| 322 Args: |
| 323 client_id: The client_id for the credential |
| 324 user_agent: The user agent for the credential |
| 325 scope: A string for the scope(s) being requested |
| 326 |
| 327 Returns: |
| 328 The credential specified or None if not present |
| 329 """ |
| 330 key = (client_id, user_agent, scope) |
| 331 |
| 332 return self._data.get(key, None) |
| 333 |
| 334 def _update_credential(self, cred, scope): |
| 335 """Update a credential and write the multistore. |
| 336 |
| 337 This must be called when the multistore is locked. |
| 338 |
| 339 Args: |
| 340 cred: The OAuth2Credential to update/set |
| 341 scope: The scope(s) that this credential covers |
| 342 """ |
| 343 key = (cred.client_id, cred.user_agent, scope) |
| 344 self._data[key] = cred |
| 345 self._write() |
| 346 |
| 347 def _get_storage(self, client_id, user_agent, scope): |
| 348 """Get a Storage object to get/set a credential. |
| 349 |
| 350 This Storage is a 'view' into the multistore. |
| 351 |
| 352 Args: |
| 353 client_id: The client_id for the credential |
| 354 user_agent: The user agent for the credential |
| 355 scope: A string for the scope(s) being requested |
| 356 |
| 357 Returns: |
| 358 A Storage object that can be used to get/set this cred |
| 359 """ |
| 360 return self._Storage(self, client_id, user_agent, scope) |
OLD | NEW |