From 4e7663079bbac333a355c6b5624c3811dba76231 Mon Sep 17 00:00:00 2001
From: Matt Wright <matt@localprojects.net>
Date: Mon, 14 May 2012 18:58:43 -0400
Subject: [PATCH] Refactor some methods

---
 example/app.py                |  3 ++-
 flask_security/confirmable.py | 25 ++++++++++++++++++-------
 flask_security/datastore.py   |  1 -
 flask_security/exceptions.py  |  5 +++++
 flask_security/views.py       | 17 ++++++++++++-----
 5 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/example/app.py b/example/app.py
index 02c6810..6eb33c1 100644
--- a/example/app.py
+++ b/example/app.py
@@ -130,7 +130,7 @@ def create_sqlalchemy_app(auth_config=None):
 
 def create_mongoengine_app(auth_config=None):
     app = create_app(auth_config)
-    app.config['MONGODB_DB'] = 'flask_security_example'
+    app.config['MONGODB_DB'] = 'flask_security_test'
     app.config['MONGODB_HOST'] = 'localhost'
     app.config['MONGODB_PORT'] = 27017
 
@@ -146,6 +146,7 @@ def create_mongoengine_app(auth_config=None):
         active = db.BooleanField(default=True)
         confirmation_token = db.StringField(max_length=255)
         confirmation_sent_at = db.DateTimeField()
+        confirmed_at = db.DateTimeField()
         roles = db.ListField(db.ReferenceField(Role), default=[])
 
     Security(app, MongoEngineUserDatastore(db, User, Role))
diff --git a/flask_security/confirmable.py b/flask_security/confirmable.py
index 3253da2..b52de60 100644
--- a/flask_security/confirmable.py
+++ b/flask_security/confirmable.py
@@ -50,18 +50,29 @@ def generate_confirmation_token(user):
     return user
 
 
+def requires_confirmation(user):
+    return (security.confirm_email and \
+            not security.login_without_confirmation and \
+            not confirmation_token_is_expired(user))
+
+
+def confirmation_token_is_expired(user):
+    token_expires = datetime.utcnow() - security.confirm_email_within
+    if user.confirmation_sent_at < token_expires:
+        return True
+    return False
+
+
 def confirm_by_token(token):
-    now = datetime.utcnow()
     user = find_user_by_confirmation_token(token)
 
-    token_expires = now - security.confirm_email_within
-
-    if user.confirmation_sent_at < token_expires:
+    if confirmation_token_is_expired(user):
         raise ConfirmationExpiredError('Confirmation token is expired', user=user)
 
-    user.confirmed_at = now
-    user.confirmation_token = None
-    user.confirmation_sent_at = None
+    user.confirmed_at = datetime.utcnow()
+    #user.confirmation_token = None
+    #user.confirmation_sent_at = None
+
     security.datastore._save_model(user)
 
     return user
diff --git a/flask_security/datastore.py b/flask_security/datastore.py
index 6bab2e0..1c3d3c6 100644
--- a/flask_security/datastore.py
+++ b/flask_security/datastore.py
@@ -91,7 +91,6 @@ class UserDatastore(object):
 
         :param user: User identifier, usually email address
         """
-        print kwargs
         user = self._do_find_user(**kwargs)
         if user:
             return user
diff --git a/flask_security/exceptions.py b/flask_security/exceptions.py
index 996fa3b..6ef84f8 100644
--- a/flask_security/exceptions.py
+++ b/flask_security/exceptions.py
@@ -67,3 +67,8 @@ class ConfirmationExpiredError(Exception):
     def __init__(self, msg, user=None):
         super(ConfirmationExpiredError, self).__init__(msg)
         self.user = user
+
+
+class ConfirmationRequiredError(Exception):
+    """Raised when a user attempts to login but requires confirmation
+    """
diff --git a/flask_security/views.py b/flask_security/views.py
index 7d09bb3..56ce9a4 100644
--- a/flask_security/views.py
+++ b/flask_security/views.py
@@ -22,6 +22,9 @@ logger = LocalProxy(lambda: current_app.logger)
 
 
 def do_login(user, remember=True):
+    if confirmable.requires_confirmation(user):
+        raise exceptions.ConfirmationRequiredError()
+
     if login_user(user, remember):
         identity_changed.send(current_app._get_current_object(),
                               identity=Identity(user.id))
@@ -41,15 +44,19 @@ def authenticate():
 
         raise exceptions.BadCredentialsError('Inactive user')
 
+    except exceptions.ConfirmationRequiredError, e:
+        msg = str(e)
+
     except exceptions.BadCredentialsError, e:
         msg = str(e)
-        utils.do_flash(msg, 'error')
-        url = request.referrer or security.login_manager.login_view
 
-        logger.debug('Unsuccessful authentication attempt: %s. '
-                     'Redirect to: %s' % (msg, url))
+    utils.do_flash(msg, 'error')
+    url = request.referrer or security.login_manager.login_view
 
-        return redirect(url)
+    logger.debug('Unsuccessful authentication attempt: %s. '
+                 'Redirect to: %s' % (msg, url))
+
+    return redirect(url)
 
 
 def logout():