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
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 rm(ctx, email):
"""Deletes a user"""
ds = SQLAlchemyUserDatastore(ctx.obj, User, Role)
u = ds.find_user(email=email)
ds.delete_user(u)
ds.commit()