OLD | NEW |
(Empty) | |
| 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['credenti
al'])) |
| 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) |
OLD | NEW |