Switch django models from pickle to jsonpickle (#614)

3 files changed, 16 insertions, 11 deletions

diff --git a/oauth2client/contrib/django_util/views.py b/oauth2client/contrib/django_util/views.py
index 4d8ae03..4858f20 100644
--- a/oauth2client/contrib/django_util/views.py
+++ b/oauth2client/contrib/django_util/views.py
@@ -22,13 +22,13 @@ in the configured storage."""
 import hashlib
 import json
 import os
-import pickle
 
 from django import http
 from django import shortcuts
 from django.conf import settings
 from django.core import urlresolvers
 from django.shortcuts import redirect
+import jsonpickle
 from six.moves.urllib import parse
 
 from oauth2client import client
@@ -71,7 +71,7 @@ def _make_flow(request, scopes, return_url=None):
             urlresolvers.reverse("google_oauth:callback")))
 
     flow_key = _FLOW_KEY.format(csrf_token)
-    request.session[flow_key] = pickle.dumps(flow)
+    request.session[flow_key] = jsonpickle.encode(flow)
     return flow
 
 
@@ -89,7 +89,7 @@ def _get_flow_for_token(csrf_token, request):
         CSRF token.
     """
     flow_pickle = request.session.get(_FLOW_KEY.format(csrf_token), None)
-    return None if flow_pickle is None else pickle.loads(flow_pickle)
+    return None if flow_pickle is None else jsonpickle.decode(flow_pickle)
 
 
 def oauth2_callback(request):
diff --git a/tests/contrib/django_util/test_views.py b/tests/contrib/django_util/test_views.py
index 3a0b752..3962924 100644
--- a/tests/contrib/django_util/test_views.py
+++ b/tests/contrib/django_util/test_views.py
@@ -152,8 +152,8 @@ class Oauth2CallbackTest(tests_django_util.TestWithDjangoEnvironment):
         self.user = django_models.User.objects.create_user(
             username='bill', email='bill@example.com', password='hunter2')
 
-    @mock.patch('oauth2client.contrib.django_util.views.pickle')
-    def test_callback_works(self, pickle):
+    @mock.patch('oauth2client.contrib.django_util.views.jsonpickle')
+    def test_callback_works(self, jsonpickle_mock):
         request = self.factory.get('oauth2/oauth2callback', data={
             'state': json.dumps(self.fake_state),
             'code': 123
@@ -169,9 +169,10 @@ class Oauth2CallbackTest(tests_django_util.TestWithDjangoEnvironment):
             redirect_uri=request.build_absolute_uri("oauth2/oauth2callback"))
 
         name = 'google_oauth2_flow_{0}'.format(self.CSRF_TOKEN)
-        self.session[name] = pickle.dumps(flow)
+        pickled_flow = object()
+        self.session[name] = pickled_flow
         flow.step2_exchange = mock.Mock()
-        pickle.loads.return_value = flow
+        jsonpickle_mock.decode.return_value = flow
 
         request.session = self.session
         request.user = self.user
@@ -180,9 +181,10 @@ class Oauth2CallbackTest(tests_django_util.TestWithDjangoEnvironment):
         self.assertEqual(
             response.status_code, django.http.HttpResponseRedirect.status_code)
         self.assertEqual(response['Location'], self.RETURN_URL)
+        jsonpickle_mock.decode.assert_called_once_with(pickled_flow)
 
-    @mock.patch('oauth2client.contrib.django_util.views.pickle')
-    def test_callback_handles_bad_flow_exchange(self, pickle):
+    @mock.patch('oauth2client.contrib.django_util.views.jsonpickle')
+    def test_callback_handles_bad_flow_exchange(self, jsonpickle_mock):
         request = self.factory.get('oauth2/oauth2callback', data={
             "state": json.dumps(self.fake_state),
             "code": 123
@@ -198,17 +200,19 @@ class Oauth2CallbackTest(tests_django_util.TestWithDjangoEnvironment):
             redirect_uri=request.build_absolute_uri('oauth2/oauth2callback'))
 
         session_key = 'google_oauth2_flow_{0}'.format(self.CSRF_TOKEN)
-        self.session[session_key] = pickle.dumps(flow)
+        pickled_flow = object()
+        self.session[session_key] = pickled_flow
 
         def local_throws(code):
             raise client.FlowExchangeError('test')
 
         flow.step2_exchange = local_throws
-        pickle.loads.return_value = flow
+        jsonpickle_mock.decode.return_value = flow
 
         request.session = self.session
         response = views.oauth2_callback(request)
         self.assertIsInstance(response, http.HttpResponseBadRequest)
+        jsonpickle_mock.decode.assert_called_once_with(pickled_flow)
 
     def test_error_returns_bad_request(self):
         request = self.factory.get('oauth2/oauth2callback', data={
diff --git a/tox.ini b/tox.ini
index 6da2719..9cf31b1 100644
--- a/tox.ini
+++ b/tox.ini
@@ -14,6 +14,7 @@ basedeps = mock>=1.3.0
 deps = {[testenv]basedeps}
        django
        keyring
+       jsonpickle
 setenv =
     pypy: with_gmp=no
     DJANGO_SETTINGS_MODULE=tests.contrib.django_util.settings