Python flask_security.Security() Examples

def test_phone_util_override(app):
    class MyPhoneUtil:
        def __init__(self, app):
            pass

        def validate_phone_number(self, input_data):
            return "call-me"

        def get_canonical_form(self, input_data):
            return "very-canonical"

    app.security = Security()
    app.security.init_app(app, phone_util_cls=MyPhoneUtil)

    with app.app_context():
        client = app.test_client()
        # trigger @before first request
        client.get("/login")
        assert uia_phone_mapper("55") == "very-canonical" 

def set_flask_security_config(app, hash_salt=None, secret_key=None):
    """Set all necessary Flask-Security configurations

    :param app: Flask app object
    :param hash_salt: The salt to be used when creating user passwords
    :param secret_key: Secret key used when hashing flask tokens
    """
    hash_salt = hash_salt or config.instance.security_hash_salt
    secret_key = secret_key or config.instance.security_secret_key

    # Make sure that it's possible to get users from the datastore
    # by username and not just by email (the default behavior)
    app.config['SECURITY_USER_IDENTITY_ATTRIBUTES'] = 'username, email'
    app.config['SECURITY_PASSWORD_HASH'] = 'pbkdf2_sha256'
    app.config['SECURITY_HASHING_SCHEMES'] = ['pbkdf2_sha256']
    app.config['SECURITY_DEPRECATED_HASHING_SCHEMES'] = []
    app.config['SECURITY_TOKEN_MAX_AGE'] = 36000  # 10 hours

    app.config['SECURITY_PASSWORD_SALT'] = hash_salt
    app.config['SECURITY_REMEMBER_SALT'] = hash_salt
    app.config['SECRET_KEY'] = secret_key 

def setup_flask_app(manager_ip='localhost',
                    driver='',
                    hash_salt=None,
                    secret_key=None):
    """Setup a functioning flask app, when working outside the rest-service

    :param manager_ip: The IP of the manager
    :param driver: SQLA driver for postgres (e.g. pg8000)
    :param hash_salt: The salt to be used when creating user passwords
    :param secret_key: Secret key used when hashing flask tokens
    :return: A Flask app
    """
    app = Flask(__name__)
    db_uri = _get_postgres_db_uri(manager_ip, driver)
    app.config['SQLALCHEMY_DATABASE_URI'] = db_uri
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
    app.config['ENV'] = 'production'
    set_flask_security_config(app, hash_salt, secret_key)
    Security(app=app, datastore=user_datastore)
    Migrate(app=app, db=db)
    db.init_app(app)
    app.app_context().push()
    return app 

def test_passwordless_custom_form(app, sqlalchemy_datastore):
    app.config["SECURITY_PASSWORDLESS"] = True

    class MyPasswordlessLoginForm(PasswordlessLoginForm):
        email = StringField("My Passwordless Email Address Field")

    app.security = Security(
        app,
        datastore=sqlalchemy_datastore,
        passwordless_login_form=MyPasswordlessLoginForm,
    )

    client = app.test_client()

    response = client.get("/login")
    assert b"My Passwordless Email Address Field" in response.data 

def test_breached_cnt(app):

    # partial response from: https://api.pwnedpasswords.com/range/07003
    pwned_response = "AF5A73CD3CBCFDCD12B0B68CB7930F3E888:2\r\n\
AFD8AA47E6FD782ADDC11D89744769F7354:2\r\n\
B04334E179537C975D0B3C72DA2E5B68E44:15\r\n\
B118F58C2373FDF97ACF93BD3339684D1EB:2\r\n\
B1ED5D27429EDF77EFD84F4EA9BDA5013FB:4\r\n\
B25C03CFBE4CBF19E0F4889711C9A488E5D:2\r\n\
B3902FD808DCA504AAAD30F3C14BD3ACE7C:10".encode(
        "utf-8"
    )

    app.security = Security()
    app.security.init_app(app)
    with app.test_request_context():
        with mock.patch("urllib.request.urlopen") as mock_urlopen:
            mock_urlopen.return_value.__enter__.return_value.read.return_value = (
                pwned_response
            )
            pbad = app.security._password_validator("flaskflask", True)
            # Still weak password, just not pwned enough. Should fail complexity
            assert len(pbad) == 1
            assert "Repeats like" in pbad[0] 

