Python flask_wtf.CSRFProtect() Examples

def test_cp_reset(app, client):
    """ Test that header based CSRF works for /reset when
    using WTF_CSRF_CHECK_DEFAULT=False.
    """
    app.config["WTF_CSRF_ENABLED"] = True
    app.config["WTF_CSRF_CHECK_DEFAULT"] = False
    CSRFProtect(app)

    with mp_validate_csrf() as mp:
        data = dict(email="matt@lp.com")
        # should fail - no CSRF token
        response = client.post("/reset", content_type="application/json", json=data)
        assert response.status_code == 400

        csrf_token = _get_csrf_token(client)
        response = client.post(
            "/reset",
            content_type="application/json",
            json=data,
            headers={"X-CSRF-Token": csrf_token},
        )
        assert response.status_code == 200
    # 2 failures since the first time it will check twice - once due to @unauth_csrf
    # which will fall-through on error to form validation (which also fails).
    assert mp.success == 1 and mp.failure == 2 

def test_cp_with_token(app, client):
    # Make sure can use returned CSRF-Token in Header.
    app.config["WTF_CSRF_ENABLED"] = True
    CSRFProtect(app)

    auth_token, csrf_token = json_login(client, use_header=True)

    # make sure returned csrf_token works in header.
    data = dict(
        password="password",
        new_password="battery staple",
        new_password_confirm="battery staple",
    )

    with mp_validate_csrf() as mp:
        response = client.post(
            "/change",
            content_type="application/json",
            json=data,
            headers={"X-CSRF-Token": csrf_token},
        )
        assert response.status_code == 200
    assert mp.success == 1 and mp.failure == 0
    json_logout(client) 

def test_csrf_cookie(app, client):
    app.config["WTF_CSRF_ENABLED"] = True
    app.config["WTF_CSRF_CHECK_DEFAULT"] = False
    CSRFProtect(app)

    json_login(client)
    found = False
    for cookie in client.cookie_jar:
        if cookie.name == "X-XSRF-Token":
            found = True
            assert cookie.path == "/"
    assert found

    # Make sure cleared on logout
    response = client.post("/logout", content_type="application/json")
    assert response.status_code == 200
    assert "X-XSRF-Token" not in [c.name for c in client.cookie_jar] 

def test_remember_login_csrf_cookie(app, client):
    # Test csrf cookie upon resuming a remember session
    app.config["WTF_CSRF_ENABLED"] = True
    CSRFProtect(app)

    # Login with remember_token generation
    json_login(client, use_header=True, remember=True)

    csrf_cookie = [c for c in client.cookie_jar if c.name == "X-XSRF-Token"][0]
    session_cookie = [c for c in client.cookie_jar if c.name == "session"][0]
    # Delete session and csrf cookie - we should always get new ones
    client.delete_cookie(csrf_cookie.domain, csrf_cookie.name)
    client.delete_cookie(session_cookie.domain, session_cookie.name)

    # Do a simple get request with the remember_token cookie present
    assert "remember_token" in [c.name for c in client.cookie_jar]
    response = client.get("/profile")
    assert response.status_code == 200
    assert "session" in [c.name for c in client.cookie_jar]
    assert "X-XSRF-Token" in [c.name for c in client.cookie_jar]
    # Logout and check that everything cleans up nicely
    json_logout(client)
    assert "remember_token" not in [c.name for c in client.cookie_jar]
    assert "session" not in [c.name for c in client.cookie_jar]
    assert "X-XSRF-Token" not in [c.name for c in client.cookie_jar] 

def test_login_csrf_json(app, client):
    app.config["WTF_CSRF_ENABLED"] = True

    with mp_validate_csrf() as mp:
        auth_token, csrf_token = json_login(client)
        assert auth_token
        assert csrf_token
    # Should be just one call to validate - since CSRFProtect not enabled.
    assert mp.success == 1 and mp.failure == 0

    response = json_logout(client)
    session = get_session(response)
    assert "csrf_token" not in session 

def test_login_csrf_json_header(app, client):
    app.config["WTF_CSRF_ENABLED"] = True
    CSRFProtect(app)

    with mp_validate_csrf() as mp:
        auth_token, csrf_token = json_login(client, use_header=True)
        assert auth_token
        assert csrf_token
    assert mp.success == 2 and mp.failure == 0
    json_logout(client) 

def test_cp_login_json_no_session(app, client_nc):
    # Test with global CSRFProtect on and not sending cookie - nothing works.
    app.config["WTF_CSRF_ENABLED"] = True
    CSRFProtect(app)

    # This shouldn't log in - and will return 400
    with mp_validate_csrf() as mp:

        data = dict(email="matt@lp.com", password="password", remember="y")
        response = client_nc.post(
            "/login",
            content_type="application/json",
            json=data,
            headers={"Accept": "application/json"},
        )
        assert response.status_code == 400

        # This still wont work since we don't send a session cookie
        response = client_nc.post(
            "/login",
            content_type="application/json",
            json=data,
            headers={"X-CSRF-Token": _get_csrf_token(client_nc)},
        )
        assert response.status_code == 400

    # Although failed - CSRF should have been called
    assert mp.failure == 2 

