mirror of
https://github.com/GSA/notifications-admin.git
synced 2026-02-05 19:03:30 -05:00
Add check that PDF file can be opened and is not malformed
This checks that the PDF file is not malformed in some way (e.g. by missing the EOF marker). We check this by trying to get the page count of the letter which will be needed to display the preview of the letter.
This commit is contained in:
@@ -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/<service_id>/preview-letter/<file_id>")
|
||||
@user_has_permissions('send_messages')
|
||||
def uploaded_letter_preview(service_id, file_id):
|
||||
|
||||
@@ -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',
|
||||
|
||||
BIN
tests/test_pdf_files/no_eof_marker.pdf
Normal file
BIN
tests/test_pdf_files/no_eof_marker.pdf
Normal file
Binary file not shown.
Reference in New Issue
Block a user