From 19f731ec0754bfb8e36b61eff06c576afa59c802 Mon Sep 17 00:00:00 2001
From: Leo Hemsted <leo.hemsted@digital.cabinet-office.gov.uk>
Date: Wed, 1 Nov 2017 15:47:05 +0000
Subject: [PATCH] add error handler that catches invalid tokens, and returns
 404

---
 app/__init__.py                      | 10 +++++++++-
 tests/app/main/test_errorhandlers.py | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/app/__init__.py b/app/__init__.py
index 34aa8523e..ce425152b 100644
--- a/app/__init__.py
+++ b/app/__init__.py
@@ -5,6 +5,7 @@ from time import monotonic
 
 import itertools
 import ago
+from itsdangerous import BadSignature
 from flask import (
     Flask,
     session,
@@ -13,7 +14,8 @@ from flask import (
     current_app,
     request,
     g,
-    url_for
+    url_for,
+    flash
 )
 from flask._compat import string_types
 from flask.globals import _lookup_req_object, _request_ctx_stack
@@ -492,6 +494,12 @@ def register_errorhandlers(application):
             raise error
         return _error_response(500)
 
+    @application.errorhandler(BadSignature)
+    def handle_bad_token(error):
+        # if someone has a malformed token
+        flash('There’s something wrong with the link you’ve used.')
+        return _error_response(404)
+
 
 def setup_event_handlers():
     from flask_login import user_logged_in
diff --git a/tests/app/main/test_errorhandlers.py b/tests/app/main/test_errorhandlers.py
index f9185a1d1..92d712259 100644
--- a/tests/app/main/test_errorhandlers.py
+++ b/tests/app/main/test_errorhandlers.py
@@ -1,3 +1,4 @@
+import pytest
 from bs4 import BeautifulSoup
 
 
@@ -6,3 +7,19 @@ def test_bad_url_returns_page_not_found(client):
     assert response.status_code == 404
     page = BeautifulSoup(response.data.decode('utf-8'), 'html.parser')
     assert page.h1.string.strip() == 'Page could not be found'
+
+
+@pytest.mark.parametrize('url', [
+    '/invitation/MALFORMED_TOKEN',
+    '/new-password/MALFORMED_TOKEN',
+    '/user-profile/email/confirm/MALFORMED_TOKEN',
+    '/verify-email/MALFORMED_TOKEN'
+])
+def test_malformed_token_returns_page_not_found(client, url):
+    response = client.get(url)
+
+    assert response.status_code == 404
+    page = BeautifulSoup(response.data.decode('utf-8'), 'html.parser')
+    assert page.h1.string.strip() == 'Page could not be found'
+    flash_banner = page.find('div', class_='banner-dangerous').string.strip()
+    assert flash_banner == "There’s something wrong with the link you’ve used."