From e239216d9247be2481d6bc0639e0285b39d364e6 Mon Sep 17 00:00:00 2001 From: Matt Wright Date: Fri, 9 Mar 2012 16:31:05 -0500 Subject: [PATCH] Added a few handy commands for Flask-Script to create users and add/remove roles. Also refactored the datastore a bit with appropriate methods --- .gitignore | 1 + example/app.py | 35 +++---- example/manage.py | 19 ++++ flask_security/__init__.py | 18 +++- flask_security/datastore/__init__.py | 117 +++++++++++++++++++----- flask_security/datastore/mongoengine.py | 54 ++++------- flask_security/datastore/sqlalchemy.py | 74 +++++++-------- flask_security/script.py | 70 ++++++++++++++ setup.py | 14 ++- 9 files changed, 280 insertions(+), 122 deletions(-) create mode 100644 example/manage.py create mode 100644 flask_security/script.py diff --git a/.gitignore b/.gitignore index e741749..8af5f8c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .DS_Store *.pyc +*.egg .project .pydevproject .settings diff --git a/example/app.py b/example/app.py index 9d0c914..a1df2d2 100644 --- a/example/app.py +++ b/example/app.py @@ -16,21 +16,21 @@ from flask.ext.security import (Security, LoginForm, user_datastore, from flask.ext.security.datastore.sqlalchemy import SQLAlchemyUserDatastore from flask.ext.security.datastore.mongoengine import MongoEngineUserDatastore +def create_roles(): + for role in ('admin', 'editor', 'author'): + user_datastore.create_role(name=role) + 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) + for u in (('matt','matt@lp.com','password',['admin'],True), + ('joe','joe@lp.com','password',['editor'],True), + ('jill','jill@lp.com','password',['author'],True), + ('tiya','tiya@lp.com','password',[],False)): + user_datastore.create_user(username=u[0], email=u[1], password=u[2], + roles=u[3], active=u[4]) + +def populate_data(): + create_roles() + create_users() def create_app(auth_config): app = Flask(__name__) @@ -81,15 +81,16 @@ def create_app(auth_config): def create_sqlalchemy_app(auth_config=None): app = create_app(auth_config) - app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:' + app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:////tmp/flask_security_example.sqlite' db = SQLAlchemy(app) Security(app, SQLAlchemyUserDatastore(db)) @app.before_first_request def before_first_request(): + db.drop_all() db.create_all() - create_users() + populate_data() return app @@ -107,7 +108,7 @@ def create_mongoengine_app(auth_config=None): from flask.ext.security import User, Role User.drop_collection() Role.drop_collection() - create_users() + populate_data() return app diff --git a/example/manage.py b/example/manage.py new file mode 100644 index 0000000..b36e732 --- /dev/null +++ b/example/manage.py @@ -0,0 +1,19 @@ +# a little trick so you can run: +# $ python example/app.py +# from the root of the security project +import sys, os +sys.path.pop(0) +sys.path.insert(0, os.getcwd()) + +from example import app +from flask.ext.script import Manager +from flask.ext.security.script import (CreateUserCommand , AddRoleCommand, + RemoveRoleCommand) + +manager = Manager(app.create_sqlalchemy_app()) +manager.add_command('create_user', CreateUserCommand()) +manager.add_command('add_role', AddRoleCommand()) +manager.add_command('remove_role', RemoveRoleCommand()) + +if __name__ == "__main__": + manager.run() \ No newline at end of file diff --git a/flask_security/__init__.py b/flask_security/__init__.py index 722e706..7174e09 100644 --- a/flask_security/__init__.py +++ b/flask_security/__init__.py @@ -81,6 +81,11 @@ class UserNotFoundError(Exception): their identifier, often username or email, and the user is not found. """ +class RoleNotFoundError(Exception): + """Raised by a user datastore when there is an attempt to find a role and + the role cannot be found. + """ + class UserIdNotFoundError(Exception): """Raised by a user datastore when there is an attempt to find a user by ID and the user is not found. @@ -91,9 +96,15 @@ class UserDatastoreError(Exception): """ class UserCreationError(Exception): - """Raise when an error occurs during user create + """Raise when an error occurs when creating a user """ +class RoleCreationError(Exception): + """Raise when an error occurs when creating a role + """ + + + #: App logger for convenience logger = LocalProxy(lambda: current_app.logger) @@ -222,6 +233,9 @@ class Security(object): app.auth_provider = Provider(Form) app.principal = Principal(app) + from flask.ext import security as s + s.User, s.Role = datastore.get_models() + setattr(app, config[USER_DATASTORE_KEY], datastore) @identity_loaded.connect_via(app) @@ -319,7 +333,7 @@ class AuthenticationProvider(object): def do_authenticate(self, user_identifier, password): try: - user = user_datastore.find(user_identifier) + user = user_datastore.find_user(user_identifier) except AttributeError, e: self.auth_error("Could not find user service: %s" % e) except UserNotFoundError, e: diff --git a/flask_security/datastore/__init__.py b/flask_security/datastore/__init__.py index 9d056b5..cf8dd2f 100644 --- a/flask_security/datastore/__init__.py +++ b/flask_security/datastore/__init__.py @@ -1,15 +1,96 @@ from datetime import datetime -from flask.ext.security import UserCreationError, pwd_context +from flask.ext import security +from flask.ext.security import UserCreationError, RoleCreationError, pwd_context class UserDatastore(object): - """A sort of abstract user service""" - def with_id(self, id): - raise NotImplementedError( - "User datastore does not implement with_id method") + """Abstracted user datastore. Always extend this and implement + missing methods""" - def find(self, user_identifier): + def _do_with_id(self, id): raise NotImplementedError( - "User datastore does not implement find_user method") + "User datastore does not implement _do_with_id method") + + def _do_find_user(self): + raise NotImplementedError( + "User datastore does not implement _do_find_user method") + + def _do_find_role(self): + raise NotImplementedError( + "User datastore does not implement _do_find_role method") + + def _do_add_role(self, user, role): + user, role = self._prepare_role_modify_args(user, role) + if role not in user.roles: + user.roles.append(role) + return user + + def _do_remove_role(self, user, role): + user, role = self._prepare_role_modify_args(user, role) + if role in user.roles: + user.roles.remove(role) + return user + + def _prepare_role_modify_args(self, user, role): + if isinstance(user, security.User): + user = user.username or user.email + + if isinstance(role, security.Role): + role = role.name + + return self.find_user(user), self.find_role(role) + + def _prepare_create_role_args(self, kwargs): + for key in ('name', 'description'): + kwargs[key] = kwargs.get(key, None) + + if kwargs['name'] is None: + raise RoleCreationError("Missing name argument") + + return kwargs + + def _prepare_create_user_args(self, kwargs): + username = kwargs.get('username', None) + email = kwargs.get('email', None) + password = kwargs.get('password', None) + + if username is None and email is None: + raise UserCreationError('Missing username and/or email arguments') + + if password is None: + raise UserCreationError('Missing password argument') + + roles = kwargs.get('roles', []) + + for i, role in enumerate(roles): + rn = role.name if isinstance(role, security.Role) else role + # see if the role exists + roles[i] = self.find_role(rn) + + kwargs['roles'] = roles + + now = datetime.utcnow() + kwargs['created_at'], kwargs['modified_at'] = now, now + + pw = kwargs['password'] + if not pwd_context.identify(pw): + kwargs['password'] = pwd_context.encrypt(pw) + + return kwargs + + def with_id(self, id): + user = self._do_with_id(id) + if user: return user + raise security.UserIdNotFoundError() + + def find_user(self, user): + user = self._do_find_user(user) + if user: return user + raise security.UserNotFoundError() + + def find_role(self, role): + role = self._do_find_role(role) + if role: return role + raise security.RoleNotFoundError() def create_role(self, **kwargs): raise NotImplementedError( @@ -19,20 +100,10 @@ class UserDatastore(object): raise NotImplementedError( "User datastore does not implement create_user method") - def _prepare_create_args(self, kwargs): - if not kwargs.has_key('username') and not kwargs.has_key('email'): - raise UserCreationError('Error creating user: username and/or ' - 'email arguments not provided') - - if not kwargs.has_key('password'): - raise UserCreationError('Error creating user: password ' - 'argument not provided') - - now = datetime.utcnow() - kwargs['created_at'], kwargs['modified_at'] = now, now + def add_role_to_user(self, user, role): + raise NotImplementedError( + "User datastore does not implement add_role_to_user method") - pw = kwargs['password'] - if not pwd_context.identify(pw): - kwargs['password'] = pwd_context.encrypt(pw) - - return kwargs \ No newline at end of file + def remove_role_from_user(self, user, role): + raise NotImplementedError( + "User datastore does not implement remove_role_from_user method") \ No newline at end of file diff --git a/flask_security/datastore/mongoengine.py b/flask_security/datastore/mongoengine.py index 39b7527..13dd9a7 100644 --- a/flask_security/datastore/mongoengine.py +++ b/flask_security/datastore/mongoengine.py @@ -8,6 +8,9 @@ class MongoEngineUserDatastore(UserDatastore): def __init__(self, db): self.db = db + def get_models(self): + db = self.db + class Role(db.Document, RoleMixin): name = db.StringField(required=True, unique=True, max_length=80) description = db.StringField(max_length=255) @@ -21,47 +24,30 @@ class MongoEngineUserDatastore(UserDatastore): created_at = db.DateTimeField() modified_at = db.DateTimeField() - security.User = User - security.Role = Role + return User, Role - def with_id(self, id): + def _do_with_id(self, id): try: return security.User.objects.get(id=id) - except: raise security.UserIdNotFoundError() + except: return None - 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 _do_find_user(self, user): + return security.User.objects(username=user).first() or \ + security.User.objects(email=user).first() + + def _do_find_role(self, role): + return security.Role.objects(name=role).first() 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() - + role = security.Role(**self._prepare_create_role_args(kwargs)) + 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 = security.User(**self._prepare_create_user_args(kwargs)) + user.save() + return user + + def add_role(self, user, role): + user = self._do_add_role(user, role) 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 1b5d1a1..4091167 100644 --- a/flask_security/datastore/sqlalchemy.py +++ b/flask_security/datastore/sqlalchemy.py @@ -8,6 +8,9 @@ class SQLAlchemyUserDatastore(UserDatastore): def __init__(self, db): self.db = db + def get_models(self): + db = self.db + roles_users = db.Table('roles_users', db.Column('user_id', db.Integer(), db.ForeignKey('role.id')), db.Column('role_id', db.Integer(), db.ForeignKey('user.id'))) @@ -44,51 +47,36 @@ class SQLAlchemyUserDatastore(UserDatastore): self.created_at = created_at self.modified_at = modified_at - security.User = User - security.Role = Role - - db.create_all() - - def with_id(self, id): - user = security.User.query.get(id) - if user: return user - raise security.UserIdNotFoundError() + return User, Role - def find(self, user_identifier): - user = security.User.query.filter_by(username=user_identifier).first() - if user: return user - user = security.User.query.filter_by(email=user_identifier).first() - if user: return user - raise security.UserNotFoundError() + def _save_model(self, model, commit=True): + self.db.session.add(model) + if commit: self.db.session.commit() + return model + + def _do_with_id(self, id): + return security.User.query.get(id) + + def _do_find_user(self, user): + return security.User.query.filter_by(username=user).first() or \ + security.User.query.filter_by(email=user).first() + + def _do_find_role(self, role): + return security.Role.query.filter_by(name=role).first() def create_role(self, commit=True, **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.query.filter_by(name=name).first() - - if role is None: - role = security.Role(name=name, description=description) - self.db.session.add(role) - if commit: self.db.session.commit() - - return role + role = security.Role(**self._prepare_create_role_args(kwargs)) + return self._save_model(role, commit) def create_user(self, commit=True, **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, commit=False)) - - kwargs['roles'] = user_roles - user = security.User(**kwargs) - self.db.session.add(user) - - if commit: self.db.session.commit() - return user \ No newline at end of file + user = security.User(**self._prepare_create_user_args(kwargs)) + return self._save_model(user, commit) + + def add_role_to_user(self, user, role, commit=True): + user = self._do_add_role(user, role) + return self._save_model(user, commit) + + def remove_role_from_user(self, user, role, commit=True): + user = self._do_remove_role(user, role) + return self._save_model(user, commit) + \ No newline at end of file diff --git a/flask_security/script.py b/flask_security/script.py new file mode 100644 index 0000000..9bcb5d8 --- /dev/null +++ b/flask_security/script.py @@ -0,0 +1,70 @@ +import json +import re +from flask.ext.script import Command, Option +from flask.ext.security import (UserCreationError, UserNotFoundError, + RoleNotFoundError, user_datastore) + +def pprint(obj): + print json.dumps(obj, sort_keys=True, indent=4) + +class CreateUserCommand(Command): + """Create a user""" + + option_list = ( + Option('-u', '--username', dest='username', default=None), + Option('-e', '--email', dest='email', default=None), + Option('-p', '--password', dest='password', default=None), + Option('-a', '--active', dest='active', default=''), + Option('-r', '--roles', dest='roles', default=''), + ) + + def run(self, **kwargs): + # sanitize active input + ai = re.sub(r'\s', '', str(kwargs['active'])) + kwargs['active'] = ai.lower() in ['', 'y','yes', '1', 'active'] + + # sanitize role input a bit + ri = re.sub(r'\s', '', kwargs['roles']) + kwargs['roles'] = [] if ri == '' else ri.split(',') + + user_datastore.create_user(**kwargs) + + print 'User created successfully.' + kwargs['password'] = '****' + pprint(kwargs) + + +class AddRoleCommand(Command): + """Add a role to a user""" + + option_list = ( + Option('-u', '--user', dest='user_identifier'), + Option('-r', '--role', dest='role_name'), + ) + + def run(self, user_identifier, role_name): + user_datastore.add_role_to_user(user_identifier, role_name) + print "Role '%s' added to user '%s' successfully" % (role_name, user_identifier) + +class RemoveRoleCommand(Command): + """Add a role to a user""" + + option_list = ( + Option('-u', '--user', dest='user_identifier'), + Option('-r', '--role', dest='role_name'), + ) + + def run(self, user_identifier, role_name): + user_datastore.remove_role_from_user(user_identifier, role_name) + print "Role '%s' removed from user '%s' successfully" % (role_name, user_identifier) + + +class DeactiveUserCommand(Command): + """Deactive a user""" + + option_list = ( + Option('-u', '--user', dest='user_identifier'), + ) + + def run(self, user_identifier): + user_datastore.deactive_user(user_identifier) \ No newline at end of file diff --git a/setup.py b/setup.py index 1676b5c..a80d7b3 100644 --- a/setup.py +++ b/setup.py @@ -8,21 +8,24 @@ Links ````` * `development version - `_ + `_ """ from setuptools import setup setup( name='Flask-Security', - version='1.0.0', + version='1.1.0-dev', url='https://github.com/mattupstate/flask-security', license='MIT', author='Matthew Wright', author_email='matt@nobien.net', description='Simple security for Flask apps', long_description=__doc__, - packages=['flask_security','flask_security.datastore'], + packages=[ + 'flask_security', + 'flask_security.datastore' + ], zip_safe=False, include_package_data=True, platforms='any', @@ -36,8 +39,13 @@ setup( test_suite='nose.collector', tests_require=[ 'nose', + 'Flask-SQLAlchemy', + 'Flask-MongoEngine', 'py-bcrypt' ], + dependency_links=[ + 'http://github.com/sbook/flask-mongoengine/tarball/master#egg=Flask-MongoEngine-0.1.3-dev' + ], classifiers=[ 'Development Status :: 4 - Beta', 'Environment :: Web Environment',