From 49d3789f984068d2c2a813507fd07cf6d92f3402 Mon Sep 17 00:00:00 2001
From: Matt Wright <mdw1980@gmail.com>
Date: Tue, 10 Jul 2012 16:23:20 -0400
Subject: [PATCH] Fix up remember token

---
 flask_security/core.py      | 1 +
 flask_security/datastore.py | 5 ++++-
 flask_security/utils.py     | 7 +++++++
 flask_security/views.py     | 8 ++++----
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/flask_security/core.py b/flask_security/core.py
index 922ed12..f3626b9 100644
--- a/flask_security/core.py
+++ b/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:
diff --git a/flask_security/datastore.py b/flask_security/datastore.py
index 90d214a..4746102 100644
--- a/flask_security/datastore.py
+++ b/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):
diff --git a/flask_security/utils.py b/flask_security/utils.py
index 77abb1f..7d51ad5 100644
--- a/flask_security/utils.py
+++ b/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.
diff --git a/flask_security/views.py b/flask_security/views.py
index 890b980..423be77 100644
--- a/flask_security/views.py
+++ b/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()