blob: ce7a5194feddfe92718fcc56de46e9a795408afc [file] [log] [blame]
Joe Gregoriod60d1e72013-06-25 15:14:28 -04001# 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.
Joe Gregorio9da2ad82011-09-11 14:04:44 -040014
15"""Multi-credential file store with lock support.
16
17This module implements a JSON credential store where multiple
Joe Gregorioe2233cd2013-01-24 15:46:23 -050018credentials can be stored in one file. That file supports locking
Joe Gregorio9da2ad82011-09-11 14:04:44 -040019both in a single process and across processes.
20
21The credential themselves are keyed off of:
22* client_id
23* user_agent
24* scope
25
26The 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 },
Joe Gregorio562b7312011-09-15 09:06:38 -040036 'credential': {
37 # JSON serialized Credentials.
38 }
Joe Gregorio9da2ad82011-09-11 14:04:44 -040039 }
40 ]
41}
42"""
43
44__author__ = 'jbeda@google.com (Joe Beda)'
45
46import base64
Joe Gregorio9b8bec62012-01-17 11:35:32 -050047import errno
Joe Gregorio9da2ad82011-09-11 14:04:44 -040048import logging
49import os
Joe Gregorio9da2ad82011-09-11 14:04:44 -040050import threading
51
Joe Gregorio549230c2012-01-11 10:38:05 -050052from anyjson import simplejson
Joe Gregorio68a8cfe2012-08-03 16:17:40 -040053from oauth2client.client import Storage as BaseStorage
54from oauth2client.client import Credentials
55from oauth2client import util
Joe Gregorio2dfd74e2012-06-13 12:18:30 -040056from locked_file import LockedFile
Joe Gregorio9da2ad82011-09-11 14:04:44 -040057
58logger = logging.getLogger(__name__)
59
60# A dict from 'filename'->_MultiStore instances
61_multistores = {}
62_multistores_lock = threading.Lock()
63
64
65class Error(Exception):
66 """Base error for this module."""
67 pass
68
69
70class NewerCredentialStoreError(Error):
71 """The credential store is a newer version that supported."""
72 pass
73
74
Joe Gregorio68a8cfe2012-08-03 16:17:40 -040075@util.positional(4)
Joe Gregorio9da2ad82011-09-11 14:04:44 -040076def 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
Joe Gregorio5cf5d122012-11-16 16:36:12 -050084 scope: string or iterable of strings, Scope(s) being requested
Joe Gregorio9da2ad82011-09-11 14:04:44 -040085 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 """
Joe Gregorioe2233cd2013-01-24 15:46:23 -050091 # 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)
99def 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)
122def 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 """
Joe Gregorio48d10b02013-05-14 10:30:40 -0400140 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)
146def 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)
167def _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 """
Joe Gregorio0fd18532012-08-24 15:54:40 -0400177 filename = os.path.expanduser(filename)
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400178 _multistores_lock.acquire()
179 try:
180 multistore = _multistores.setdefault(
Joe Gregorio68a8cfe2012-08-03 16:17:40 -0400181 filename, _MultiStore(filename, warn_on_readonly=warn_on_readonly))
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400182 finally:
183 _multistores_lock.release()
Joe Gregorio48d10b02013-05-14 10:30:40 -0400184 return multistore
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400185
186
187class _MultiStore(object):
188 """A file backed store for multiple credentials."""
189
Joe Gregorio68a8cfe2012-08-03 16:17:40 -0400190 @util.positional(2)
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400191 def __init__(self, filename, warn_on_readonly=True):
192 """Initialize the class.
193
194 This will create the file if necessary.
195 """
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400196 self._file = LockedFile(filename, 'r+b', 'rb')
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400197 self._thread_lock = threading.Lock()
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400198 self._read_only = False
199 self._warn_on_readonly = warn_on_readonly
200
201 self._create_file_if_needed()
202
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500203 # Cache of deserialized store. This is only valid after the
204 # _MultiStore is locked or _refresh_data_cache is called. This is
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400205 # of the form of:
206 #
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500207 # ((key, value), (key, value)...) -> OAuth2Credential
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400208 #
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
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500215 def __init__(self, multistore, key):
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400216 self._multistore = multistore
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500217 self._key = key
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400218
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 """
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500242 credential = self._multistore._get_credential(self._key)
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400243 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 """
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500255 self._multistore._update_credential(self._key, credentials)
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400256
Joe Gregorioec75dc12012-02-06 13:40:42 -0500257 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 """
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500265 self._multistore._delete_credential(self._key)
Joe Gregorioec75dc12012-02-06 13:40:42 -0500266
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400267 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 """
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400273 if not os.path.exists(self._file.filename()):
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400274 old_umask = os.umask(0177)
275 try:
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400276 open(self._file.filename(), 'a+b').close()
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400277 finally:
278 os.umask(old_umask)
279
280 def _lock(self):
281 """Lock the entire multistore."""
282 self._thread_lock.acquire()
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400283 self._file.open_and_lock()
284 if not self._file.is_locked():
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400285 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 '
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400289 'valid for this run.' % self._file.filename())
290 if os.path.getsize(self._file.filename()) == 0:
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400291 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
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500297 # cached the data yet. If we are readonly, we assume is isn't
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400298 # changing out from under us and that we only have to read it
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500299 # once. This prevents us from whacking any new access keys that
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400300 # 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."""
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400305 self._file.unlock_and_close()
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400306 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()
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400317 self._file.file_handle().seek(0)
318 return simplejson.load(self._file.file_handle())
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400319
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
Joe Gregorio2dfd74e2012-06-13 12:18:30 -0400331 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()
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400334
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']
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500388 key = util.dict_to_tuple_key(raw_key)
Joe Gregorio562b7312011-09-15 09:06:38 -0400389 credential = None
390 credential = Credentials.new_from_json(simplejson.dumps(cred_entry['credential']))
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400391 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():
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500402 raw_key = dict(cred_key)
Joe Gregorio562b7312011-09-15 09:06:38 -0400403 raw_cred = simplejson.loads(cred.to_json())
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400404 raw_creds.append({'key': raw_key, 'credential': raw_cred})
405 self._locked_json_write(raw_data)
406
Joe Gregorio48d10b02013-05-14 10:30:40 -0400407 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
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500415 def _get_credential(self, key):
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400416 """Get a credential from the multistore.
417
418 The multistore must be locked.
419
420 Args:
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500421 key: The key used to retrieve the credential
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400422
423 Returns:
424 The credential specified or None if not present
425 """
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400426 return self._data.get(key, None)
427
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500428 def _update_credential(self, key, cred):
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400429 """Update a credential and write the multistore.
430
431 This must be called when the multistore is locked.
432
433 Args:
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500434 key: The key used to retrieve the credential
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400435 cred: The OAuth2Credential to update/set
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400436 """
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400437 self._data[key] = cred
438 self._write()
439
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500440 def _delete_credential(self, key):
Joe Gregorioec75dc12012-02-06 13:40:42 -0500441 """Delete a credential and write the multistore.
442
443 This must be called when the multistore is locked.
444
445 Args:
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500446 key: The key used to retrieve the credential
Joe Gregorioec75dc12012-02-06 13:40:42 -0500447 """
Joe Gregorioec75dc12012-02-06 13:40:42 -0500448 try:
449 del self._data[key]
450 except KeyError:
451 pass
452 self._write()
453
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500454 def _get_storage(self, key):
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400455 """Get a Storage object to get/set a credential.
456
457 This Storage is a 'view' into the multistore.
458
459 Args:
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500460 key: The key used to retrieve the credential
Joe Gregorio9da2ad82011-09-11 14:04:44 -0400461
462 Returns:
463 A Storage object that can be used to get/set this cred
464 """
Joe Gregorioe2233cd2013-01-24 15:46:23 -0500465 return self._Storage(self, key)