def add_role(ctx, email, role): ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) try: ds.add_role_to_user(email, role) ds.commit() except: log.exception("Couldn't add role")
def register_users_from_manifest(user_datastore: SQLAlchemyUserDatastore, manifest: Dict) -> None: """Register the users and roles in a manifest. :param user_datastore: A user data store :param dict manifest: A manifest dictionary, which contains two keys: ``roles`` and ``users``. The ``roles`` key corresponds to a list of dictionaries containing ``name`` and ``description`` entries. The ``users`` key corresponds to a list of dictionaries containing ``email``, ``password``, and ``name`` entries as well as a optional ``roles`` entry with a corresponding list relational to the names in the ``roles`` entry in the manifest. """ for role in manifest['roles']: user_datastore.find_or_create_role(**role) for user_manifest in manifest['users']: email = user_manifest['email'] user = user_datastore.find_user(email=email) if user is None: logger.info(f'creating user: {email}') user = user_datastore.create_user( confirmed_at=datetime.datetime.now(), email=email, password=user_manifest['password'], name=user_manifest['name'], ) for role_name in user_manifest.get('roles', []): if user_datastore.add_role_to_user(user, role_name): logger.info(f'registered {user} as {role_name}') user_datastore.commit()
def rm(ctx, name): """Deletes a user""" ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) u = ds.find_role(name) if u: ds.delete(u) ds.commit()
def _init_users(): """Initialize Afterglow user datastore if AUTH_ENABLED = True""" # noinspection PyUnresolvedReferences from .. import oauth2 # register oauth token-related models # All imports put here to avoid unnecessary loading of packages on startup # if user auth is disabled from alembic import (config as alembic_config, context as alembic_context) from alembic.script import ScriptDirectory from alembic.runtime.environment import EnvironmentContext global user_datastore, security user_datastore = SQLAlchemyUserDatastore(db, DbUser, DbRole) security = Security(app, user_datastore, register_blueprint=False) # Make sure that the database directory exists try: os.makedirs(os.path.abspath(app.config['DATA_ROOT'])) except OSError as e: if e.errno != errno.EEXIST: raise # Create/upgrade tables via Alembic cfg = alembic_config.Config() cfg.set_main_option( 'script_location', os.path.abspath( os.path.join(__file__, '../..', 'db_migration', 'users'))) script = ScriptDirectory.from_config(cfg) # noinspection PyProtectedMember with EnvironmentContext( cfg, script, fn=lambda rev, _: script._upgrade_revs('head', rev), as_sql=False, starting_rev=None, destination_rev='head', tag=None, ), db.engine.connect() as connection: alembic_context.configure(connection=connection) with alembic_context.begin_transaction(): alembic_context.run_migrations() # Initialize user roles if missing try: roles_created = False for name, descr in [('admin', 'Afterglow Administrator'), ('user', 'Afterglow User')]: if not user_datastore.find_role(name): user_datastore.create_role(name=name, description=descr) roles_created = True if roles_created: user_datastore.commit() except Exception: db.session.rollback() raise
def add(ctx, email, password): """Creates a new user""" ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) try: ds.create_user(email=email, password=password) ds.commit() except: log.exception("Couldn't create user")
def make_admin(ctx, email): """Makes a given user an admin""" ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) try: ds.add_role_to_user(email, 'admin') ds.commit() except: log.exception("Couldn't make admin")
def add(ctx, name, description): """Creates a new role""" ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) try: ds.create_role(name=name, description=description) ds.commit() except: log.exception("Couldn't create role")
def test_auth_uniquifier(app): # If add fs_token_uniquifier to user model - change password shouldn't invalidate # auth tokens. from sqlalchemy import Column, String from flask_sqlalchemy import SQLAlchemy from flask_security.models import fsqla_v2 as fsqla from flask_security import Security, SQLAlchemyUserDatastore app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" db = SQLAlchemy(app) fsqla.FsModels.set_db_info(db) class Role(db.Model, fsqla.FsRoleMixin): pass class User(db.Model, fsqla.FsUserMixin): fs_token_uniquifier = Column(String(64), unique=True, nullable=False) with app.app_context(): db.create_all() ds = SQLAlchemyUserDatastore(db, User, Role) app.security = Security(app, datastore=ds) with app.app_context(): ds.create_user( email="*****@*****.**", password=hash_password("password"), ) ds.commit() client = app.test_client() # standard login with auth token response = json_authenticate(client) token = response.json["response"]["user"]["authentication_token"] headers = {"Authentication-Token": token} # make sure can access restricted page response = client.get("/token", headers=headers) assert b"Token Authentication" in response.data # change password response = client.post( "/change", data={ "password": "******", "new_password": "******", "new_password_confirm": "new strong password", }, follow_redirects=True, ) assert response.status_code == 200 # authtoken should still be valid response = client.get("/token", headers=headers) assert response.status_code == 200
def test_replace_send_code(app, get_message): # replace tf_send_code - and have it return an error to check that. from flask_sqlalchemy import SQLAlchemy from flask_security.models import fsqla_v2 as fsqla from flask_security import Security, hash_password app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" db = SQLAlchemy(app) fsqla.FsModels.set_db_info(db) class Role(db.Model, fsqla.FsRoleMixin): pass class User(db.Model, fsqla.FsUserMixin): rv = [None, "That didnt work out as we planned", "Failed Again"] def tf_send_security_token(self, method, **kwargs): return User.rv.pop(0) with app.app_context(): db.create_all() ds = SQLAlchemyUserDatastore(db, User, Role) app.security = Security(app, datastore=ds) with app.app_context(): client = app.test_client() ds.create_user( email="*****@*****.**", password=hash_password("password"), tf_primary_method="sms", tf_totp_secret=app.security._totp_factory.generate_totp_secret(), ) ds.commit() data = dict(email="*****@*****.**", password="******") response = client.post("/login", data=data, follow_redirects=True) assert b"Please enter your authentication code" in response.data rescue_data = dict(help_setup="lost_device") response = client.post("/tf-rescue", data=rescue_data, follow_redirects=True) assert b"That didnt work out as we planned" in response.data # Test JSON headers = { "Accept": "application/json", "Content-Type": "application/json" } response = client.post("/tf-rescue", json=rescue_data, headers=headers) assert response.status_code == 500 assert response.json["response"]["errors"]["help_setup"][ 0] == "Failed Again"
def test_change_token_uniquifier(app): # make sure that existing token no longer works once we change the token uniquifier from sqlalchemy import Column, String from flask_sqlalchemy import SQLAlchemy from flask_security.models import fsqla_v2 as fsqla from flask_security import Security, SQLAlchemyUserDatastore app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" db = SQLAlchemy(app) fsqla.FsModels.set_db_info(db) class Role(db.Model, fsqla.FsRoleMixin): pass class User(db.Model, fsqla.FsUserMixin): fs_token_uniquifier = Column(String(64), unique=True, nullable=False) with app.app_context(): db.create_all() ds = SQLAlchemyUserDatastore(db, User, Role) app.security = Security(app, datastore=ds) with app.app_context(): ds.create_user( email="*****@*****.**", password=hash_password("password"), ) ds.commit() client_nc = app.test_client(use_cookies=False) response = json_authenticate(client_nc) token = response.json["response"]["user"]["authentication_token"] verify_token(client_nc, token) # now change uniquifier with app.test_request_context("/"): user = app.security.datastore.find_user(email="*****@*****.**") app.security.datastore.reset_user_access(user) app.security.datastore.commit() verify_token(client_nc, token, status=401) # get new token and verify it works response = json_authenticate(client_nc) token = response.json["response"]["user"]["authentication_token"] verify_token(client_nc, token)
def make_user(manager, email, password): """Create a pre-existing user an admin.""" # Example: python3 -m compath make_admin [email protected] password ds = SQLAlchemyUserDatastore(manager, User, Role) user = ds.find_user(email=email) if user is None: ds.create_user(email=email, password=password, confirmed_at=datetime.datetime.utcnow()) ds.commit() click.echo('User {} was successfully created'.format(email)) else: click.echo('User {} already exists'.format(email))
def make_admin(manager, email): """Make a pre-existing user an admin.""" # Example: python3 -m compath make_admin [email protected] ds = SQLAlchemyUserDatastore(manager, User, Role) user = ds.find_user(email=email) if user is None: click.echo('User not found') sys.exit(0) admin = ds.find_or_create_role('admin') ds.add_role_to_user(user, admin) ds.commit() click.echo('User {} is now admin'.format(email))
def test_null_token_uniquifier(app): # If existing record has a null fs_token_uniquifier, should be set on first use. from sqlalchemy import Column, String from flask_sqlalchemy import SQLAlchemy from flask_security.models import fsqla_v2 as fsqla from flask_security import Security, SQLAlchemyUserDatastore app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" db = SQLAlchemy(app) fsqla.FsModels.set_db_info(db) class Role(db.Model, fsqla.FsRoleMixin): pass class User(db.Model, fsqla.FsUserMixin): fs_token_uniquifier = Column(String(64), unique=True, nullable=True) with app.app_context(): db.create_all() ds = SQLAlchemyUserDatastore(db, User, Role) app.security = Security(app, datastore=ds) with app.app_context(): ds.create_user( email="*****@*****.**", password=hash_password("password"), ) ds.commit() # manually null out fs_token_uniquifier user = ds.find_user(email="*****@*****.**") user.fs_token_uniquifier = None ds.put(user) ds.commit() client_nc = app.test_client(use_cookies=False) response = json_authenticate(client_nc) token = response.json["response"]["user"]["authentication_token"] verify_token(client_nc, token)
def load(ctx, file): """Dump stuff for loading later (in lieu of having proper migrations)""" ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) for line in file: email, first, last, roles, password = line.strip().split('\t') u = ds.find_user(email=email) if not u: u = ds.create_user(email=email, first_name=first, last_name=last, password=password) log.info('added %s', u) ds.commit() for role_name in roles.strip().split(','): r = ds.find_role(role_name) if not r: r = ds.create_role(name=role_name) ds.commit() if not u.has_role(r): ds.add_role_to_user(u, r) ds.commit()
user_datastore = SQLAlchemyUserDatastore(_db, User, Role) try: _db.create_all() except OperationalError: pass Migrate(app, _db) Security(app, user_datastore, register_form=RegisterForm, login_form=LoginForm) admin_role = 'admin' user_datastore.find_or_create_role(name=admin_role) try: user_datastore.commit() except IntegrityError: user_datastore.db.session.rollback() def login_required(func): if AppConfig.TESTING: return func return _login_required(func) def admin_required(func): if AppConfig.TESTING: return func
def test_no_sms(app, get_message): # Make sure that don't require tf_phone_number if SMS isn't an option. from sqlalchemy import ( Boolean, Column, Integer, String, ) from sqlalchemy.orm import relationship, backref from flask_sqlalchemy import SQLAlchemy from flask_security.models import fsqla_v2 as fsqla from flask_security import Security, UserMixin, hash_password app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" db = SQLAlchemy(app) fsqla.FsModels.set_db_info(db) class Role(db.Model, fsqla.FsRoleMixin): pass class User(db.Model, UserMixin): id = Column(Integer, primary_key=True) email = Column(String(255), unique=True, nullable=False) password = Column(String(255), nullable=False) active = Column(Boolean(), nullable=False) # Faster token checking fs_uniquifier = Column(String(64), unique=True, nullable=False) # 2FA tf_primary_method = Column(String(64), nullable=True) tf_totp_secret = Column(String(255), nullable=True) roles = relationship( "Role", secondary="roles_users", backref=backref("users", lazy="dynamic") ) with app.app_context(): db.create_all() ds = SQLAlchemyUserDatastore(db, User, Role) app.security = Security(app, datastore=ds) with app.app_context(): client = app.test_client() ds.create_user( email="*****@*****.**", password=hash_password("password"), ) ds.commit() data = dict(email="*****@*****.**", password="******") client.post("/login", data=data, follow_redirects=True) with app.mail.record_messages() as outbox: response = client.post( "/tf-setup", data=dict(setup="email"), follow_redirects=True ) msg = b"To complete logging in, please enter the code sent to your mail" assert msg in response.data code = outbox[0].body.split()[-1] # sumbit right token and show appropriate response response = client.post( "/tf-validate", data=dict(code=code), follow_redirects=True ) assert b"You successfully changed your two-factor method" in response.data
def create_app(test_config=None): # create and configure the app app = Flask(__name__, instance_relative_config=True) app.config.from_object(os.environ['APP_SETTINGS']) db.init_app(app) db.app = app # hook up mail mail = Mail(app) try: # initialize data store and setup roles, security # this will fail if the db does not exist ... user_datastore = SQLAlchemyUserDatastore(db, User, Role) user_datastore.find_or_create_role(name='admin', description='Administrator') user_datastore.find_or_create_role(name='client', description='Client') user_datastore.commit() security = Security(app, user_datastore) except: pass with app.app_context(): # Create admin if admin.app is None: admin.init_app(app) app.config['FLASK_ADMIN_SWATCH'] = 'cosmo' admin.add_view(UserAdmin(User, db.session)) admin.add_view(RoleAdmin(Role, db.session)) admin.add_view(BaseAdmin(Company, db.session)) admin.add_view(BaseAdmin(UserRequest, db.session)) admin.add_view(BaseAdmin(Project, db.session)) admin.add_view(BaseAdmin(ServiceAgreement, db.session)) # Initialize the plugin manager plugin_manager = PluginManager(app) plugins = get_enabled_plugins() if test_config is None: # load the instance config, if it exists, when not testing app.config.from_pyfile('config.py', silent=True) else: # load the test config if passed in app.config.from_mapping(test_config) # ensure the instance folder exists try: os.makedirs(app.instance_path) except OSError: pass app.register_blueprint(user_bp) app.register_blueprint(company_bp) app.register_blueprint(request_bp) app.register_blueprint(project_bp) # internal version @app.route('/version') def version(): return app_version # index @app.route("/") def index(): return render_template("index.html") @app.route("/plugins") @login_required def plugins(): return render_template("plugins.html", plugins=get_enabled_plugins()) @app.route("/disable/<plugin>") def disable(plugin): plugin = get_plugin(plugin) plugin_manager.disable_plugins([plugin]) return redirect(url_for("index")) @app.route("/enable/<plugin>") def enable(plugin): plugin = get_plugin(plugin) plugin_manager.enable_plugins([plugin]) return redirect(url_for("index")) return app
def rm(ctx, email): """Deletes a user""" ds = SQLAlchemyUserDatastore(ctx.obj, User, Role) u = ds.find_user(email=email) ds.delete_user(u) ds.commit()