Switching flask_util and django_util to use DictionaryStorage instead…

… of custom classes.

oauth2client/contrib/django_util/storage.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from oauth2client import client
+from oauth2client.contrib.dictionary_storage import DictionaryStorage
+
+_CREDENTIALS_KEY = 'google_oauth2_credentials'
 
 
 def get_storage(request):
@@ -22,32 +24,4 @@ def get_storage(request):
     :param request: Reference to the current request object
     :return: A OAuth2Client Storage implementation based on sessions
     """
-    return DjangoSessionStorage(request.session)
-
-_CREDENTIALS_KEY = 'google_oauth2_credentials'
-
-
-class DjangoSessionStorage(client.Storage):
-    """Storage implementation that uses Django sessions."""
-
-    def __init__(self, session):
-        super(DjangoSessionStorage, self).__init__()
-        self.session = session
-
-    def locked_get(self):
-        serialized = self.session.get(_CREDENTIALS_KEY)
-
-        if serialized is None:
-            return None
-
-        credentials = client.OAuth2Credentials.from_json(serialized)
-        credentials.set_store(self)
-
-        return credentials
-
-    def locked_put(self, credentials):
-        self.session[_CREDENTIALS_KEY] = credentials.to_json()
-
-    def locked_delete(self):
-        if _CREDENTIALS_KEY in self.session:
-            del self.session[_CREDENTIALS_KEY]
+    return DictionaryStorage(request.session, key=_CREDENTIALS_KEY)

oauth2client/contrib/flask_util.py

@@ -183,9 +183,8 @@ def requires_calendar():
     raise ImportError('The flask utilities require flask 0.9 or newer.')
 
 from oauth2client.client import FlowExchangeError
-from oauth2client.client import OAuth2Credentials
 from oauth2client.client import OAuth2WebServerFlow
-from oauth2client.client import Storage
+from oauth2client.contrib.dictionary_storage import DictionaryStorage
 from oauth2client import clientsecrets
 
 
@@ -264,7 +263,7 @@ def init_app(self, app, scopes=None, client_secrets_file=None,
         self.flow_kwargs = kwargs
 
         if storage is None:
-            storage = FlaskSessionStorage()
+            storage = DictionaryStorage(session, key=_CREDENTIALS_KEY)
         self.storage = storage
 
         if scopes is None:
@@ -548,31 +547,3 @@ def http(self, *args, **kwargs):
         if not self.credentials:
             raise ValueError('No credentials available.')
         return self.credentials.authorize(httplib2.Http(*args, **kwargs))
-
-
-class FlaskSessionStorage(Storage):
-    """Storage implementation that uses Flask sessions.
-
-    Note that flask's default sessions are signed but not encrypted. Users
-    can see their own credentials and non-https connections can intercept user
-    credentials. We strongly recommend using a server-side session
-    implementation.
-    """
-
-    def locked_get(self):
-        serialized = session.get(_CREDENTIALS_KEY)
-
-        if serialized is None:
-            return None
-
-        credentials = OAuth2Credentials.from_json(serialized)
-        credentials.set_store(self)
-
-        return credentials
-
-    def locked_put(self, credentials):
-        session[_CREDENTIALS_KEY] = credentials.to_json()
-
-    def locked_delete(self):
-        if _CREDENTIALS_KEY in session:
-            del session[_CREDENTIALS_KEY]

tests/contrib/test_django_util.py

@@ -121,7 +121,7 @@ def test_view(request):
         self.assertFalse(request.oauth.has_credentials())
         self.assertIsNone(request.oauth.http)
 
-    @mock.patch("oauth2client.client.OAuth2Credentials")
+    @mock.patch('oauth2client.contrib.dictionary_storage.OAuth2Credentials')
     def test_has_credentials_in_storage(self, OAuth2Credentials):
         request = self.factory.get('/test')
         request.session = mock.MagicMock()
@@ -142,7 +142,7 @@ def test_view(request):
         self.assertTrue(request.oauth.has_credentials())
         self.assertIsNotNone(request.oauth.http)
 
-    @mock.patch("oauth2client.client.OAuth2Credentials")
+    @mock.patch('oauth2client.contrib.dictionary_storage.OAuth2Credentials')
     def test_specified_scopes(self, OAuth2Credentials):
         request = self.factory.get('/test')
         request.session = mock.MagicMock()
@@ -181,7 +181,7 @@ def test_view(request):
 
         self.assertEquals(response.status_code, 302)
 
-    @mock.patch("oauth2client.contrib.django_util.UserOAuth2", autospec=True)
+    @mock.patch('oauth2client.contrib.django_util.UserOAuth2', autospec=True)
     def test_has_credentials_in_storage(self, UserOAuth2):
         request = self.factory.get('/test')
         request.session = mock.MagicMock()
@@ -199,7 +199,7 @@ def test_view(request):
         self.assertEquals(response.status_code, 200)
         self.assertEquals(response.content, b"test")
 
-    @mock.patch("oauth2client.client.OAuth2Credentials")
+    @mock.patch('oauth2client.contrib.dictionary_storage.OAuth2Credentials')
     def test_has_credentials_in_storage_no_scopes(self, OAuth2Credentials):
         request = self.factory.get('/test')
 
@@ -217,7 +217,7 @@ def test_view(request):
         response = test_view(request)
         self.assertEquals(response.status_code, 302)
 
-    @mock.patch("oauth2client.client.OAuth2Credentials")
+    @mock.patch('oauth2client.contrib.dictionary_storage.OAuth2Credentials')
     def test_specified_scopes(self, OAuth2Credentials):
         request = self.factory.get('/test')
         request.session = mock.MagicMock()
@@ -387,14 +387,21 @@ def test_no_saved_flow(self):
         self.assertEquals(response.content, b'Missing Oauth2 flow.')
 
 
+class MockObjectWithSession(object):
+    def __init__(self, session):
+        self.session = session
+
+
 class StorageTest(TestWithSession):
 
     def test_session_delete(self):
         self.session[storage._CREDENTIALS_KEY] = "test_val"
-        django_storage = storage.DjangoSessionStorage(self.session)
+        request = MockObjectWithSession(self.session)
+        django_storage = storage.get_storage(request)
         django_storage.delete()
         self.assertIsNone(self.session.get(storage._CREDENTIALS_KEY))
 
     def test_session_delete_nothing(self):
-        django_storage = storage.DjangoSessionStorage(self.session)
+        request = MockObjectWithSession(self.session)
+        django_storage = storage.get_storage(request)
         django_storage.delete()