diff --git a/keg_auth/core.py b/keg_auth/core.py index 66de692..7417402 100644 --- a/keg_auth/core.py +++ b/keg_auth/core.py @@ -1,3 +1,5 @@ +import socket + import arrow import flask import flask_login @@ -437,10 +439,36 @@ def clear_session(app, user): flask.session.clear() -def fix_session_cookies(app, **extra): - cookie_values = flask.request.cookies.getlist( +def is_ip(value: str) -> bool: + """Determine if the given string is an IP address. + + :param value: value to check + :type value: str + + :return: True if string is an IP address + :rtype: bool + + NOTE: This function was removed from Flask in 2.4 + """ + for family in (socket.AF_INET, socket.AF_INET6): + try: + socket.inet_pton(family, value) + except OSError: + pass + else: + return True + + return False + + +def get_cookie_values(app): + return flask.request.cookies.getlist( app.config.get('SESSION_COOKIE_NAME') ) + + +def fix_session_cookies(app, **extra): + cookie_values = get_cookie_values(app) server_name = app.config.get('SERVER_NAME') if len(cookie_values) > 1 and server_name: # werkzeug update has breaking session, since it matches both @@ -448,7 +476,7 @@ def fix_session_cookies(app, **extra): # chop off the port which is usually not supported by browsers cookie_domain = server_name.rsplit(':', 1)[0].lstrip('.') - if flask.helpers.is_ip(cookie_domain): + if is_ip(cookie_domain): return cookie_domain = '.' + cookie_domain diff --git a/keg_auth/tests/test_core.py b/keg_auth/tests/test_core.py index d689b6e..8118cda 100644 --- a/keg_auth/tests/test_core.py +++ b/keg_auth/tests/test_core.py @@ -4,7 +4,7 @@ import flask_login import freezegun import arrow -from keg_auth.core import update_last_login +from keg_auth.core import fix_session_cookies, update_last_login from keg_auth_ta.model import entities as ents @@ -34,3 +34,14 @@ def test_session_not_cleared_after_signing_out(self): assert 'foo' in list(flask.session) flask_login.logout_user() assert len(list(flask.session)) == 1 + + +def test_fix_session_cookies(): + with flask.current_app.test_request_context(): + with mock.patch( + 'keg_auth.core.get_cookie_values', autospec=True, spec_set=True + ) as m_cookies: + with mock.patch('flask.abort', autospec=True, spec_set=True) as m_abort: + m_cookies.return_value = ['FoObAr', 'BarFOo'] + fix_session_cookies(flask.current_app) + m_abort.assert_called_once()