Modify oauth2client.multistore_file to store and retrieve credentials using an arbitrary key.

Reviewed in https://codereview.appspot.com/7067054/.
This commit is contained in:
Joe Gregorio
2013-01-24 15:46:23 -05:00
parent 164f37ed3a
commit e2233cd7a2
7 changed files with 196 additions and 115 deletions

View File

@@ -72,6 +72,55 @@ def get_credential_storage(filename, client_id, user_agent, scope,
scope: string or iterable of strings, Scope(s) being requested
warn_on_readonly: if True, log a warning if the store is readonly
Returns:
An object derived from client.Storage for getting/setting the
credential.
"""
# Recreate the legacy key with these specific parameters
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):
"""Get a Storage instance for a credential using a single string as a key.
Allows you to provide a string as a custom key that will be used for
credential storage and retrieval.
Args:
filename: The JSON file storing a set of credentials
key_string: A string to use as the key for storing this credential.
warn_on_readonly: if True, log a warning if the store is readonly
Returns:
An object derived from client.Storage for getting/setting the
credential.
"""
# Create a key dictionary that can be used
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):
"""Get a Storage instance for a credential using a dictionary as a key.
Allows you to provide a dictionary as a custom key that will be used for
credential storage and retrieval.
Args:
filename: The JSON file storing a set of credentials
key_dict: A dictionary to use as the key for storing this credential. There
is no ordering of the keys in the dictionary. Logically equivalent
dictionaries will produce equivalent storage keys.
warn_on_readonly: if True, log a warning if the store is readonly
Returns:
An object derived from client.Storage for getting/setting the
credential.
@@ -83,8 +132,8 @@ def get_credential_storage(filename, client_id, user_agent, scope,
filename, _MultiStore(filename, warn_on_readonly=warn_on_readonly))
finally:
_multistores_lock.release()
scope = util.scopes_to_string(scope)
return multistore._get_storage(client_id, user_agent, scope)
key = util.dict_to_tuple_key(key_dict)
return multistore._get_storage(key)
class _MultiStore(object):
@@ -107,7 +156,7 @@ class _MultiStore(object):
# _MultiStore is locked or _refresh_data_cache is called. This is
# of the form of:
#
# (client_id, user_agent, scope) -> OAuth2Credential
# ((key, value), (key, value)...) -> OAuth2Credential
#
# If this is None, then the store hasn't been read yet.
self._data = None
@@ -115,11 +164,9 @@ class _MultiStore(object):
class _Storage(BaseStorage):
"""A Storage object that knows how to read/write a single credential."""
def __init__(self, multistore, client_id, user_agent, scope):
def __init__(self, multistore, key):
self._multistore = multistore
self._client_id = client_id
self._user_agent = user_agent
self._scope = scope
self._key = key
def acquire_lock(self):
"""Acquires any lock necessary to access this Storage.
@@ -144,8 +191,7 @@ class _MultiStore(object):
Returns:
oauth2client.client.Credentials
"""
credential = self._multistore._get_credential(
self._client_id, self._user_agent, self._scope)
credential = self._multistore._get_credential(self._key)
if credential:
credential.set_store(self)
return credential
@@ -158,7 +204,7 @@ class _MultiStore(object):
Args:
credentials: Credentials, the credentials to store.
"""
self._multistore._update_credential(credentials, self._scope)
self._multistore._update_credential(self._key, credentials)
def locked_delete(self):
"""Delete a credential.
@@ -168,8 +214,7 @@ class _MultiStore(object):
Args:
credentials: Credentials, the credentials to store.
"""
self._multistore._delete_credential(self._client_id, self._user_agent,
self._scope)
self._multistore._delete_credential(self._key)
def _create_file_if_needed(self):
"""Create an empty file if necessary.
@@ -292,10 +337,7 @@ class _MultiStore(object):
OAuth2Credential object.
"""
raw_key = cred_entry['key']
client_id = raw_key['clientId']
user_agent = raw_key['userAgent']
scope = raw_key['scope']
key = (client_id, user_agent, scope)
key = util.dict_to_tuple_key(raw_key)
credential = None
credential = Credentials.new_from_json(simplejson.dumps(cred_entry['credential']))
return (key, credential)
@@ -309,73 +351,59 @@ class _MultiStore(object):
raw_creds = []
raw_data['data'] = raw_creds
for (cred_key, cred) in self._data.items():
raw_key = {
'clientId': cred_key[0],
'userAgent': cred_key[1],
'scope': cred_key[2]
}
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_credential(self, client_id, user_agent, scope):
def _get_credential(self, key):
"""Get a credential from the multistore.
The multistore must be locked.
Args:
client_id: The client_id for the credential
user_agent: The user agent for the credential
scope: A string for the scope(s) being requested
key: The key used to retrieve the credential
Returns:
The credential specified or None if not present
"""
key = (client_id, user_agent, scope)
return self._data.get(key, None)
def _update_credential(self, cred, scope):
def _update_credential(self, key, cred):
"""Update a credential and write the multistore.
This must be called when the multistore is locked.
Args:
key: The key used to retrieve the credential
cred: The OAuth2Credential to update/set
scope: The scope(s) that this credential covers
"""
key = (cred.client_id, cred.user_agent, scope)
self._data[key] = cred
self._write()
def _delete_credential(self, client_id, user_agent, scope):
def _delete_credential(self, key):
"""Delete a credential and write the multistore.
This must be called when the multistore is locked.
Args:
client_id: The client_id for the credential
user_agent: The user agent for the credential
scope: The scope(s) that this credential covers
key: The key used to retrieve the credential
"""
key = (client_id, user_agent, scope)
try:
del self._data[key]
except KeyError:
pass
self._write()
def _get_storage(self, client_id, user_agent, scope):
def _get_storage(self, key):
"""Get a Storage object to get/set a credential.
This Storage is a 'view' into the multistore.
Args:
client_id: The client_id for the credential
user_agent: The user agent for the credential
scope: A string for the scope(s) being requested
key: The key used to retrieve the credential
Returns:
A Storage object that can be used to get/set this cred
"""
return self._Storage(self, client_id, user_agent, scope)
return self._Storage(self, key)

View File

@@ -145,3 +145,18 @@ def scopes_to_string(scopes):
return scopes
else:
return ' '.join(scopes)
def dict_to_tuple_key(dictionary):
"""Converts a dictionary to a tuple that can be used as an immutable key.
The resulting key is always sorted so that logically equivalent dictionaries
always produce an identical tuple for a key.
Args:
dictionary: the dictionary to use as the key.
Returns:
A tuple representing the dictionary in it's naturally sorted ordering.
"""
return tuple(sorted(dictionary.items()))

View File

@@ -35,6 +35,7 @@ from apiclient.http import HttpMockSequence
from oauth2client import file
from oauth2client import locked_file
from oauth2client import multistore_file
from oauth2client import util
from oauth2client.anyjson import simplejson
from oauth2client.client import AccessTokenCredentials
from oauth2client.client import AssertionCredentials
@@ -58,6 +59,21 @@ class OAuth2ClientFileTests(unittest.TestCase):
except OSError:
pass
def create_test_credentials(self):
access_token = 'foo'
client_secret = 'cOuDdkfjxxnv+'
refresh_token = '1/0/a.df219fjls0'
token_expiry = datetime.datetime.utcnow()
token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
user_agent = 'refresh_checker/1.0'
client_id = 'some_client_id'
credentials = OAuth2Credentials(
access_token, client_id, client_secret,
refresh_token, token_expiry, token_uri,
user_agent)
return credentials
def test_non_existent_file_storage(self):
s = file.Storage(FILENAME)
credentials = s.get()
@@ -78,18 +94,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
def test_pickle_and_json_interop(self):
# Write a file with a pickled OAuth2Credentials.
access_token = 'foo'
client_id = 'some_client_id'
client_secret = 'cOuDdkfjxxnv+'
refresh_token = '1/0/a.df219fjls0'
token_expiry = datetime.datetime.utcnow()
token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
user_agent = 'refresh_checker/1.0'
credentials = OAuth2Credentials(
access_token, client_id, client_secret,
refresh_token, token_expiry, token_uri,
user_agent)
credentials = self.create_test_credentials()
f = open(FILENAME, 'w')
pickle.dump(credentials, f)
@@ -112,18 +117,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
self.assertEquals(data['_module'], OAuth2Credentials.__module__)
def test_token_refresh(self):
access_token = 'foo'
client_id = 'some_client_id'
client_secret = 'cOuDdkfjxxnv+'
refresh_token = '1/0/a.df219fjls0'
token_expiry = datetime.datetime.utcnow()
token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
user_agent = 'refresh_checker/1.0'
credentials = OAuth2Credentials(
access_token, client_id, client_secret,
refresh_token, token_expiry, token_uri,
user_agent)
credentials = self.create_test_credentials()
s = file.Storage(FILENAME)
s.put(credentials)
@@ -136,18 +130,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
self.assertEquals(credentials.access_token, 'bar')
def test_credentials_delete(self):
access_token = 'foo'
client_id = 'some_client_id'
client_secret = 'cOuDdkfjxxnv+'
refresh_token = '1/0/a.df219fjls0'
token_expiry = datetime.datetime.utcnow()
token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
user_agent = 'refresh_checker/1.0'
credentials = OAuth2Credentials(
access_token, client_id, client_secret,
refresh_token, token_expiry, token_uri,
user_agent)
credentials = self.create_test_credentials()
s = file.Storage(FILENAME)
s.put(credentials)
@@ -175,18 +158,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
self.assertEquals('0600', oct(stat.S_IMODE(os.stat(FILENAME).st_mode)))
def test_read_only_file_fail_lock(self):
access_token = 'foo'
client_secret = 'cOuDdkfjxxnv+'
refresh_token = '1/0/a.df219fjls0'
token_expiry = datetime.datetime.utcnow()
token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
user_agent = 'refresh_checker/1.0'
client_id = 'some_client_id'
credentials = OAuth2Credentials(
access_token, client_id, client_secret,
refresh_token, token_expiry, token_uri,
user_agent)
credentials = self.create_test_credentials()
open(FILENAME, 'a+b').close()
os.chmod(FILENAME, 0400)
@@ -230,18 +202,7 @@ class OAuth2ClientFileTests(unittest.TestCase):
self.assertEquals(None, credentials)
def test_multistore_file(self):
access_token = 'foo'
client_secret = 'cOuDdkfjxxnv+'
refresh_token = '1/0/a.df219fjls0'
token_expiry = datetime.datetime.utcnow()
token_uri = 'https://www.google.com/accounts/o8/oauth2/token'
user_agent = 'refresh_checker/1.0'
client_id = 'some_client_id'
credentials = OAuth2Credentials(
access_token, client_id, client_secret,
refresh_token, token_expiry, token_uri,
user_agent)
credentials = self.create_test_credentials()
store = multistore_file.get_credential_storage(
FILENAME,
@@ -263,5 +224,65 @@ class OAuth2ClientFileTests(unittest.TestCase):
if os.name == 'posix':
self.assertEquals('0600', oct(stat.S_IMODE(os.stat(FILENAME).st_mode)))
def test_multistore_file_custom_key(self):
credentials = self.create_test_credentials()
custom_key = {'myapp': 'testing', 'clientid': 'some client'}
store = multistore_file.get_credential_storage_custom_key(
FILENAME, custom_key)
store.put(credentials)
stored_credentials = store.get()
self.assertNotEquals(None, stored_credentials)
self.assertEqual(credentials.access_token, stored_credentials.access_token)
store.delete()
stored_credentials = store.get()
self.assertEquals(None, stored_credentials)
def test_multistore_file_custom_string_key(self):
credentials = self.create_test_credentials()
# store with string key
store = multistore_file.get_credential_storage_custom_string_key(
FILENAME, 'mykey')
store.put(credentials)
stored_credentials = store.get()
self.assertNotEquals(None, stored_credentials)
self.assertEqual(credentials.access_token, stored_credentials.access_token)
# try retrieving with a dictionary
store_dict = multistore_file.get_credential_storage_custom_string_key(
FILENAME, {'key': 'mykey'})
stored_credentials = store.get()
self.assertNotEquals(None, stored_credentials)
self.assertEqual(credentials.access_token, stored_credentials.access_token)
store.delete()
stored_credentials = store.get()
self.assertEquals(None, stored_credentials)
def test_multistore_file_backwards_compatibility(self):
credentials = self.create_test_credentials()
scopes = ['scope1', 'scope2']
# store the credentials using the legacy key method
store = multistore_file.get_credential_storage(
FILENAME, 'client_id', 'user_agent', scopes)
store.put(credentials)
# retrieve the credentials using a custom key that matches the legacy key
key = {'clientId': 'client_id', 'userAgent': 'user_agent',
'scope': util.scopes_to_string(scopes)}
store = multistore_file.get_credential_storage_custom_key(FILENAME, key)
stored_credentials = store.get()
self.assertEqual(credentials.access_token, stored_credentials.access_token)
if __name__ == '__main__':
unittest.main()

View File

@@ -25,3 +25,20 @@ class ScopeToStringTests(unittest.TestCase):
]
for expected, case in cases:
self.assertEqual(expected, util.scopes_to_string(case))
class KeyConversionTests(unittest.TestCase):
def test_key_conversions(self):
d = {'somekey': 'some value', 'another': 'something else', 'onemore': 'foo'}
tuple_key = util.dict_to_tuple_key(d)
# the resulting key should be naturally sorted
self.assertEqual(
(('another', 'something else'),
('onemore', 'foo'),
('somekey', 'some value')),
tuple_key)
# check we get the original dictionary back
self.assertEqual(d, dict(tuple_key))