mirror of
https://github.com/wassname/flask-security.git
synced 2026-06-30 16:40:04 +08:00
Matt Wright
67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
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 |