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