|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | + |
| 3 | +import time |
| 4 | + |
| 5 | +from http import HTTPStatus |
| 6 | + |
| 7 | +from tests.common.db.accounts import UserFactory |
| 8 | +from warehouse.accounts.models import UniqueLoginStatus, UserUniqueLogin |
| 9 | +from warehouse.utils.otp import _get_totp |
| 10 | + |
| 11 | + |
| 12 | +def test_unrecognized_login_with_totp(webtest): |
| 13 | + """ |
| 14 | + Tests that a user with TOTP logging in from an unrecognized IP address |
| 15 | + is required to confirm their email address. |
| 16 | + """ |
| 17 | + |
| 18 | + # Arrange: Create a user with a verified email and TOTP enabled |
| 19 | + user = UserFactory.create( |
| 20 | + with_verified_primary_email=True, |
| 21 | + clear_pwd="password", |
| 22 | + ) |
| 23 | + |
| 24 | + # Act 1: Go to the login page |
| 25 | + login_page = webtest.get("/account/login/") |
| 26 | + assert login_page.status_code == HTTPStatus.OK |
| 27 | + |
| 28 | + # Fill out the login form and submit it |
| 29 | + login_form = login_page.forms["login-form"] |
| 30 | + login_form["username"] = user.username |
| 31 | + login_form["password"] = "password" |
| 32 | + resp = login_form.submit() |
| 33 | + |
| 34 | + # We should be redirected to the two-factor authentication page |
| 35 | + assert resp.status_code == HTTPStatus.SEE_OTHER |
| 36 | + assert resp.headers["Location"].startswith("http://localhost/account/two-factor/") |
| 37 | + two_factor_page = resp.follow() |
| 38 | + assert "Two-factor authentication" in two_factor_page |
| 39 | + |
| 40 | + # This is the first time we're logging in from this IP, so we should |
| 41 | + # be redirected to the "unrecognized login" page after submitting the |
| 42 | + # TOTP code. |
| 43 | + two_factor_form = two_factor_page.forms["totp-auth-form"] |
| 44 | + two_factor_form["totp_value"] = ( |
| 45 | + _get_totp(user.totp_secret).generate(time.time()).decode() |
| 46 | + ) |
| 47 | + resp = two_factor_form.submit() |
| 48 | + |
| 49 | + assert resp.status_code == HTTPStatus.SEE_OTHER |
| 50 | + assert resp.headers["Location"].endswith("/account/confirm-login/") |
| 51 | + unrecognized_page = resp.follow() |
| 52 | + assert "Unrecognized device" in unrecognized_page |
| 53 | + |
| 54 | + # This is a hack because the functional test doesn't have another way to |
| 55 | + # determine the magic link that was sent in the email. Instead, find the |
| 56 | + # UserUniqueLogin that was created for this user and manually confirm it |
| 57 | + db_session = webtest.extra_environ["warehouse.db_session"] |
| 58 | + unique_login = ( |
| 59 | + db_session.query(UserUniqueLogin).filter(UserUniqueLogin.user == user).one() |
| 60 | + ) |
| 61 | + assert unique_login.status == UniqueLoginStatus.PENDING |
| 62 | + unique_login.status = UniqueLoginStatus.CONFIRMED |
| 63 | + db_session.commit() |
| 64 | + |
| 65 | + # Act 2: Try to log in again |
| 66 | + login_page = webtest.get("/account/login/") |
| 67 | + login_form = login_page.forms["login-form"] |
| 68 | + login_form["username"] = user.username |
| 69 | + login_form["password"] = "password" |
| 70 | + resp = login_form.submit() |
| 71 | + |
| 72 | + # We should be redirected to the two-factor authentication page |
| 73 | + assert resp.status_code == HTTPStatus.SEE_OTHER |
| 74 | + assert resp.headers["Location"].startswith("http://localhost/account/two-factor/") |
| 75 | + two_factor_page = resp.follow() |
| 76 | + assert "Two-factor authentication" in two_factor_page |
| 77 | + |
| 78 | + # Fill out the TOTP form and submit it |
| 79 | + two_factor_form = two_factor_page.forms["totp-auth-form"] |
| 80 | + two_factor_form["totp_value"] = ( |
| 81 | + _get_totp(user.totp_secret).generate(time.time()).decode() |
| 82 | + ) |
| 83 | + |
| 84 | + # We should be able to successfully log in |
| 85 | + logged_in = two_factor_form.submit().follow(status=HTTPStatus.OK) |
| 86 | + assert logged_in.html.find("title", string="Warehouse · The Python Package Index") |
0 commit comments