codekingpro/portable-devtools
114k
1"""2 flask_security.datastore3 ~~~~~~~~~~~~~~~~~~~~~~~~4 5 This module contains an user datastore classes.6 7 :copyright: (c) 2012 by Matt Wright.8 :copyright: (c) 2019-2024 by J. Christopher Wagner (jwag).9 :license: MIT, see LICENSE for more details.10"""11 12from __future__ import annotations13 14from datetime import datetime15import json16import typing as t17import uuid18from copy import copy19 20from .utils import config_value as cv21 22if t.TYPE_CHECKING: # pragma: no cover23 import flask_sqlalchemy24 import mongoengine25 import sqlalchemy.orm.scoping26 27 28class Datastore:29 def __init__(self, db):30 self.db = db31 32 def commit(self):33 pass34 35 def put(self, model):36 raise NotImplementedError37 38 def delete(self, model):39 raise NotImplementedError40 41 42try:43 import sqlalchemy.types as types44 45 class AsaList(types.TypeDecorator):46 """47 SQL-like DBs don't have a List type - so do that here by converting to a comma48 separate string.49 For SQLAlchemy-based datastores, this can be used as::50 51 Column(MutableList.as_mutable(AsaList()), nullable=True)52 """53 54 impl = types.UnicodeText55 56 def process_bind_param(self, value, dialect):57 # produce a string from an iterable58 try:59 return ",".join(value)60 except TypeError:61 return value62 63 def process_result_value(self, value, dialect):64 if value:65 return value.split(",")66 return []67 68except ImportError: # pragma: no cover69 70 class AsaList: # type: ignore71 """72 SQL-like DBs don't have a List type - so do that here by converting to a comma73 separate string.74 For SQLAlchemy-based datastores, this can be used as::75 76 Column(MutableList.as_mutable(AsaList()), nullable=True)77 """78 79 pass80 81 82class SQLAlchemyDatastore(Datastore):83 def commit(self):84 self.db.session.commit()85 86 def put(self, model):87 self.db.session.add(model)88 return model89 90 def delete(self, model):91 self.db.session.delete(model)92 93 94class MongoEngineDatastore(Datastore):95 def put(self, model):96 model.save()97 return model98 99 def delete(self, model):100 model.delete()101 102 103class PeeweeDatastore(Datastore):104 def put(self, model):105 model.save()106 return model107 108 def delete(self, model):109 model.delete_instance(recursive=True)110 111 112def with_pony_session(f):113 from functools import wraps114 115 @wraps(f)116 def decorator(*args, **kwargs):117 from pony.orm import db_session118 from pony.orm.core import local119 from flask import (120 after_this_request,121 current_app,122 has_app_context,123 has_request_context,124 )125 from flask.signals import appcontext_popped126 127 register = local.db_context_counter == 0128 if register and (has_app_context() or has_request_context()):129 db_session.__enter__()130 131 result = f(*args, **kwargs)132 133 if register:134 if has_request_context():135 136 @after_this_request137 def pop(request):138 db_session.__exit__()139 return request140 141 elif has_app_context():142 143 @appcontext_popped.connect_via(current_app._get_current_object())144 def pop(sender, *args, **kwargs):145 while local.db_context_counter:146 db_session.__exit__()147 148 else:149 raise RuntimeError("Needs app or request context")150 return result151 152 return decorator153 154 155class PonyDatastore(Datastore):156 def commit(self):157 self.db.commit()158 159 @with_pony_session160 def put(self, model):161 return model162 163 @with_pony_session164 def delete(self, model):165 model.delete()166 167 168class UserDatastore:169 """Abstracted user datastore.170 171 :param user_model: A user model class definition172 :param role_model: A role model class definition173 :param webauthn_model: A model used to store webauthn registrations174 175 .. important::176 For mutating operations, the user/role will be added to the177 datastore (by calling self.put(<object>). If the datastore is session based178 (such as for SQLAlchemyDatastore) it is up to caller to actually179 commit the transaction by calling datastore.commit().180 181 182 .. note::183 You must implement get_user_mapping in your WebAuthn model184 if your User model doesn't have a primary key Column called 'id'185 """186 187 def __init__(188 self,189 user_model: t.Type[User],190 role_model: t.Type[Role],191 webauthn_model: t.Type[WebAuthn] | None = None,192 ):193 self.user_model = user_model194 self.role_model = role_model195 self.webauthn_model = webauthn_model196 197 if t.TYPE_CHECKING: # pragma: no cover198 # These are available from a DataStore implementation199 def delete(self, model):200 pass201 202 def put(self, model):203 pass204 205 def _prepare_role_modify_args(self, role: str | Role) -> Role | None:206 if isinstance(role, str):207 return self.find_role(role)208 return role209 210 def _prepare_create_user_args(self, **kwargs):211 kwargs.setdefault("active", True)212 roles = copy(kwargs.get("roles", []))213 for i, role in enumerate(roles):214 rn = role.name if isinstance(role, self.role_model) else role215 # see if the role exists216 roles[i] = self.find_role(rn)217 kwargs["roles"] = roles218 kwargs.setdefault("fs_uniquifier", uuid.uuid4().hex)219 if hasattr(self.user_model, "fs_token_uniquifier"):220 kwargs.setdefault("fs_token_uniquifier", uuid.uuid4().hex)221 if hasattr(self.user_model, "fs_webauthn_user_handle"):222 kwargs.setdefault("fs_webauthn_user_handle", uuid.uuid4().hex)223 224 return kwargs225 226 def find_user(self, **kwargs: t.Any) -> User | None:227 """Returns a user matching the provided parameters.228 Besides keyword arguments used to filter the results,229 'case_insensitive' can be passed (defaults to False)230 """231 raise NotImplementedError232 233 def find_role(self, role: str) -> Role | None:234 """Returns a role matching the provided name."""235 raise NotImplementedError236 237 def add_role_to_user(self, user: User, role: Role | str) -> bool:238 """Adds a role to a user.239 240 :param user: The user to manipulate.241 :param role: The role to add to the user. Can be a Role object or242 string role name243 :return: True is role was added, False if role already existed.244 """245 if not (role_obj := self._prepare_role_modify_args(role)):246 raise ValueError(f"Role: {role} doesn't exist")247 if role_obj not in user.roles:248 user.roles.append(role_obj)249 self.put(user)250 return True251 return False252 253 def remove_role_from_user(self, user: User, role: Role | str) -> bool:254 """Removes a role from a user.255 256 :param user: The user to manipulate. Can be an User object or email257 :param role: The role to remove from the user. Can be a Role object or258 string role name259 :return: True if role was removed, False if role doesn't exist or user didn't260 have role.261 """262 rv = False263 role_obj = self._prepare_role_modify_args(role)264 if role_obj in user.roles:265 rv = True266 user.roles.remove(role_obj)267 self.put(user)268 return rv269 270 def add_permissions_to_role(271 self, role: Role | str, permissions: set | list | tuple | str272 ) -> bool:273 """Add one or more permissions to role.274 275 :param role: The role to modify. Can be a Role object or276 string role name277 :param permissions: a set, list, tuple or comma separated string.278 :return: True if permissions added, False if role doesn't exist.279 280 Caller must commit to DB.281 282 .. versionadded:: 4.0.0283 """284 285 rv = False286 if role_obj := self._prepare_role_modify_args(role):287 rv = True288 current_perms = role_obj.get_permissions()289 if isinstance(permissions, set) or isinstance(permissions, tuple):290 permissions = list(permissions)291 elif isinstance(permissions, str):292 permissions = [p.strip() for p in permissions.split(",")]293 # always give a list to DB - some (e.g. Mongo) only take list/tuple294 role_obj.permissions = list(current_perms.union(set(permissions)))295 self.put(role_obj)296 return rv297 298 def remove_permissions_from_role(299 self, role: Role | str, permissions: set | list | tuple | str300 ) -> bool:301 """Remove one or more permissions from a role.302 303 :param role: The role to modify. Can be a Role object or304 string role name305 :param permissions: a set, list, tuple or a comma separated string.306 :return: True if permissions removed, False if role doesn't exist.307 308 Caller must commit to DB.309 310 .. versionadded:: 4.0.0311 """312 313 rv = False314 if role_obj := self._prepare_role_modify_args(role):315 rv = True316 current_perms = role_obj.get_permissions()317 if isinstance(permissions, set) or isinstance(permissions, tuple):318 permissions = list(permissions)319 elif isinstance(permissions, str):320 permissions = [p.strip() for p in permissions.split(",")]321 role_obj.permissions = list(current_perms.difference(set(permissions)))322 self.put(role_obj)323 return rv324 325 def toggle_active(self, user: User) -> bool:326 """Toggles a user's active status. Always returns True."""327 user.active = not user.active328 self.put(user)329 return True330 331 def deactivate_user(self, user: User) -> bool:332 """Deactivates a specified user. Returns `True` if a change was made.333 334 This will immediately disallow access to all endpoints that require335 authentication either via session or tokens.336 The user will not be able to log in again.337 338 :param user: The user to deactivate339 """340 if user.active:341 user.active = False342 self.put(user)343 return True344 return False345 346 def activate_user(self, user: User) -> bool:347 """Activates a specified user. Returns `True` if a change was made.348 349 :param user: The user to activate350 """351 if not user.active:352 user.active = True353 self.put(user)354 return True355 return False356 357 def set_uniquifier(self, user: User, uniquifier: str | None = None) -> None:358 """Set user's Flask-Security identity key.359 This will immediately render outstanding auth tokens,360 session cookies and remember cookies invalid.361 362 :param user: User to modify363 :param uniquifier: Unique value - if none then uuid.uuid4().hex is used364 365 .. versionadded:: 3.3.0366 """367 if not uniquifier:368 uniquifier = uuid.uuid4().hex369 user.fs_uniquifier = uniquifier370 self.put(user)371 372 def set_token_uniquifier(self, user: User, uniquifier: str | None = None) -> None:373 """Set user's auth token identity key.374 This will immediately render outstanding auth tokens invalid.375 376 :param user: User to modify377 :param uniquifier: Unique value - if none then uuid.uuid4().hex is used378 379 This method is a no-op if the user model doesn't contain the attribute380 ``fs_token_uniquifier``381 382 .. versionadded:: 4.0.0383 """384 if not uniquifier:385 uniquifier = uuid.uuid4().hex386 if hasattr(user, "fs_token_uniquifier"):387 user.fs_token_uniquifier = uniquifier388 self.put(user)389 390 def create_role(self, **kwargs: t.Any) -> Role:391 """392 Creates and returns a new role from the given parameters.393 Supported params (depending on RoleModel):394 395 :kwparam name: Role name396 :kwparam permissions: a list, set, tuple or comma separated string.397 These are user-defined strings that correspond to args used with398 @permissions_required()399 400 .. versionadded:: 3.3.0401 402 """403 404 # Usually we just use raw DB model create - for permissions we want to405 # be nicer and allow sending in a list or set or a single string.406 if "permissions" in kwargs and hasattr(self.role_model, "permissions"):407 perms = kwargs["permissions"]408 if isinstance(perms, set) or isinstance(perms, tuple):409 perms = list(perms)410 elif isinstance(perms, str):411 perms = [p.strip() for p in perms.split(",")]412 kwargs["permissions"] = perms413 414 role = self.role_model(**kwargs)415 return self.put(role)416 417 def find_or_create_role(self, name: str, **kwargs: t.Any) -> Role:418 """Returns a role matching the given name or creates it with any419 additionally provided parameters.420 """421 return self.find_role(name) or self.create_role(name=name, **kwargs)422 423 def create_user(self, **kwargs: t.Any) -> User:424 """Creates and returns a new user from the given parameters.425 426 :kwparam email: required.427 :kwparam password: Hashed password.428 :kwparam roles: list of roles to be added to user.429 Can be Role objects or strings430 431 Any other element of the User data model may be supplied as well.432 433 .. note::434 No normalization is done on email - it is assumed the caller has already435 done that.436 437 Best practice is::438 439 try:440 enorm = app.security._mail_util.validate(email)441 except ValueError:442 443 .. danger::444 Be aware that whatever `password` is passed in will445 be stored directly in the DB. Do NOT pass in a plaintext password!446 Best practice is to pass in ``hash_password(plaintext_password)``.447 448 Furthermore, no validation nor normalization is done on the password449 (e.g for minimum length).450 451 Best practice is::452 453 pbad, pnorm = app.security._password_util.validate(password, True)454 455 Look for `pbad` being None. Pass the normalized password `pnorm` to this456 method.457 458 The new user's ``active`` property will be set to ``True``459 unless explicitly set to ``False`` in `kwargs` (e.g. active = False)460 """461 kwargs = self._prepare_create_user_args(**kwargs)462 user = self.user_model(**kwargs)463 return self.put(user)464 465 def delete_user(self, user: User) -> None:466 """Deletes the specified user.467 468 :param user: The user to delete469 """470 self.delete(user) # type: ignore471 472 def reset_user_access(self, user: User) -> None:473 """474 Use this method to reset user authentication methods in the case of compromise.475 This will:476 477 * reset fs_uniquifier - which causes session cookie, remember cookie, auth478 tokens to be unusable479 * reset fs_token_uniquifier (if present) - cause auth tokens to be unusable480 * remove all unified signin TOTP secrets so those can't be used481 * remove all two-factor secrets so those can't be used482 * remove all registered webauthn credentials483 * remove all one-time recovery codes484 * will NOT affect password485 486 Note that if using unified sign in and allow 'email' as a way to receive a code;487 this will also get reset. If the user registered w/o a password then they likely488 will have no way to authenticate.489 490 Note - this method isn't used directly by Flask-Security - it is provided491 as a helper for an application's administrative needs.492 493 Remember to call commit on DB if needed.494 495 .. versionadded:: 3.4.1496 497 .. versionchanged:: 5.0.0498 Added webauthn and recovery codes reset.499 """500 self.set_uniquifier(user)501 self.set_token_uniquifier(user)502 if hasattr(user, "us_totp_secrets"):503 self.us_reset(user)504 if hasattr(user, "tf_primary_method"):505 self.tf_reset(user)506 if hasattr(user, "webauthn"):507 self.webauthn_reset(user)508 if hasattr(user, "mf_recovery_codes"):509 self.mf_set_recovery_codes(user, None)510 511 def tf_set(512 self,513 user: User,514 primary_method: str,515 totp_secret: str | None = None,516 phone: str | None = None,517 ) -> None:518 """Set two-factor info into user record.519 This carefully only changes things if different.520 521 If totp_secret isn't provided - existing one won't be changed.522 If phone isn't provided, the existing phone number won't be changed.523 524 This could be called from an application to apiori setup a user for two factor525 without the user having to go through the setup process.526 527 To get a totp_secret - use ``app.security._totp_factory.generate_totp_secret()``528 529 .. versionadded: 3.4.1530 """531 532 changed = False533 if user.tf_primary_method != primary_method:534 user.tf_primary_method = primary_method535 changed = True536 if totp_secret and user.tf_totp_secret != totp_secret:537 user.tf_totp_secret = totp_secret538 changed = True539 if phone and user.tf_phone_number != phone:540 user.tf_phone_number = phone541 changed = True542 if changed:543 self.put(user)544 545 def tf_reset(self, user: User) -> None:546 """Disable two-factor auth for user.547 548 .. versionadded: 3.4.1549 """550 user.tf_primary_method = None551 user.tf_totp_secret = None552 user.tf_phone_number = None553 self.put(user)554 555 def mf_set_recovery_codes(self, user: User, rcs: list[str] | None) -> None:556 """Set MF recovery codes into user record.557 Any existing codes will be erased.558 559 .. versionadded: 5.0.0560 """561 user.mf_recovery_codes = rcs562 self.put(user)563 564 def mf_get_recovery_codes(self, user: User) -> list[str]:565 codes = getattr(user, "mf_recovery_codes", [])566 return codes if codes else []567 568 def mf_delete_recovery_code(self, user: User, idx: int) -> bool:569 """Delete a single recovery code.570 Recovery codes are single-use - so delete after using!571 572 Return True if code found and deleted, False otherwise.573 574 .. versionadded: 5.0.0575 """576 if not user.mf_recovery_codes:577 return False578 try:579 user.mf_recovery_codes.pop(idx)580 self.put(user)581 return True582 except IndexError:583 return False584 585 def us_get_totp_secrets(self, user: User) -> dict[str, str]:586 """Return totp secrets.587 These are json encoded in the DB.588 589 Returns a dict with methods as keys and secrets as values.590 591 .. versionadded:: 3.4.0592 """593 if not user.us_totp_secrets:594 return {}595 return json.loads(user.us_totp_secrets)596 597 def us_put_totp_secrets(self, user: User, secrets: dict[str, str] | None) -> None:598 """Save secrets. Assume to be a dict (or None)599 with keys as methods, and values as (encrypted) secrets.600 601 .. versionadded:: 3.4.0602 """603 user.us_totp_secrets = json.dumps(secrets) if secrets else None604 self.put(user) # type: ignore605 606 def us_set(607 self,608 user: User,609 method: str,610 totp_secret: str | None = None,611 phone: str | None = None,612 ) -> None:613 """Set unified sign in info into user record.614 615 If totp_secret isn't provided - existing one won't be changed.616 If phone isn't provided, the existing phone number won't be changed.617 618 This could be called from an application to apiori setup a user for unified619 sign in without the user having to go through the setup process.620 621 To get a totp_secret - use ``app.security._totp_factory.generate_totp_secret()``622 623 .. versionadded:: 3.4.1624 """625 626 if totp_secret:627 totp_secrets = self.us_get_totp_secrets(user)628 totp_secrets[method] = totp_secret629 self.us_put_totp_secrets(user, totp_secrets)630 if phone and user.us_phone_number != phone:631 user.us_phone_number = phone632 self.put(user)633 634 def us_reset(self, user: User, method: str | None = None) -> None:635 """Disable unified sign in for user.636 This will disable authenticator app and SMS, and email.637 N.B. if user has no password they may not be able to authenticate at all.638 639 .. versionadded:: 3.4.1640 641 .. versionchanged:: 5.0.0642 Added optional method argument to delete just a single method643 644 """645 if not method:646 # delete all647 self.us_put_totp_secrets(user, None)648 user.us_phone_number = None649 self.put(user)650 else:651 totp_secrets = self.us_get_totp_secrets(user)652 del totp_secrets[method]653 self.us_put_totp_secrets(user, totp_secrets)654 if method == "sms":655 user.us_phone_number = None656 self.put(user)657 658 def us_setup_email(self, user: User) -> bool:659 # setup email (if allowed) for user for unified sign in.660 from .proxies import _security661 662 if not cv("UNIFIED_SIGNIN") or "email" not in cv("US_ENABLED_METHODS"):663 return False664 totp_secrets = self.us_get_totp_secrets(user)665 totp_secrets["email"] = _security._totp_factory.generate_totp_secret()666 self.us_put_totp_secrets(user, totp_secrets)667 return True668 669 def set_webauthn_user_handle(670 self, user: User, user_handle: str | None = None671 ) -> None:672 """Set the value for the Relaying Party's (that's us)673 UserHandle (user.id)674 If no value is passed in, a UUID is generated.675 """676 if not user_handle:677 user_handle = uuid.uuid4().hex678 user.fs_webauthn_user_handle = user_handle679 self.put(user)680 681 def create_webauthn(682 self,683 user: User,684 credential_id: bytes,685 public_key: bytes,686 name: str,687 sign_count: int,688 usage: str,689 device_type: str,690 backup_state: bool,691 transports: list[str] | None = None,692 extensions: str | None = None,693 **kwargs: t.Any,694 ) -> None:695 """696 Create a new webauthn registration record.697 Note that we need to find webauthn records per user as well as698 find a user from a given webauthn (credential_id) record.699 700 .. versionadded: 5.0.0701 """702 raise NotImplementedError703 704 def delete_webauthn(self, webauthn: WebAuthn) -> None:705 """706 .. versionadded: 5.0.0707 """708 self.delete(webauthn)709 710 def find_webauthn(self, credential_id: bytes) -> WebAuthn | None:711 """Returns a credential matching the id.712 713 .. versionadded: 5.0.0714 """715 raise NotImplementedError716 717 def find_user_from_webauthn(self, webauthn: WebAuthn) -> User | None:718 """Returns user associated with this webauthn credential719 720 .. versionadded: 5.0.0721 """722 if not self.webauthn_model:723 raise NotImplementedError724 user_filter = webauthn.get_user_mapping()725 return self.find_user(**user_filter)726 727 def webauthn_reset(self, user: User) -> None:728 """Reset access via webauthn credentials.729 This will DELETE all registered credentials.730 There doesn't appear to be any reason to change the user's731 fs_webauthn_user_handle.732 733 .. versionadded: 5.0.0734 """735 for cred in user.webauthn:736 self.delete(cred)737 self.put(user)738 739 740class SQLAlchemyUserDatastore(SQLAlchemyDatastore, UserDatastore):741 """A UserDatastore implementation that assumes the742 use of743 `Flask-SQLAlchemy <https://pypi.python.org/pypi/flask-sqlalchemy/>`_744 for datastore transactions.745 746 :param db:747 :param user_model: See :ref:`Models <models_topic>`.748 :param role_model: See :ref:`Models <models_topic>`.749 :param webauthn_model: See :ref:`Models <models_topic>`.750 """751 752 def __init__(753 self,754 db: flask_sqlalchemy.SQLAlchemy,755 user_model: t.Type[User],756 role_model: t.Type[Role],757 webauthn_model: t.Type[WebAuthn] | None = None,758 ):759 SQLAlchemyDatastore.__init__(self, db)760 UserDatastore.__init__(self, user_model, role_model, webauthn_model)761 762 def find_user(self, case_insensitive: bool = False, **kwargs: t.Any) -> User | None:763 from sqlalchemy import func as alchemyFn764 765 query = self.user_model.query766 if cv("JOIN_USER_ROLES") and hasattr(self.user_model, "roles"):767 from sqlalchemy.orm import joinedload768 769 query = query.options(joinedload(self.user_model.roles)) # type: ignore770 771 if case_insensitive:772 # While it is of course possible to pass in multiple keys to filter on773 # that isn't the normal use case. If caller asks for case_insensitive774 # AND gives multiple keys - throw an error.775 if len(kwargs) > 1:776 raise ValueError("Case insensitive option only supports single key")777 attr, identifier = kwargs.popitem()778 subquery = alchemyFn.lower(779 getattr(self.user_model, attr)780 ) == alchemyFn.lower(identifier)781 return query.filter(subquery).first()782 else:783 return query.filter_by(**kwargs).first()784 785 def find_role(self, role: str) -> Role | None:786 return self.role_model.query.filter_by(name=role).first() # type: ignore787 788 def find_webauthn(self, credential_id: bytes) -> WebAuthn | None:789 return self.webauthn_model.query.filter_by( # type: ignore790 credential_id=credential_id791 ).first()792 793 def create_webauthn(794 self,795 user: User,796 credential_id: bytes,797 public_key: bytes,798 name: str,799 sign_count: int,800 usage: str,801 device_type: str,802 backup_state: bool,803 transports: list[str] | None = None,804 extensions: str | None = None,805 **kwargs: t.Any,806 ) -> None:807 from .proxies import _security808 809 if not hasattr(self, "webauthn_model") or not self.webauthn_model:810 raise NotImplementedError811 812 webauthn = self.webauthn_model(813 credential_id=credential_id,814 public_key=public_key,815 name=name,816 sign_count=sign_count,817 usage=usage,818 device_type=device_type,819 backup_state=backup_state,820 transports=transports,821 extensions=extensions,822 lastuse_datetime=_security.datetime_factory(),823 **kwargs,824 )825 user.webauthn.append(webauthn)826 self.put(webauthn)827 self.put(user)828 829 830class SQLAlchemySessionUserDatastore(SQLAlchemyUserDatastore, SQLAlchemyDatastore):831 """A UserDatastore implementation that directly uses832 `SQLAlchemy's <https://docs.sqlalchemy.org/en/14/orm/session_basics.html>`_833 session API.834 835 :param session:836 :param user_model: See :ref:`Models <models_topic>`.837 :param role_model: See :ref:`Models <models_topic>`.838 :param webauthn_model: See :ref:`Models <models_topic>`.839 """840 841 def __init__(842 self,843 session: sqlalchemy.orm.scoping.scoped_session,844 user_model: t.Type[User],845 role_model: t.Type[Role],846 webauthn_model: t.Type[WebAuthn] | None = None,847 ):848 class PretendFlaskSQLAlchemyDb:849 """This is a pretend db object, so we can just pass in a session."""850 851 def __init__(self, session):852 self.session = session853 854 SQLAlchemyUserDatastore.__init__(855 self,856 PretendFlaskSQLAlchemyDb(session), # type: ignore857 user_model,858 role_model,859 webauthn_model,860 )861 862 def commit(self):863 super().commit()864 865 866class MongoEngineUserDatastore(MongoEngineDatastore, UserDatastore):867 """A UserDatastore implementation that assumes the868 use of869 `MongoEngine <https://pypi.org/project/mongoengine/>`_870 for datastore transactions.871 872 :param db:873 :param user_model: See :ref:`Models <models_topic>`.874 :param role_model: See :ref:`Models <models_topic>`.875 :param webauthn_model: See :ref:`Models <models_topic>`.876 """877 878 def __init__(879 self,880 db: mongoengine.connection,881 user_model: t.Type[User],882 role_model: t.Type[Role],883 webauthn_model: t.Type[WebAuthn] | None = None,884 ):885 MongoEngineDatastore.__init__(self, db)886 UserDatastore.__init__(self, user_model, role_model, webauthn_model)887 888 def find_user(self, case_insensitive=False, **kwargs):889 from mongoengine.queryset.visitor import Q, QCombination890 from mongoengine.errors import ValidationError891 892 try:893 if case_insensitive:894 # While it is of course possible to pass in multiple keys to filter on895 # that isn't the normal use case. If caller asks for case_insensitive896 # AND gives multiple keys - throw an error.897 if len(kwargs) > 1:898 raise ValueError("Case insensitive option only supports single key")899 attr, identifier = kwargs.popitem()900 query = {f"{attr}__iexact": identifier}901 obj = self.user_model.objects(**query).first()902 else:903 queries = map(lambda i: Q(**{i[0]: i[1]}), kwargs.items())904 query = QCombination(QCombination.AND, queries)905 obj = self.user_model.objects(query).first()906 except ValidationError: # pragma: no cover907 return None908 return obj909 910 def find_role(self, role):911 return self.role_model.objects(name=role).first()912 913 def find_webauthn(self, credential_id: bytes) -> WebAuthn | None:914 if not self.webauthn_model:915 raise NotImplementedError916 917 obj = self.webauthn_model.objects( # type: ignore918 credential_id=credential_id919 ).first()920 return obj921 922 def create_webauthn(923 self,924 user: User,925 credential_id: bytes,926 public_key: bytes,927 name: str,928 sign_count: int,929 usage: str,930 device_type: str,931 backup_state: bool,932 transports: list[str] | None = None,933 extensions: str | None = None,934 **kwargs: t.Any,935 ) -> None:936 from .proxies import _security937 938 if not hasattr(self, "webauthn_model") or not self.webauthn_model:939 raise NotImplementedError940 webauthn = self.webauthn_model(941 user=user,942 credential_id=credential_id,943 public_key=public_key,944 name=name,945 sign_count=sign_count,946 usage=usage,947 device_type=device_type,948 backup_state=backup_state,949 transports=transports,950 extensions=extensions,951 lastuse_datetime=_security.datetime_factory(),952 **kwargs,953 )954 user.webauthn.append(webauthn)955 self.put(webauthn) # type: ignore956 self.put(user) # type: ignore957 958 959class PeeweeUserDatastore(PeeweeDatastore, UserDatastore):960 """A UserDatastore implementation that assumes the961 use of962 `Peewee Flask utils \963 <https://docs.peewee-orm.com/en/latest/peewee/playhouse.html#flask-utils>`_964 for datastore transactions.965 """966 967 def __init__(self, db, user_model, role_model, role_link, webauthn_model=None):968 """969 :param db:970 :param user_model: A user model class definition971 :param role_model: A role model class definition972 :param role_link: A model implementing the many-to-many user-role relation973 :param webauthn_model: A webauthn model class definition974 975 """976 PeeweeDatastore.__init__(self, db)977 UserDatastore.__init__(self, user_model, role_model, webauthn_model)978 self.UserRole = role_link979 980 def find_user(self, case_insensitive=False, **kwargs):981 from peewee import fn as peeweeFn982 983 try:984 if case_insensitive:985 # While it is of course possible to pass in multiple keys to filter on986 # that isn't the normal use case. If caller asks for case_insensitive987 # AND gives multiple keys - throw an error.988 if len(kwargs) > 1:989 raise ValueError("Case insensitive option only supports single key")990 attr, identifier = kwargs.popitem()991 return self.user_model.get(992 peeweeFn.lower(getattr(self.user_model, attr))993 == peeweeFn.lower(identifier)994 )995 else:996 return self.user_model.filter(**kwargs).get()997 except self.user_model.DoesNotExist:998 return None999 1000 def find_role(self, role):1001 try:1002 return self.role_model.filter(name=role).get()1003 except self.role_model.DoesNotExist:1004 return None1005 1006 def create_user(self, **kwargs):1007 """Creates and returns a new user from the given parameters."""1008 roles = kwargs.pop("roles", [])1009 user = self.user_model(**self._prepare_create_user_args(**kwargs))1010 user = self.put(user)1011 for role in roles:1012 self.add_role_to_user(user, role)1013 self.put(user)1014 return user1015 1016 def add_role_to_user(self, user, role):1017 """Adds a role to a user.1018 1019 :param user: The user to manipulate1020 :param role: The role to add to the user1021 """1022 role = self._prepare_role_modify_args(role)1023 result = self.UserRole.select().where(1024 self.UserRole.user == user.id, self.UserRole.role == role.id1025 )1026 if result.count():1027 return False1028 else:1029 self.put(self.UserRole.create(user=user.id, role=role.id))1030 return True1031 1032 def remove_role_from_user(self, user, role):1033 """Removes a role from a user.1034 1035 :param user: The user to manipulate1036 :param role: The role to remove from the user1037 """1038 role = self._prepare_role_modify_args(role)1039 result = self.UserRole.select().where(1040 self.UserRole.user == user, self.UserRole.role == role1041 )1042 if result.count():1043 query = self.UserRole.delete().where(1044 self.UserRole.user == user, self.UserRole.role == role1045 )1046 query.execute()1047 return True1048 else:1049 return False1050 1051 def find_webauthn(self, credential_id):1052 if not self.webauthn_model:1053 raise NotImplementedError1054 try:1055 return self.webauthn_model.filter(credential_id=credential_id).get()1056 except self.webauthn_model.DoesNotExist:1057 return None1058 1059 def create_webauthn(1060 self,1061 user: User,1062 credential_id: bytes,1063 public_key: bytes,1064 name: str,1065 sign_count: int,1066 usage: str,1067 device_type: str,1068 backup_state: bool,1069 transports: list[str] | None = None,1070 extensions: str | None = None,1071 **kwargs: t.Any,1072 ) -> None:1073 from .proxies import _security1074 1075 if not hasattr(self, "webauthn_model") or not self.webauthn_model:1076 raise NotImplementedError1077 webauthn = self.webauthn_model(1078 user=user,1079 credential_id=credential_id,1080 public_key=public_key,1081 name=name,1082 sign_count=sign_count,1083 usage=usage,1084 device_type=device_type,1085 backup_state=backup_state,1086 transports=transports,1087 extensions=extensions,1088 lastuse_datetime=_security.datetime_factory(),1089 **kwargs,1090 )1091 self.put(webauthn) # type: ignore1092 1093 1094class PonyUserDatastore(PonyDatastore, UserDatastore):1095 """A UserDatastore implementation that assumes the1096 use of1097 `PonyORM <https://pypi.python.org/pypi/pony/>`_1098 for datastore transactions.1099 1100 Code primarily from https://github.com/ET-CS but taken over after1101 being abandoned.1102 1103 :param db:1104 :param user_model: See :ref:`Models <models_topic>`.1105 :param role_model: See :ref:`Models <models_topic>`.1106 :param webauthn_model: See :ref:`Models <models_topic>`.1107 """1108 1109 def __init__(self, db, user_model, role_model, webauthn_model=None):1110 PonyDatastore.__init__(self, db)1111 UserDatastore.__init__(self, user_model, role_model, webauthn_model)1112 1113 @with_pony_session1114 def find_user(self, case_insensitive=False, **kwargs):1115 if case_insensitive:1116 # While it is of course possible to pass in multiple keys to filter on1117 # that isn't the normal use case. If caller asks for case_insensitive1118 # AND gives multiple keys - throw an error.1119 if len(kwargs) > 1:1120 raise ValueError("Case insensitive option only supports single key")1121 # TODO - implement case insensitive look ups.1122 1123 return self.user_model.get(**kwargs)1124 1125 @with_pony_session1126 def find_role(self, role):1127 return self.role_model.get(name=role)1128 1129 @with_pony_session1130 def add_role_to_user(self, *args, **kwargs):1131 return super().add_role_to_user(*args, **kwargs)1132 1133 @with_pony_session1134 def create_user(self, **kwargs):1135 return super().create_user(**kwargs)1136 1137 @with_pony_session1138 def create_role(self, **kwargs):1139 return super().create_role(**kwargs)1140 1141 1142if t.TYPE_CHECKING: # pragma: no cover1143 # Normally - the application creates the Models and glues them together1144 # For typing we do that here since we don't know which DB interface they1145 # will pick.1146 from .core import UserMixin, RoleMixin, WebAuthnMixin1147 1148 class CanonicalUserDatastore(Datastore, UserDatastore):1149 pass1150 1151 class User(UserMixin):1152 id: int1153 email: str1154 username: str | None1155 password: str | None1156 active: bool1157 fs_uniquifier: str1158 fs_token_uniquifier: str1159 fs_webauthn_user_handle: str1160 confirmed_at: datetime | None1161 last_login_at: datetime1162 current_login_at: datetime1163 last_login_ip: str | None1164 current_login_ip: str | None1165 login_count: int1166 tf_primary_method: str | None1167 tf_totp_secret: str | None1168 tf_phone_number: str | None1169 mf_recovery_codes: list[str] | None1170 us_phone_number: str | None1171 us_totp_secrets: str | bytes | None1172 create_datetime: datetime1173 update_datetime: datetime1174 roles: list[Role]1175 webauthn: list[WebAuthn]1176 1177 def __init__(self, **kwargs): ...1178 1179 class Role(RoleMixin):1180 id: int1181 name: str1182 description: str | None1183 permissions: list[str] | None1184 update_datetime: datetime1185 1186 def __init__(self, **kwargs): ...1187 1188 class WebAuthn(WebAuthnMixin):1189 id: int1190 name: str1191 credential_id: bytes1192 public_key: bytes1193 sign_count: int1194 transports: list[str] | None1195 backup_state: bool1196 device_type: str1197 extensions: str | None1198 lastuse_datetime: datetime1199 user_id: int1200 usage: str