def test_custom_forms_via_config(app, sqlalchemy_datastore):
    class MyLoginForm(LoginForm):
        email = StringField("My Login Email Address Field")

    class MyRegisterForm(RegisterForm):
        email = StringField("My Register Email Address Field")

    app.config["SECURITY_LOGIN_FORM"] = MyLoginForm
    app.config["SECURITY_REGISTER_FORM"] = MyRegisterForm

    security = Security(datastore=sqlalchemy_datastore)
    security.init_app(app)

    client = app.test_client()

    response = client.get("/login")
    assert b"My Login Email Address Field" in response.data

    response = client.get("/register")
    assert b"My Register Email Address Field" in response.data 

def test_breached(app):

    # partial response from: https://api.pwnedpasswords.com/range/07003
    pwned_response = "AF5A73CD3CBCFDCD12B0B68CB7930F3E888:2\r\n\
AFD8AA47E6FD782ADDC11D89744769F7354:2\r\n\
B04334E179537C975D0B3C72DA2E5B68E44:15\r\n\
B118F58C2373FDF97ACF93BD3339684D1EB:2\r\n\
B1ED5D27429EDF77EFD84F4EA9BDA5013FB:4\r\n\
B25C03CFBE4CBF19E0F4889711C9A488E5D:2\r\n\
B3902FD808DCA504AAAD30F3C14BD3ACE7C:10".encode(
        "utf-8"
    )

    app.security = Security()
    app.security.init_app(app)
    with app.test_request_context():
        with mock.patch("urllib.request.urlopen") as mock_urlopen:
            mock_urlopen.return_value.__enter__.return_value.read.return_value = (
                pwned_response
            )
            pbad = app.security._password_validator("flaskflask", False)
            assert len(pbad) == 1
            assert app.config["SECURITY_MSG_PASSWORD_BREACHED"][0] in pbad[0] 

def test_form_labels(app):
    app.config["BABEL_DEFAULT_LOCALE"] = "fr_FR"
    app.security = Security()
    app.security.init_app(app)
    assert check_xlation(app, "fr_FR"), "You must run python setup.py compile_catalog"

    with app.test_request_context():
        rform = RegisterForm()
        assert str(rform.password.label.text) == "Mot de passe"
        assert str(rform.password_confirm.label.text) == "Confirmer le mot de passe"
        assert str(rform.email.label.text) == "Adresse email"
        assert str(rform.submit.label.text) == "Inscription"

        form = LoginForm()
        assert str(form.password.label.text) == "Mot de passe"
        assert str(form.remember.label.text) == "Se souvenir de moi"
        assert str(form.email.label.text) == "Adresse email"
        assert str(form.submit.label.text) == "Connexion"

        form = ChangePasswordForm()
        assert str(form.password.label.text) == "Mot de passe"
        assert str(form.new_password.label.text) == "Nouveau mot de passe"
        assert str(form.new_password_confirm.label.text) == "Confirmer le mot de passe"
        assert str(form.submit.label.text) == "Changer le mot de passe" 

def test_access_datastore_from_factory(app, datastore):
    security = Security()
    security.init_app(app, datastore)

    assert security.datastore is not None
    assert security.app is not None 

def add_flask_security_to_app(app, config):
    _add_configuration_to_app(app, config)

    db = SQLAlchemy(app)

    user_interface = create_user_interface(db)
    security = Security(app, user_interface)

    _add_apikey_handler(security, user_interface)
    return db, user_interface 

