From d47f9933d2fb7317ce8155bf4211dc10327f4de7 Mon Sep 17 00:00:00 2001
From: Kenneth Kehl <@kkehl@flexion.us>
Date: Mon, 13 May 2024 13:39:51 -0700
Subject: [PATCH] add check for government email addresses

---
 app/main/views/register.py            | 20 +++++++++++++++++
 app/main/views/sign_in.py             | 31 ++++++++++++++++++++++++++-
 tests/app/main/views/test_register.py | 20 +++++++++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/app/main/views/register.py b/app/main/views/register.py
index 4638eaeab..1910c5c5b 100644
--- a/app/main/views/register.py
+++ b/app/main/views/register.py
@@ -146,6 +146,26 @@ def check_invited_user_email_address_matches_expected(
         debug_msg("invited user email did not match expected email, abort(403)")
         flash("You cannot accept an invite for another person.")
         abort(403)
+    check_for_gov_email_address(user_email)
+
+
+def check_for_gov_email_address(user_email):
+    # We could try to check that it is a government email at the time the invite is
+    # sent, but due to the way login.gov allows multiple emails, it would not be effective.
+    # We track the login.gov user by their uuid, so if they have a login.gov account
+    # with a .gov email address and a .com email address, the .com address will work without
+    # having a check here.
+    if (
+        user_email.lower().endswith(".gov")
+        or user_email.lower().endswith(".mil")
+        or user_email.lower().endswith(".si.edu")
+    ):
+        # everything is good, proceed
+        pass
+    else:
+        debug_msg("invited user has a non-government email address.")
+        flash("You must use a government email address.")
+        abort(403)
 
 
 @main.route("/set-up-your-profile", methods=["GET", "POST"])
diff --git a/app/main/views/sign_in.py b/app/main/views/sign_in.py
index 618a35654..23a9cefd4 100644
--- a/app/main/views/sign_in.py
+++ b/app/main/views/sign_in.py
@@ -4,7 +4,16 @@ import uuid
 
 import jwt
 import requests
-from flask import Response, current_app, redirect, render_template, request, url_for
+from flask import (
+    Response,
+    abort,
+    current_app,
+    flash,
+    redirect,
+    render_template,
+    request,
+    url_for,
+)
 from flask_login import current_user
 from notifications_utils.url_safe_token import generate_token
 
@@ -88,6 +97,7 @@ def _do_login_dot_gov():
         try:
             access_token = _get_access_token(code, state)
             user_email, user_uuid = _get_user_email_and_uuid(access_token)
+            check_for_gov_email_address(user_email)
             redirect_url = request.args.get("next")
             user = user_api_client.get_user_by_uuid_or_email(user_uuid, user_email)
 
@@ -171,3 +181,22 @@ def sign_in():
 @login_manager.unauthorized_handler
 def sign_in_again():
     return redirect(url_for("main.sign_in", next=request.path))
+
+
+def check_for_gov_email_address(user_email):
+    # We could try to check that it is a government email at the time the invite is
+    # sent, but due to the way login.gov allows multiple emails, it would not be effective.
+    # We track the login.gov user by their uuid, so if they have a login.gov account
+    # with a .gov email address and a .com email address, the .com address will work without
+    # having a check here.
+    if (
+        user_email.lower().endswith(".gov")
+        or user_email.lower().endswith(".mil")
+        or user_email.lower().endswith(".si.edu")
+    ):
+        # everything is good, proceed
+        pass
+    else:
+        current_app.logger.warning("invited user has a non-government email address.")
+        flash("You must use a government email address.")
+        abort(403)
diff --git a/tests/app/main/views/test_register.py b/tests/app/main/views/test_register.py
index 688ab1623..194ef000d 100644
--- a/tests/app/main/views/test_register.py
+++ b/tests/app/main/views/test_register.py
@@ -401,6 +401,26 @@ def test_check_invited_user_email_address_doesnt_match_expected(mocker):
     mock_abort.assert_called_once_with(403)
 
 
+def test_check_user_email_address_fails_if_not_government_address(mocker):
+    mock_flash = mocker.patch("app.main.views.register.flash")
+    mock_abort = mocker.patch("app.main.views.register.abort")
+
+    check_invited_user_email_address_matches_expected(
+        "fake@fake.bogus", "Fake@Fake.BOGUS"
+    )
+    mock_flash.assert_called_once_with("You must use a government email address.")
+    mock_abort.assert_called_once_with(403)
+
+
+def test_check_user_email_address_succeeds_if_government_address(mocker):
+    mock_flash = mocker.patch("app.main.views.register.flash")
+    mock_abort = mocker.patch("app.main.views.register.abort")
+
+    check_invited_user_email_address_matches_expected("fake@fake.mil", "Fake@Fake.MIL")
+    mock_flash.assert_not_called()
+    mock_abort.assert_not_called()
+
+
 def decode_invite_data(state):
     state = state.encode("utf8")
     state = base64.b64decode(state)