From 1ae6bc3cf133887771ed5523ea295e859cd65c59 Mon Sep 17 00:00:00 2001 From: Matt Wright Date: Wed, 16 Oct 2013 14:00:36 -0400 Subject: [PATCH] Add the ability to specify additional fields on the user model that can be used for logging in. --- flask_security/core.py | 3 ++- flask_security/datastore.py | 46 +++++++++++++++++++++-------------- flask_security/forms.py | 1 + flask_security/utils.py | 10 ++++++++ tests/configured_tests.py | 11 +++++++++ tests/test_app/__init__.py | 16 ++++++------ tests/test_app/mongoengine.py | 1 + tests/test_app/peewee_app.py | 1 + tests/test_app/sqlalchemy.py | 1 + 9 files changed, 63 insertions(+), 27 deletions(-) diff --git a/flask_security/core.py b/flask_security/core.py index 9ec4284..00b7c41 100644 --- a/flask_security/core.py +++ b/flask_security/core.py @@ -84,7 +84,8 @@ _default_config = { 'EMAIL_SUBJECT_PASSWORDLESS': 'Login instructions', 'EMAIL_SUBJECT_PASSWORD_NOTICE': 'Your password has been reset', 'EMAIL_SUBJECT_PASSWORD_CHANGE_NOTICE': 'Your password has been changed', - 'EMAIL_SUBJECT_PASSWORD_RESET': 'Password reset instructions' + 'EMAIL_SUBJECT_PASSWORD_RESET': 'Password reset instructions', + 'USER_IDENTITY_ATTRIBUTES': ['email'] } #: Default Flask-Security messages diff --git a/flask_security/datastore.py b/flask_security/datastore.py index 4716931..7342d41 100644 --- a/flask_security/datastore.py +++ b/flask_security/datastore.py @@ -9,6 +9,8 @@ :license: MIT, see LICENSE for more details. """ +from .utils import get_identity_attributes + class Datastore(object): def __init__(self, db): @@ -179,14 +181,14 @@ class SQLAlchemyUserDatastore(SQLAlchemyDatastore, UserDatastore): SQLAlchemyDatastore.__init__(self, db) UserDatastore.__init__(self, user_model, role_model) - def get_user(self, id_or_email): - returned = None - if self._is_numeric(id_or_email): - returned = self.user_model.query.get(id_or_email) - if not returned: - returned = self.user_model.query.filter( - self.user_model.email.ilike(id_or_email)).first() - return returned + def get_user(self, identifier): + if self._is_numeric(identifier): + return self.user_model.query.get(identifier) + for attr in get_identity_attributes(): + query = getattr(self.user_model, attr).ilike(identifier) + rv = self.user_model.query.filter(query).first() + if rv is not None: + return rv def _is_numeric(self, value): try: @@ -210,12 +212,18 @@ class MongoEngineUserDatastore(MongoEngineDatastore, UserDatastore): MongoEngineDatastore.__init__(self, db) UserDatastore.__init__(self, user_model, role_model) - def get_user(self, id_or_email): + def get_user(self, identifier): from mongoengine import ValidationError try: - return self.user_model.objects(id=id_or_email).first() + return self.user_model.objects(id=identifier).first() except ValidationError: - return self.user_model.objects(email__iexact=id_or_email).first() + pass + for attr in get_identity_attributes(): + query_key = '%s__iexact' % attr + query = {query_key: identifier} + rv = self.user_model.objects(**query).first() + if rv is not None: + return rv def find_user(self, **kwargs): try: @@ -254,16 +262,18 @@ class PeeweeUserDatastore(PeeweeDatastore, UserDatastore): UserDatastore.__init__(self, user_model, role_model) self.UserRole = role_link - def get_user(self, id_or_email): + def get_user(self, identifier): try: - return self.user_model.get(self.user_model.id == id_or_email) + return self.user_model.get(self.user_model.id == identifier) except ValueError: pass - try: - return self.user_model.get(self.user_model.email ** id_or_email) - except self.user_model.DoesNotExist: - pass - return None + + for attr in get_identity_attributes(): + column = getattr(self.user_model, attr) + try: + return self.user_model.get(column ** identifier) + except self.user_model.DoesNotExist: + pass def find_user(self, **kwargs): try: diff --git a/flask_security/forms.py b/flask_security/forms.py index 57ebb3e..debdf40 100644 --- a/flask_security/forms.py +++ b/flask_security/forms.py @@ -220,6 +220,7 @@ class LoginForm(Form, NextFormMixin): self.password.errors.append(get_message('PASSWORD_NOT_PROVIDED')[0]) return False + self.user = _datastore.get_user(self.email.data) if self.user is None: diff --git a/flask_security/utils.py b/flask_security/utils.py index 06effff..1f7fa14 100644 --- a/flask_security/utils.py +++ b/flask_security/utils.py @@ -279,6 +279,16 @@ def get_token_status(token, serializer, max_age=None): return expired, invalid, user +def get_identity_attributes(app=None): + app = app or current_app + attrs = app.config['SECURITY_USER_IDENTITY_ATTRIBUTES'] + try: + attrs = [f.strip() for f in attrs.split(',')] + except AttributeError: + pass + return attrs + + @contextmanager def capture_passwordless_login_requests(): login_requests = [] diff --git a/tests/configured_tests.py b/tests/configured_tests.py index 0cc9031..6be025e 100644 --- a/tests/configured_tests.py +++ b/tests/configured_tests.py @@ -826,3 +826,14 @@ class ConfirmableExtendFormsTest(SecurityTest): def test_send_confirmation(self): r = self._get('/confirm', follow_redirects=True) self.assertIn("My Send Confirmation Email Address Field", r.data) + + +class AdditionalUserIdentityAttributes(SecurityTest): + + AUTH_CONFIG = { + 'SECURITY_USER_IDENTITY_ATTRIBUTES': ('email', 'username') + } + + def test_authenticate(self): + r = self.authenticate(email='matt') + self.assertIn('Hello matt@lp.com', r.data) diff --git a/tests/test_app/__init__.py b/tests/test_app/__init__.py index 379e4ec..3d0c798 100644 --- a/tests/test_app/__init__.py +++ b/tests/test_app/__init__.py @@ -128,17 +128,17 @@ def create_roles(): def create_users(count=None): - users = [('matt@lp.com', 'password', ['admin'], True), - ('joe@lp.com', 'password', ['editor'], True), - ('dave@lp.com', 'password', ['admin', 'editor'], True), - ('jill@lp.com', 'password', ['author'], True), - ('tiya@lp.com', 'password', [], False)] + users = [('matt@lp.com', 'matt', 'password', ['admin'], True), + ('joe@lp.com', 'joe', 'password', ['editor'], True), + ('dave@lp.com', 'dave', 'password', ['admin', 'editor'], True), + ('jill@lp.com', 'jill', 'password', ['author'], True), + ('tiya@lp.com', 'tiya', 'password', [], False)] count = count or len(users) for u in users[:count]: - pw = encrypt_password(u[1]) - ds.create_user(email=u[0], password=pw, - roles=u[2], active=u[3]) + pw = encrypt_password(u[2]) + ds.create_user(email=u[0], username=u[1], password=pw, + roles=u[3], active=u[4]) ds.commit() diff --git a/tests/test_app/mongoengine.py b/tests/test_app/mongoengine.py index d40f666..9be6f5f 100644 --- a/tests/test_app/mongoengine.py +++ b/tests/test_app/mongoengine.py @@ -30,6 +30,7 @@ def create_app(config, **kwargs): class User(db.Document, UserMixin): email = db.StringField(unique=True, max_length=255) + username = db.StringField(max_length=255) password = db.StringField(required=True, max_length=255) last_login_at = db.DateTimeField() current_login_at = db.DateTimeField() diff --git a/tests/test_app/peewee_app.py b/tests/test_app/peewee_app.py index 8b4fc05..b33ab29 100644 --- a/tests/test_app/peewee_app.py +++ b/tests/test_app/peewee_app.py @@ -29,6 +29,7 @@ def create_app(config, **kwargs): class User(db.Model, UserMixin): email = TextField() + username = TextField() password = TextField() last_login_at = DateTimeField(null=True) current_login_at = DateTimeField(null=True) diff --git a/tests/test_app/sqlalchemy.py b/tests/test_app/sqlalchemy.py index d1af62b..fcce996 100644 --- a/tests/test_app/sqlalchemy.py +++ b/tests/test_app/sqlalchemy.py @@ -32,6 +32,7 @@ def create_app(config, **kwargs): class User(db.Model, UserMixin): id = db.Column(db.Integer, primary_key=True) email = db.Column(db.String(255), unique=True) + username = db.Column(db.String(255)) password = db.Column(db.String(255)) last_login_at = db.Column(db.DateTime()) current_login_at = db.Column(db.DateTime())