def test_cp_config2(app, client):
    # Test improper config (must have CSRFProtect configured if setting
    # CSRF_PROTECT_MECHANISMS
    app.config["WTF_CSRF_ENABLED"] = True

    # The check is done on first request.
    with pytest.raises(ValueError) as ev:
        logout(client)
    assert "CsrfProtect not part of application" in str(ev.value) 

def test_different_mechanisms(app, client):
    # Verify that using token doesn't require CSRF, but sessions do
    app.config["WTF_CSRF_ENABLED"] = True
    app.config["WTF_CSRF_CHECK_DEFAULT"] = False
    CSRFProtect(app)

    with mp_validate_csrf() as mp:
        auth_token, csrf_token = json_login(client)

        # session based change password should fail
        data = dict(
            password="password",
            new_password="battery staple",
            new_password_confirm="battery staple",
        )

        response = client.post(
            "/change", json=data, headers={"Content-Type": "application/json"}
        )
        assert response.status_code == 400
        assert b"The CSRF token is missing" in response.data

        # token based should work
        response = client.post(
            "/change",
            json=data,
            headers={
                "Content-Type": "application/json",
                "Authentication-Token": auth_token,
            },
        )
        assert response.status_code == 200
    assert mp.success == 1 and mp.failure == 2 

def test_different_mechanisms_nc(app, client_nc):
    # Verify that using token and no session cookie works
    # Note that we had to disable unauth_endpoints since you can't log in
    # w/ CSRF if you don't send in the session cookie.
    app.config["WTF_CSRF_ENABLED"] = True
    app.config["WTF_CSRF_CHECK_DEFAULT"] = False
    CSRFProtect(app)

    with mp_validate_csrf() as mp:
        auth_token, csrf_token = json_login(client_nc)

        # token based should work
        data = dict(
            password="password",
            new_password="battery staple",
            new_password_confirm="battery staple",
        )
        response = client_nc.post(
            "/change",
            json=data,
            headers={
                "Content-Type": "application/json",
                "Authentication-Token": auth_token,
            },
        )
        assert response.status_code == 200
    assert mp.success == 0 and mp.failure == 0 

def test_cp_with_token_cookie(app, client):
    # Make sure can use returned CSRF-Token cookie in Header.
    app.config["WTF_CSRF_ENABLED"] = True
    CSRFProtect(app)

    json_login(client, use_header=True)

    # make sure returned csrf_token works in header.
    data = dict(
        password="password",
        new_password="battery staple",
        new_password_confirm="battery staple",
    )
    csrf_token = [c.value for c in client.cookie_jar if c.name == "X-XSRF-Token"][0]
    with mp_validate_csrf() as mp:
        response = client.post(
            "/change",
            content_type="application/json",
            json=data,
            headers={"X-XSRF-Token": csrf_token},
        )
        assert response.status_code == 200
    # 2 successes since the utils:csrf_cookie_handler will check
    assert mp.success == 2 and mp.failure == 0
    json_logout(client)
    assert "X-XSRF-Token" not in [c.name for c in client.cookie_jar] 

def test_cp_with_token_cookie_refresh(app, client):
    # Test CSRF_COOKIE_REFRESH_EACH_REQUEST
    app.config["WTF_CSRF_ENABLED"] = True
    CSRFProtect(app)

    json_login(client, use_header=True)

    # make sure returned csrf_token works in header.
    data = dict(
        password="password",
        new_password="battery staple",
        new_password_confirm="battery staple",
    )

    csrf_cookie = [c for c in client.cookie_jar if c.name == "X-XSRF-Token"][0]
    with mp_validate_csrf() as mp:
        # Delete cookie - we should always get a new one
        client.delete_cookie(csrf_cookie.domain, csrf_cookie.name)
        response = client.post(
            "/change",
            content_type="application/json",
            json=data,
            headers={"X-XSRF-Token": csrf_cookie.value},
        )
        assert response.status_code == 200
        csrf_cookie = [c for c in client.cookie_jar if c.name == "X-XSRF-Token"][0]
        assert csrf_cookie
    assert mp.success == 1 and mp.failure == 0
    json_logout(client)
    assert "X-XSRF-Token" not in [c.name for c in client.cookie_jar] 

def configure_wtf(self) -> None:
        if self.config["WTF_CSRF_ENABLED"]:
            csrf = CSRFProtect(self.flask_app)
            csrf_exempt_list = self.config["WTF_CSRF_EXEMPT_LIST"]
            for ex in csrf_exempt_list:
                csrf.exempt(ex) 

def setUp(self):
        self.app = Flask(__name__)
        self.app.testing = True
        self.app.secret_key = 'for test'
        dropzone = Dropzone(self.app)  # noqa
        csrf = CSRFProtect(self.app)  # noqa

        self.dropzone = _Dropzone

        @self.app.route('/upload')
        def upload():
            pass

        @self.app.route('/')
        def index():
            return render_template_string('''
                    {{ dropzone.load_css() }}\n{{ dropzone.create(action_view='upload') }}
                    {{ dropzone.load_js() }}\n{{ dropzone.config() }}''')

        @self.app.route('/load')
        def load():
            return render_template_string('''
                            {{ dropzone.load() }}\n{{ dropzone.create(action_view='upload') }}''')

        self.context = self.app.test_request_context()
        self.context.push()
        self.client = self.app.test_client()