Fix up remember token

flask_security/core.py

@@ -113,6 +113,7 @@ def _user_loader(user_id):
 
 
 def _token_loader(token):
+    print 'token loader!'
     try:
         return current_app.security.datastore.find_user(remember_token=token)
     except Exception, e:

flask_security/datastore.py

@@ -11,7 +11,7 @@
 
 from flask import current_app
 
-from . import exceptions, confirmable
+from . import exceptions, confirmable, utils
 
 
 class UserDatastore(object):
@@ -113,6 +113,9 @@ class UserDatastore(object):
         if not pwd_context.identify(pw):
             kwargs['password'] = pwd_context.encrypt(pw)
 
+        kwargs['remember_token'] = utils.get_remember_token(kwargs['email'],
+                                                            kwargs['password'])
+
         return kwargs
 
     def with_id(self, id):

flask_security/utils.py

@@ -15,6 +15,7 @@ from contextlib import contextmanager
 from importlib import import_module
 
 from flask import url_for, flash, current_app, request, session, render_template
+from flask.ext.login import make_secure_token
 
 from .signals import user_registered, password_reset_requested
 
@@ -24,6 +25,12 @@ def generate_token():
     return base64.urlsafe_b64encode(os.urandom(30))
 
 
+def get_remember_token(email, password):
+    assert email is not None
+    assert password is not None
+    return make_secure_token(email, password)
+
+
 def do_flash(message, category=None):
     """Flash a message depending on if the `FLASH_MESSAGES` configuration
     value is set.

flask_security/views.py

@@ -13,7 +13,7 @@ from datetime import datetime
 
 from flask import current_app as app, redirect, request, session, \
      render_template
-from flask.ext.login import login_user, logout_user, make_secure_token
+from flask.ext.login import login_user, logout_user
 from flask.ext.principal import Identity, AnonymousIdentity, identity_changed
 from werkzeug.local import LocalProxy
 
@@ -26,7 +26,7 @@ from .forms import LoginForm, RegisterForm, ForgotPasswordForm, \
 from .recoverable import reset_by_token, \
      reset_password_reset_token
 from .signals import user_registered
-from .utils import get_post_login_redirect, do_flash
+from .utils import get_post_login_redirect, do_flash, get_remember_token
 
 
 # Convenient references
@@ -43,8 +43,8 @@ def _do_login(user, remember=True):
     if not login_user(user, remember):
         return False
 
-    user.remember_token = None if not remember else \
-                          make_secure_token(user.email, user.password)
+    if remember:
+        user.remember_token = get_remember_token(user.email, user.password)
 
     if _security.trackable:
         old_current, new_current = user.current_login_at, datetime.utcnow()