mirror of
https://github.com/GSA/notifications-admin.git
synced 2026-08-23 15:57:23 -04:00
Add virus scan stage when uploading a letter
This commit is contained in:
@@ -1,8 +1,10 @@
|
|||||||
import uuid
|
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 import current_service
|
||||||
|
from app.extensions import antivirus_client
|
||||||
from app.main import main
|
from app.main import main
|
||||||
from app.main.forms import PDFUploadForm
|
from app.main.forms import PDFUploadForm
|
||||||
from app.utils import user_has_permissions
|
from app.utils import user_has_permissions
|
||||||
@@ -20,6 +22,13 @@ def upload_letter(service_id):
|
|||||||
form = PDFUploadForm()
|
form = PDFUploadForm()
|
||||||
|
|
||||||
if form.validate_on_submit():
|
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()
|
upload_id = uuid.uuid4()
|
||||||
|
|
||||||
return redirect(
|
return redirect(
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from flask import url_for
|
from flask import url_for
|
||||||
|
|
||||||
|
from app.utils import normalize_spaces
|
||||||
from tests.conftest import SERVICE_ONE_ID
|
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):
|
def test_post_upload_letter_redirects_for_valid_file(mocker, client_request):
|
||||||
mocker.patch('uuid.uuid4', return_value='fake-uuid')
|
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:
|
with open('tests/test_pdf_files/one_page_pdf.pdf', 'rb') as file:
|
||||||
client_request.post(
|
client_request.post(
|
||||||
@@ -36,6 +38,7 @@ def test_post_upload_letter_redirects_for_valid_file(mocker, client_request):
|
|||||||
_external=True
|
_external=True
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
assert antivirus_mock.called
|
||||||
|
|
||||||
|
|
||||||
def test_post_upload_letter_shows_error_when_file_is_not_a_pdf(client_request):
|
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"
|
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):
|
def test_uploaded_letter_preview(client_request):
|
||||||
page = client_request.get(
|
page = client_request.get(
|
||||||
'main.uploaded_letter_preview',
|
'main.uploaded_letter_preview',
|
||||||
|
|||||||
Reference in New Issue
Block a user