def test_access_datastore_from_app_factory_pattern(app, datastore):
    security = Security(datastore=datastore)
    security.init_app(app)

    assert security.datastore is not None
    assert security.app is not None 

def test_init_app_kwargs_override_constructor_kwargs(app, datastore):
    security = Security(
        datastore=datastore,
        login_form="__init__login_form",
        register_form="__init__register_form",
    )
    security.init_app(app, login_form="init_app_login_form")

    assert security.login_form == "init_app_login_form"
    assert security.register_form == "__init__register_form" 

def init_app_with_options(app, datastore, **options):
    security_args = options.pop("security_args", {})
    app.config.update(**options)
    app.security = Security(app, datastore=datastore, **security_args)
    populate_data(app) 

def _set_flask_security(self):
        """Set Flask-Security specific configurations and init the extension
        """
        set_flask_security_config(self)
        Security(app=self, datastore=user_datastore)

        # Get the login manager and set our own callback to be the user getter
        login_manager = self.extensions['security'].login_manager
        login_manager.request_loader(user_loader)

        self.token_serializer = self.extensions[
            'security'].remember_token_serializer 

def init_flask_security(flask_app, database):
    from flask_security import SQLAlchemyUserDatastore, Security
    from psi.app.models.user import User
    from psi.app.models.role import Role
    import psi.app.config as config
    for key, value in config.BaseConfig.security_messages.items():
        flask_app.config['SECURITY_MSG_' + key] = value
    user_datastore = SQLAlchemyUserDatastore(database, User, Role)
    from psi.app.views.login_form import LoginForm
    security = Security(flask_app, user_datastore, login_form=LoginForm)
    return security 

def init():
    global security

    app.config['SECURITY_CONFIRMABLE'] = not app.config['MAIL_SUPPRESS_SEND']
    app.config['SECURITY_CHANGEABLE'] = True
    app.config['SECURITY_SEND_PASSWORD_CHANGE_EMAIL'] = not app.config['MAIL_SUPPRESS_SEND']
    app.config['SECURITY_POST_CHANGE_VIEW'] = "profile.html"
    app.config['SECURITY_PASSWORD_HASH'] = "bcrypt"
    app.config['SECURITY_MSG_CONFIRMATION_REQUIRED'] = (
            flask.Markup('Email requires confirmation. <a href="/confirm">Resend confirmation instructions</a>.'),
            'error')
    # This comes from config: app.config['SECURITY_REGISTERABLE']

    # Update all salts with SECRET_KEY if they are not set
    secret_key = app.config['SECRET_KEY']
    for salt in ('SECURITY_PASSWORD_SALT', 'SECURITY_CONFIRM_SALT',
            'SECURITY_RESET_SALT', 'SECURITY_LOGIN_SALT',
            'SECURITY_REMEMBER_SALT'):
        app.config[salt] = app.config.get(salt, secret_key)

    app.config['SECURITY_EMAIL_SENDER'] = app.config['MAIL_DEFAULT_SENDER']

    app.config['SECURITY_POST_LOGIN_VIEW'] = "/"

    security = flask_security.Security(app, CustomUserDatastore(),
            login_form=CustomLoginForm,
            register_form=CustomRegisterForm,
            confirm_register_form=CustomRegisterForm)

    security.send_mail_task(send_security_mail)

    if app.config['SECURITY_CONFIRMABLE'] and app.config['NEW_USER_NOTIFICATION']:
        flask_security.signals.user_confirmed.connect(new_user_notification, app) 

def add_flask_security(app):
    with app.app_context():
        app.config['SECURITY_UNAUTHORIZED_VIEW'] = '/'
        app.config['SECRET_KEY'] = f3c_global_config.secret_key
        app.config['SECURITY_PASSWORD_SALT'] = f3c_global_config.secret_key

        user_datastore = MongoEngineUserDatastore(db, User, Role)
        security = Security(app, user_datastore)
        create_default_user_and_roles(user_datastore)
        _add_apikey_handler(security, user_datastore) 

