From ebe807b87b6fbf2b195791019513c695dd1d38ef Mon Sep 17 00:00:00 2001 From: Matt Wright Date: Thu, 8 Mar 2012 17:06:23 -0500 Subject: [PATCH] MongoEngine datastore completed. Also refactored model classes to use a mixin --- example/app.py | 95 +++++++++++++++++-------- flask_security/__init__.py | 29 +++++++- flask_security/datastore/mongoengine.py | 68 +++++++++++++++++- flask_security/datastore/sqlalchemy.py | 25 +------ tests/functional_tests.py | 20 ++++-- 5 files changed, 174 insertions(+), 63 deletions(-) diff --git a/example/app.py b/example/app.py index f2c036c..38e66a8 100644 --- a/example/app.py +++ b/example/app.py @@ -6,23 +6,17 @@ sys.path.pop(0) sys.path.insert(0, os.getcwd()) from flask import Flask, render_template -from flask.ext.sqlalchemy import SQLAlchemy -from flask.ext.security import (Security, LoginForm, user_datastore, - login_required, roles_required, roles_accepted) -from flask.ext.security.datastore.sqlalchemy import SQLAlchemyUserDatastore -def create_app(auth_config=None): - - app = Flask(__name__) - app.debug = True - app.config['SECRET_KEY'] = 'secret' - app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:' - app.config['AUTH'] = auth_config or {} - - db = SQLAlchemy(app) - - Security(app, SQLAlchemyUserDatastore(db)) - +from flask.ext.mongoengine import MongoEngine +from flask.ext.sqlalchemy import SQLAlchemy + +from flask.ext.security import (Security, LoginForm, user_datastore, + login_required, roles_required, roles_accepted) + +from flask.ext.security.datastore.sqlalchemy import SQLAlchemyUserDatastore +from flask.ext.security.datastore.mongoengine import MongoEngineUserDatastore + +def add_endpoints(app): @app.route('/') def index(): return render_template('index.html', content='Home Page') @@ -59,27 +53,66 @@ def create_app(auth_config=None): def admin_or_editor(): return render_template('index.html', content='Admin or Editor Page') +def create_users(): + user_datastore.create_user(username='matt', email='matt@lp.com', + password='password', + roles=['admin']) + + user_datastore.create_user(username='joe', email='joe@lp.com', + password='password', + roles=['editor']) + + user_datastore.create_user(username='jill', email='jill@lp.com', + password='password', + roles=['author']) + + user_datastore.create_user(username='tiya', email='tiya@lp.com', + password='password', active=False) + +def create_sqlalchemy_app(auth_config=None): + app = Flask(__name__) + app.debug = True + app.config['SECRET_KEY'] = 'secret' + app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:' + app.config['AUTH'] = auth_config or {} + + db = SQLAlchemy(app) + + Security(app, SQLAlchemyUserDatastore(db)) + + add_endpoints(app) + @app.before_first_request def before_first_request(): db.create_all() + create_users() - user_datastore.create_user(username='matt', email='matt@lp.com', - password='password', - roles=['admin']) - - user_datastore.create_user(username='joe', email='joe@lp.com', - password='password', - roles=['editor']) - - user_datastore.create_user(username='jill', email='jill@lp.com', - password='password', - roles=['author']) - - user_datastore.create_user(username='tiya', email='tiya@lp.com', - password='password', active=False) + return app + +def create_mongoengine_app(auth_config=None): + app = Flask(__name__) + app.debug = True + app.config['SECRET_KEY'] = 'secret' + app.config['AUTH'] = auth_config or {} + app.config['MONGODB_DB'] = 'flask_security_example' + app.config['MONGODB_HOST'] = 'localhost' + app.config['MONGODB_PORT'] = 27017 + + db = MongoEngine(app) + + Security(app, MongoEngineUserDatastore(db)) + + add_endpoints(app) + + @app.before_first_request + def before_first_request(): + from flask.ext.security import User, Role + User.drop_collection() + Role.drop_collection() + create_users() return app if __name__ == '__main__': - app = create_app() + app = create_sqlalchemy_app() app.run() \ No newline at end of file diff --git a/flask_security/__init__.py b/flask_security/__init__.py index 1f6666f..5638825 100644 --- a/flask_security/__init__.py +++ b/flask_security/__init__.py @@ -14,13 +14,14 @@ from __future__ import absolute_import import sys from datetime import datetime +from types import StringType from flask import (current_app, Blueprint, flash, redirect, request, session, _request_ctx_stack, url_for, abort, g) -from flask.ext.login import (AnonymousUser as AnonymousUserBase, UserMixin, - LoginManager, login_required, login_user, logout_user, - current_user, user_logged_in, user_logged_out) +from flask.ext.login import (AnonymousUser as AnonymousUserBase, + UserMixin as BaseUserMixin, LoginManager, login_required, login_user, + logout_user, current_user, user_logged_in, user_logged_out) from flask.ext.principal import (Identity, Principal, RoleNeed, UserNeed, Permission, AnonymousIdentity, identity_changed, identity_loaded) @@ -140,6 +141,28 @@ def roles_accepted(*args): return decorated_view return wrapper +class RoleMixin(object): + def __eq__(self, other): + return self.name == other.name + + def __ne__(self, other): + return self.name != other.name + + def __str__(self): + return '' % (self.name, self.description) + +class UserMixin(BaseUserMixin): + def is_active(self): + return self.active + + def has_role(self, role): + if type(role) is StringType: + role = Role(name=role) + return role in self.roles + + def __str__(self): + ctx = (str(self.id), self.username, self.email) + return '' % ctx class AnonymousUser(AnonymousUserBase): def __init__(self): diff --git a/flask_security/datastore/mongoengine.py b/flask_security/datastore/mongoengine.py index 2f6b29e..dc61a18 100644 --- a/flask_security/datastore/mongoengine.py +++ b/flask_security/datastore/mongoengine.py @@ -1 +1,67 @@ -# TODO: Add mongoengine datastore \ No newline at end of file +from flask.ext import security +from flask.ext.security import UserMixin, RoleMixin +from flask.ext.security.datastore import UserDatastore + +class MongoEngineUserDatastore(UserDatastore): + """MongoEngine datastore""" + + def __init__(self, db): + self.db = db + + class Role(db.Document, RoleMixin): + name = db.StringField(required=True, max_length=80) + description = db.StringField(max_length=255) + + class User(db.Document, UserMixin): + username = db.StringField(max_length=255) + email = db.StringField(max_length=255) + password = db.StringField(max_length=120) + active = db.BooleanField(default=True) + roles= db.ListField(db.ReferenceField(Role), default=[]) + created_at = db.DateTimeField() + modified_at = db.DateTimeField() + + security.User = User + security.Role = Role + + def with_id(self, id): + try: return security.User.objects.get(id=id) + except: raise security.UserIdNotFoundError() + + def find(self, user_identifier): + user = security.User.objects(username=user_identifier).first() + if user: return user + user = security.User.objects(email=user_identifier).first() + if user: return user + raise security.UserNotFoundError() + + def create_role(self, **kwargs): + if not kwargs.has_key('name'): + raise TypeError("create_role() did not receive " + "keyword argument 'name'") + + name = kwargs.get('name') + description = kwargs.get('description', None) + + role = security.Role.objects(name=name).first() + + if role is None: + role = security.Role(name=name, description=description) + role.save() + + return role + + def create_user(self, **kwargs): + kwargs = self._prepare_create_args(kwargs) + + roles = kwargs.get('roles', []) + user_roles = [] + for role in roles: + user_roles.append(self.create_role(name=role)) + + kwargs['roles'] = user_roles + + user = security.User(**kwargs) + user.save() + + return user \ No newline at end of file diff --git a/flask_security/datastore/sqlalchemy.py b/flask_security/datastore/sqlalchemy.py index f7747a6..1b5d1a1 100644 --- a/flask_security/datastore/sqlalchemy.py +++ b/flask_security/datastore/sqlalchemy.py @@ -1,6 +1,5 @@ -from types import StringType from flask.ext import security -from flask.ext.login import UserMixin +from flask.ext.security import UserMixin, RoleMixin from flask.ext.security.datastore import UserDatastore class SQLAlchemyUserDatastore(UserDatastore): @@ -13,7 +12,7 @@ class SQLAlchemyUserDatastore(UserDatastore): db.Column('user_id', db.Integer(), db.ForeignKey('role.id')), db.Column('role_id', db.Integer(), db.ForeignKey('user.id'))) - class Role(db.Model): + class Role(db.Model, RoleMixin): id = db.Column(db.Integer(), primary_key=True) name = db.Column(db.String(80), unique=True) description = db.Column(db.String(255)) @@ -21,15 +20,6 @@ class SQLAlchemyUserDatastore(UserDatastore): def __init__(self, name=None, description=None): self.name = name self.description = description - - def __eq__(self, other): - return self.name == other.name - - def __ne__(self, other): - return self.name != other.name - - def __str__(self): - return '' % (self.name, self.description) class User(db.Model, UserMixin): id = db.Column(db.Integer, primary_key=True) @@ -53,17 +43,6 @@ class SQLAlchemyUserDatastore(UserDatastore): self.roles = roles or [] self.created_at = created_at self.modified_at = modified_at - - def is_active(self): - return self.active - - def has_role(self, role): - if type(role) is StringType: - role = security.Role(name=role) - return role in self.roles - - def __str__(self): - return '' % self.__dict__ security.User = User security.Role = Role diff --git a/tests/functional_tests.py b/tests/functional_tests.py index 9f88daf..50ffa5d 100644 --- a/tests/functional_tests.py +++ b/tests/functional_tests.py @@ -1,19 +1,22 @@ import unittest from example import app -class SecurityTests(unittest.TestCase): +class SecurityTest(unittest.TestCase): AUTH_CONFIG = None def setUp(self): - super(SecurityTests, self).setUp() + super(SecurityTest, self).setUp() - self.app = app.create_app(self.AUTH_CONFIG or None) + self.app = self._create_app(self.AUTH_CONFIG or None) self.app.debug = False self.app.config['TESTING'] = True self.client = self.app.test_client() + def _create_app(self, auth_config): + return app.create_sqlalchemy_app(auth_config) + def _get(self, route, content_type=None, follow_redirects=None): return self.client.get(route, follow_redirects=follow_redirects, content_type=content_type or 'text/html') @@ -22,6 +25,7 @@ class SecurityTests(unittest.TestCase): return self.client.post(route, data=data, follow_redirects=follow_redirects, content_type=content_type or 'text/html') + def authenticate(self, username, password, endpoint=None): data = dict(username=username, password=password) @@ -31,7 +35,7 @@ class SecurityTests(unittest.TestCase): def logout(self, endpoint=None): return self._get(endpoint or '/logout', follow_redirects=True) -class DefaultSecurityTests(SecurityTests): +class DefaultSecurityTests(SecurityTest): def test_login_view(self): r = self._get('/login') @@ -97,7 +101,7 @@ class DefaultSecurityTests(SecurityTests): self.assertIn('Login Page', r.data) -class ConfiguredSecurityTests(SecurityTests): +class ConfiguredSecurityTests(SecurityTest): AUTH_CONFIG = { 'password_hash': 'bcrypt', @@ -121,3 +125,9 @@ class ConfiguredSecurityTests(SecurityTests): self.authenticate("matt", "password", endpoint="/custom_auth") r = self.logout(endpoint="/custom_logout") assert 'Post Logout' in r.data + + +class MongoEngineSecurityTests(DefaultSecurityTests): + + def _create_app(self, auth_config): + return app.create_mongoengine_app(auth_config)