diff --git a/app/main/views/uploads.py b/app/main/views/uploads.py index f35ae95f9..a93fc7427 100644 --- a/app/main/views/uploads.py +++ b/app/main/views/uploads.py @@ -1,8 +1,10 @@ import uuid +from io import BytesIO -from flask import redirect, render_template, request, url_for +from flask import flash, redirect, render_template, request, url_for from app import current_service +from app.extensions import antivirus_client from app.main import main from app.main.forms import PDFUploadForm from app.utils import user_has_permissions @@ -20,6 +22,13 @@ def upload_letter(service_id): form = PDFUploadForm() if form.validate_on_submit(): + pdf_file_bytes = form.file.data.read() + + 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 + upload_id = uuid.uuid4() return redirect( diff --git a/tests/app/main/views/test_uploads.py b/tests/app/main/views/test_uploads.py index d02b65a9e..67e5999b0 100644 --- a/tests/app/main/views/test_uploads.py +++ b/tests/app/main/views/test_uploads.py @@ -1,5 +1,6 @@ from flask import url_for +from app.utils import normalize_spaces from tests.conftest import SERVICE_ONE_ID @@ -22,6 +23,7 @@ def test_get_upload_letter(client_request): def test_post_upload_letter_redirects_for_valid_file(mocker, client_request): mocker.patch('uuid.uuid4', return_value='fake-uuid') + antivirus_mock = mocker.patch('app.main.views.uploads.antivirus_client.scan', return_value=True) with open('tests/test_pdf_files/one_page_pdf.pdf', 'rb') as file: client_request.post( @@ -36,6 +38,7 @@ def test_post_upload_letter_redirects_for_valid_file(mocker, client_request): _external=True ) ) + assert antivirus_mock.called def test_post_upload_letter_shows_error_when_file_is_not_a_pdf(client_request): @@ -59,6 +62,20 @@ def test_post_upload_letter_shows_error_when_no_file_uploaded(client_request): assert page.find('span', class_='error-message').text.strip() == "You need to upload a file to submit" +def test_post_upload_letter_shows_error_when_file_contains_virus(mocker, client_request): + mocker.patch('app.main.views.uploads.antivirus_client.scan', return_value=False) + + with open('tests/test_pdf_files/one_page_pdf.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 has failed the virus check' + + def test_uploaded_letter_preview(client_request): page = client_request.get( 'main.uploaded_letter_preview',