def test_breached_real(app):
    """ Actually go out to internet.. """

    app.security = Security()
    app.security.init_app(app)
    with app.test_request_context():
        pbad = app.security._password_validator("flaskflask", True)
        assert len(pbad) == 1
        assert app.config["SECURITY_MSG_PASSWORD_BREACHED"][0] in pbad[0] 

def test_form_required_local_message(app, sqlalchemy_datastore):
    """ Test having a local message (not xlatable and not part of MSG_ config."""

    class MyLoginForm(LoginForm):
        myfield = StringField("My Custom Field", validators=[Required(message="hi")])

    app.config["SECURITY_LOGIN_FORM"] = MyLoginForm

    security = Security(datastore=sqlalchemy_datastore)
    security.init_app(app)

    client = app.test_client()

    response = client.post("/login", content_type="application/json")
    assert response.status_code == 400
    assert b"myfield" in response.data 

def test_form_required(app, sqlalchemy_datastore):
    class MyLoginForm(LoginForm):
        myfield = StringField("My Custom Field", validators=[Required()])

    app.config["SECURITY_LOGIN_FORM"] = MyLoginForm

    security = Security(datastore=sqlalchemy_datastore)
    security.init_app(app)

    client = app.test_client()

    response = client.post("/login", content_type="application/json")
    assert response.status_code == 400
    assert b"myfield" in response.data 

def test_change_hash_type(app, sqlalchemy_datastore):
    init_app_with_options(
        app,
        sqlalchemy_datastore,
        **{
            "SECURITY_PASSWORD_HASH": "plaintext",
            "SECURITY_PASSWORD_SALT": None,
            "SECURITY_PASSWORD_SCHEMES": ["bcrypt", "plaintext"],
        }
    )

    app.config["SECURITY_PASSWORD_HASH"] = "bcrypt"
    app.config["SECURITY_PASSWORD_SALT"] = "salty"

    app.security = Security(
        app, datastore=sqlalchemy_datastore, register_blueprint=False
    )

    client = app.test_client()

    response = client.post(
        "/login", data=dict(email="matt@lp.com", password="password")
    )
    assert response.status_code == 302

    response = client.get("/logout")

    response = client.post(
        "/login", data=dict(email="matt@lp.com", password="password")
    )
    assert response.status_code == 302 

def test_register_blueprint_flag(app, sqlalchemy_datastore):
    app.security = Security(app, datastore=Security, register_blueprint=False)
    client = app.test_client()
    response = client.get("/login")
    assert response.status_code == 404 

def test_nullable_username(app, sqlalchemy_datastore):
    # sqlalchemy datastore uses fsqlav2 which has username as unique and nullable
    # make sure can register multiple users with no username
    # Note that current WTForms (2.2.1) has a bug where StringFields can never be
    # None - it changes them to an empty string. DBs don't like that if you have
    # your column be 'nullable'.
    class NullableStringField(StringField):
        def process_formdata(self, valuelist):
            if valuelist:
                self.data = valuelist[0]

    class MyRegisterForm(ConfirmRegisterForm):
        username = NullableStringField("Username")

    app.config["SECURITY_CONFIRM_REGISTER_FORM"] = MyRegisterForm
    security = Security(datastore=sqlalchemy_datastore)
    security.init_app(app)

    client = app.test_client()

    data = dict(email="u1@test.com", password="password", password_confirm="password")
    response = client.post(
        "/register", json=data, headers={"Content-Type": "application/json"}
    )
    assert response.status_code == 200
    logout(client)

    data = dict(email="u2@test.com", password="password", password_confirm="password")
    response = client.post(
        "/register", json=data, headers={"Content-Type": "application/json"}
    )
    assert response.status_code == 200 

