diff --git a/app/main/views/uploads.py b/app/main/views/uploads.py index 052329766..6a2315a90 100644 --- a/app/main/views/uploads.py +++ b/app/main/views/uploads.py @@ -1,7 +1,16 @@ import uuid from io import BytesIO -from flask import flash, redirect, render_template, request, url_for +from flask import ( + current_app, + flash, + redirect, + render_template, + request, + url_for, +) +from notifications_utils.pdf import pdf_page_count +from PyPDF2.utils import PdfReadError from app import current_service from app.extensions import antivirus_client @@ -28,12 +37,16 @@ def upload_letter(service_id): virus_free = antivirus_client.scan(BytesIO(pdf_file_bytes)) if not virus_free: - flash('Your file has failed the virus check', 'dangerous') - return render_template('views/uploads/choose-file.html', form=form), 400 + return invalid_upload_error('Your file has failed the virus check') if len(pdf_file_bytes) > MAX_FILE_UPLOAD_SIZE: - flash('Your file must be smaller than 2MB', 'dangerous') - return render_template('views/uploads/choose-file.html', form=form), 400 + return invalid_upload_error('Your file must be smaller than 2MB') + + try: + pdf_page_count(BytesIO(pdf_file_bytes)) + except PdfReadError: + current_app.logger.info('Invalid PDF uploaded for service_id: {}'.format(service_id)) + return invalid_upload_error('Your file must be a valid PDF') upload_id = uuid.uuid4() @@ -49,6 +62,11 @@ def upload_letter(service_id): return render_template('views/uploads/choose-file.html', form=form) +def invalid_upload_error(message): + flash(message, 'dangerous') + return render_template('views/uploads/choose-file.html', form=PDFUploadForm()), 400 + + @main.route("/services//preview-letter/") @user_has_permissions('send_messages') def uploaded_letter_preview(service_id, file_id): diff --git a/tests/app/main/views/test_uploads.py b/tests/app/main/views/test_uploads.py index 530130e74..d8daf1398 100644 --- a/tests/app/main/views/test_uploads.py +++ b/tests/app/main/views/test_uploads.py @@ -90,6 +90,20 @@ def test_post_choose_upload_file_when_file_is_too_big(mocker, client_request): assert normalize_spaces(page.select('.banner-dangerous')[0].text) == 'Your file must be smaller than 2MB' +def test_post_choose_upload_file_when_file_is_malformed(mocker, client_request): + mocker.patch('app.main.views.uploads.antivirus_client.scan', return_value=True) + + with open('tests/test_pdf_files/no_eof_marker.pdf', 'rb') as file: + page = client_request.post( + 'main.upload_letter', + service_id=SERVICE_ONE_ID, + _data={'file': file}, + _expected_status=400 + ) + assert page.find('h1').text == 'Upload a letter' + assert normalize_spaces(page.select('.banner-dangerous')[0].text) == 'Your file must be a valid PDF' + + def test_uploaded_letter_preview(client_request): page = client_request.get( 'main.uploaded_letter_preview', diff --git a/tests/test_pdf_files/no_eof_marker.pdf b/tests/test_pdf_files/no_eof_marker.pdf new file mode 100644 index 000000000..857ae1fbf Binary files /dev/null and b/tests/test_pdf_files/no_eof_marker.pdf differ