From d655a38e148d5b955a233b52b11b1a47470f3cae Mon Sep 17 00:00:00 2001
From: Leo Hemsted <leo.hemsted@digital.cabinet-office.gov.uk>
Date: Wed, 12 Apr 2017 12:12:11 +0100
Subject: [PATCH] don't allow bad filetypes to preview

---
 app/main/views/send.py                 |  3 +++
 app/main/views/templates.py            |  3 +++
 tests/app/main/views/test_send.py      | 19 +++++++++++++++++++
 tests/app/main/views/test_templates.py | 18 ++++++++++++++++++
 4 files changed, 43 insertions(+)

diff --git a/app/main/views/send.py b/app/main/views/send.py
index 1e4f388eb..dd20ff41f 100644
--- a/app/main/views/send.py
+++ b/app/main/views/send.py
@@ -294,6 +294,9 @@ def check_messages(service_id, template_type, upload_id):
 @login_required
 @user_has_permissions('send_texts', 'send_emails', 'send_letters')
 def check_messages_preview(service_id, template_type, upload_id, filetype):
+    if filetype not in ('pdf', 'png'):
+        abort(404)
+
     template = _check_messages(
         service_id, template_type, upload_id, letters_as_pdf=True
     )['template']
diff --git a/app/main/views/templates.py b/app/main/views/templates.py
index 54826d5fb..866970d70 100644
--- a/app/main/views/templates.py
+++ b/app/main/views/templates.py
@@ -88,6 +88,9 @@ def view_template(service_id, template_id):
 @login_required
 @user_has_permissions('view_activity', admin_override=True)
 def view_letter_template_preview(service_id, template_id, filetype):
+    if filetype not in ('pdf', 'png'):
+        abort(404)
+
     db_template = service_api_client.get_service_template(service_id, template_id)['data']
     return TemplatePreview.from_database_object(db_template, filetype)
 
diff --git a/tests/app/main/views/test_send.py b/tests/app/main/views/test_send.py
index cea8680c7..f643cce6a 100644
--- a/tests/app/main/views/test_send.py
+++ b/tests/app/main/views/test_send.py
@@ -577,6 +577,25 @@ def test_should_show_preview_letter_message(
     assert mocked_preview.call_args[0][1] == filetype
 
 
+def test_dont_show_preview_letter_templates_for_bad_filetype(
+    logged_in_client,
+    mock_get_service_template,
+    service_one,
+    fake_uuid
+):
+    resp = logged_in_client.get(
+        url_for(
+            'main.check_messages_preview',
+            service_id=service_one['id'],
+            template_type='letter',
+            upload_id=fake_uuid,
+            filetype='blah'
+        )
+    )
+    assert resp.status_code == 404
+    assert mock_get_service_template.called is False
+
+
 def test_check_messages_should_revalidate_file_when_uploading_file(
     logged_in_client,
     service_one,
diff --git a/tests/app/main/views/test_templates.py b/tests/app/main/views/test_templates.py
index bc966ea54..065d51f23 100644
--- a/tests/app/main/views/test_templates.py
+++ b/tests/app/main/views/test_templates.py
@@ -185,6 +185,24 @@ def test_should_show_preview_letter_templates(
     assert mocked_preview.call_args[0][1] == filetype
 
 
+def test_dont_show_preview_letter_templates_for_bad_filetype(
+    logged_in_client,
+    mock_get_service_template,
+    service_one,
+    fake_uuid
+):
+    resp = logged_in_client.get(
+        url_for(
+            '.view_letter_template_preview',
+            service_id=service_one['id'],
+            template_id=fake_uuid,
+            filetype='blah'
+        )
+    )
+    assert resp.status_code == 404
+    assert mock_get_service_template.called is False
+
+
 def test_should_redirect_when_saving_a_template(
     logged_in_client,
     active_user_with_permissions,