def test_form_data_is_passed_to_user_registered_signal(app, sqlalchemy_datastore):
    class MyRegisterForm(RegisterForm):
        additional_field = StringField("additional_field")

    app.security = Security(
        app, datastore=sqlalchemy_datastore, register_form=MyRegisterForm
    )

    recorded = []

    @user_registered.connect_via(app)
    def on_user_registered(app, user, confirm_token, form_data):

        assert isinstance(app, Flask)
        assert isinstance(user, UserMixin)
        assert confirm_token is None
        assert form_data["additional_field"] == "additional_data"

        recorded.append(user)

    client = app.test_client()

    data = dict(
        email="dude@lp.com",
        password="password",
        password_confirm="password",
        additional_field="additional_data",
    )
    response = client.post("/register", data=data, follow_redirects=True)

    assert response.status_code == 200
    assert len(recorded) == 1 

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="trp@lp.com",
            password=hash_password("password"),
            tf_primary_method="sms",
            tf_totp_secret=app.security._totp_factory.generate_totp_secret(),
        )
        ds.commit()

        data = dict(email="trp@lp.com", password="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_basic_custom_forms(app, sqlalchemy_datastore):
    class MyLoginForm(LoginForm):
        email = StringField("My Login Email Address Field")

    class MyRegisterForm(RegisterForm):
        email = StringField("My Register Email Address Field")

    class MyForgotPasswordForm(ForgotPasswordForm):
        email = StringField(
            "My Forgot Email Address Field",
            validators=[email_required, email_validator, valid_user_email],
        )

    class MyResetPasswordForm(ResetPasswordForm):
        password = StringField("My Reset Password Field")

    class MyChangePasswordForm(ChangePasswordForm):
        password = PasswordField("My Change Password Field")

    app.security = Security(
        app,
        datastore=sqlalchemy_datastore,
        login_form=MyLoginForm,
        register_form=MyRegisterForm,
        forgot_password_form=MyForgotPasswordForm,
        reset_password_form=MyResetPasswordForm,
        change_password_form=MyChangePasswordForm,
    )

    populate_data(app)
    client = app.test_client()

    response = client.get("/login")
    assert b"My Login Email Address Field" in response.data

    response = client.get("/register")
    assert b"My Register Email Address Field" in response.data

    response = client.get("/reset")
    assert b"My Forgot Email Address Field" in response.data

    with capture_reset_password_requests() as requests:
        response = client.post("/reset", data=dict(email="matt@lp.com"))

    token = requests[0]["token"]
    response = client.get("/reset/" + token)
    assert b"My Reset Password Field" in response.data

    authenticate(client)

    response = client.get("/change")
    assert b"My Change Password Field" in response.data 

def test_easy_password(app, sqlalchemy_datastore):
    class MyRegisterForm(ConfirmRegisterForm):
        username = StringField("Username")

    app.config["SECURITY_CONFIRM_REGISTER_FORM"] = MyRegisterForm
    security = Security(datastore=sqlalchemy_datastore)
    security.init_app(app)

    client = app.test_client()

    # With zxcvbn
    data = dict(
        email="dude@lp.com",
        username="dude",
        password="mattmatt2",
        password_confirm="mattmatt2",
    )
    response = client.post(
        "/register", json=data, headers={"Content-Type": "application/json"}
    )
    assert response.headers["Content-Type"] == "application/json"
    assert response.status_code == 400
    # Response from zxcvbn
    assert "Repeats like" in response.json["response"]["errors"]["password"][0]

    # Test that username affects password selection
    data = dict(
        email="dude@lp.com",
        username="Joe",
        password="JoeTheDude",
        password_confirm="JoeTheDude",
    )
    response = client.post(
        "/register", json=data, headers={"Content-Type": "application/json"}
    )
    assert response.headers["Content-Type"] == "application/json"
    assert response.status_code == 400
    # Response from zxcvbn
    assert (
        "Password not complex enough"
        in response.json["response"]["errors"]["password"][0]
    )