From a8b916b57ff249827343f553aa52a7b53535570b Mon Sep 17 00:00:00 2001
From: Chris Hill-Scott <me@quis.cc>
Date: Wed, 12 Dec 2018 13:10:46 +0000
Subject: [PATCH] Refactor gov user check into a decorator

We quite often use it in the same way as `@user_has_permissions`.
---
 app/main/views/add_service.py      |  9 +++------
 app/main/views/choose_account.py   |  1 -
 app/main/views/service_settings.py |  5 ++---
 app/main/views/user_profile.py     | 15 +++------------
 app/models/user.py                 |  1 -
 app/utils.py                       |  9 +++++++++
 6 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/app/main/views/add_service.py b/app/main/views/add_service.py
index d959fb798..4466114e2 100644
--- a/app/main/views/add_service.py
+++ b/app/main/views/add_service.py
@@ -1,7 +1,6 @@
 from flask import current_app, redirect, render_template, session, url_for
-from flask_login import current_user, login_required
+from flask_login import login_required
 from notifications_python_client.errors import HTTPError
-from werkzeug.exceptions import abort
 
 from app import (
     billing_api_client,
@@ -13,7 +12,7 @@ from app import (
 from app.main import main
 from app.main.forms import CreateServiceForm
 from app.models.user import InvitedUser
-from app.utils import AgreementInfo, email_safe, is_gov_user
+from app.utils import AgreementInfo, email_safe, user_is_gov_user
 
 
 def _add_invited_user_to_service(invited_user):
@@ -69,15 +68,13 @@ def _create_example_template(service_id):
 
 @main.route("/add-service", methods=['GET', 'POST'])
 @login_required
+@user_is_gov_user
 def add_service():
     invited_user = session.get('invited_user')
     if invited_user:
         service_id = _add_invited_user_to_service(invited_user)
         return redirect(url_for('main.service_dashboard', service_id=service_id))
 
-    if not current_user.is_gov_user:
-        abort(403)
-
     form = CreateServiceForm()
     heading = 'About your service'
 
diff --git a/app/main/views/choose_account.py b/app/main/views/choose_account.py
index 30a035bfa..82d6d96ab 100644
--- a/app/main/views/choose_account.py
+++ b/app/main/views/choose_account.py
@@ -4,7 +4,6 @@ from werkzeug.routing import RequestRedirect
 
 from app import user_api_client
 from app.main import main
-from app.utils import is_gov_user
 
 
 @main.route("/services")
diff --git a/app/main/views/service_settings.py b/app/main/views/service_settings.py
index 73201f079..6df92b023 100644
--- a/app/main/views/service_settings.py
+++ b/app/main/views/service_settings.py
@@ -56,6 +56,7 @@ from app.utils import (
     email_safe,
     get_logo_cdn_domain,
     user_has_permissions,
+    user_is_gov_user,
     user_is_platform_admin,
 )
 
@@ -146,11 +147,9 @@ def request_to_go_live(service_id):
 @main.route("/services/<service_id>/service-settings/submit-request-to-go-live", methods=['GET', 'POST'])
 @login_required
 @user_has_permissions('manage_service')
+@user_is_gov_user
 def submit_request_to_go_live(service_id):
 
-    if not current_user.is_gov_user:
-        abort(403)
-
     form = RequestToGoLiveForm()
 
     if form.validate_on_submit():
diff --git a/app/main/views/user_profile.py b/app/main/views/user_profile.py
index f9a50810a..a2caf8d2e 100644
--- a/app/main/views/user_profile.py
+++ b/app/main/views/user_profile.py
@@ -1,13 +1,6 @@
 import json
 
-from flask import (
-    abort,
-    current_app,
-    redirect,
-    render_template,
-    session,
-    url_for,
-)
+from flask import current_app, redirect, render_template, session, url_for
 from flask_login import current_user, login_required
 from notifications_utils.url_safe_token import check_token
 
@@ -21,7 +14,7 @@ from app.main.forms import (
     ConfirmPasswordForm,
     TwoFactorForm,
 )
-from app.utils import is_gov_user
+from app.utils import user_is_gov_user
 
 NEW_EMAIL = 'new-email'
 NEW_MOBILE = 'new-mob'
@@ -56,11 +49,9 @@ def user_profile_name():
 
 @main.route("/user-profile/email", methods=['GET', 'POST'])
 @login_required
+@user_is_gov_user
 def user_profile_email():
 
-    if not current_user.is_gov_user:
-        abort(403)
-
     def _is_email_already_in_use(email):
         return user_api_client.is_email_already_in_use(email)
     form = ChangeEmailForm(_is_email_already_in_use,
diff --git a/app/models/user.py b/app/models/user.py
index c1226b2ff..7b9ffd6ec 100644
--- a/app/models/user.py
+++ b/app/models/user.py
@@ -5,7 +5,6 @@ from flask_login import AnonymousUserMixin, UserMixin
 
 from app.utils import is_gov_user
 
-
 roles = {
     'send_messages': ['send_texts', 'send_emails', 'send_letters'],
     'manage_templates': ['manage_templates'],
diff --git a/app/utils.py b/app/utils.py
index a9536787e..7e20d7680 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -62,6 +62,15 @@ def user_has_permissions(*permissions, **permission_kwargs):
     return wrap
 
 
+def user_is_gov_user(f):
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        if not current_user.is_gov_user:
+            abort(403)
+        return f(*args, **kwargs)
+    return wrapped
+
+
 def user_is_platform_admin(f):
     @wraps(f)
     def wrapped(*args, **kwargs):