From 0ee5c330840b5c28c8d39acd4fbc72fe7149a62a Mon Sep 17 00:00:00 2001 From: Ken Tsang Date: Fri, 23 Mar 2018 12:04:37 +0000 Subject: [PATCH] Add antivirus check on precompiled letters sent with test key - precompiled PDFs sent by test key uploaded to scan bucket - set status to VIRUS-SCAN-FAILED for pdfs failing virus scan rather than PERMANENT-FAILURE - Make call to AV app for precompiled letters sent via a test key, and set notification status to PENDING-VIRUS-SCAN --- app/celery/letters_pdf_tasks.py | 30 ++++++++++++------- app/letters/utils.py | 22 +++++++------- app/v2/notifications/post_notifications.py | 14 +++++++-- tests/app/celery/test_letters_pdf_tasks.py | 22 ++++++++++---- tests/app/letters/test_letter_utils.py | 18 +++++++---- .../test_post_letter_notifications.py | 14 ++++++--- 6 files changed, 78 insertions(+), 42 deletions(-) diff --git a/app/celery/letters_pdf_tasks.py b/app/celery/letters_pdf_tasks.py index 2d13019de..c526848ea 100644 --- a/app/celery/letters_pdf_tasks.py +++ b/app/celery/letters_pdf_tasks.py @@ -17,16 +17,22 @@ from app.dao.notifications_dao import ( get_notification_by_id, update_notification_status_by_id, dao_update_notification, + dao_get_notification_by_reference, dao_get_notifications_by_references, dao_update_notifications_by_reference, ) from app.letters.utils import ( delete_pdf_from_letters_scan_bucket, get_reference_from_filename, - move_scanned_pdf_to_letters_pdf_bucket, + move_scanned_pdf_to_test_or_live_pdf_bucket, upload_letter_pdf ) -from app.models import NOTIFICATION_CREATED, NOTIFICATION_PERMANENT_FAILURE +from app.models import ( + KEY_TYPE_TEST, + NOTIFICATION_CREATED, + NOTIFICATION_DELIVERED, + NOTIFICATION_VIRUS_SCAN_FAILED, +) @notify_celery.task(bind=True, name="create-letters-pdf", max_retries=15, default_retry_delay=300) @@ -157,16 +163,18 @@ def letter_in_created_state(filename): @notify_celery.task(name='process-virus-scan-passed') def process_virus_scan_passed(filename): current_app.logger.info('Virus scan passed: {}'.format(filename)) - move_scanned_pdf_to_letters_pdf_bucket(filename) reference = get_reference_from_filename(filename) - updated_count = update_letter_pdf_status(reference, NOTIFICATION_CREATED) + notification = dao_get_notification_by_reference(reference) - if updated_count != 1: - raise Exception( - "There should only be one letter notification for each reference. Found {} notifications".format( - updated_count - ) - ) + is_test_key = notification.key_type == KEY_TYPE_TEST + move_scanned_pdf_to_test_or_live_pdf_bucket( + filename, + is_test_letter=is_test_key + ) + update_letter_pdf_status( + reference, + NOTIFICATION_DELIVERED if is_test_key else NOTIFICATION_CREATED + ) @notify_celery.task(name='process-virus-scan-failed') @@ -174,7 +182,7 @@ def process_virus_scan_failed(filename): current_app.logger.error('Virus scan failed: {}'.format(filename)) delete_pdf_from_letters_scan_bucket(filename) reference = get_reference_from_filename(filename) - updated_count = update_letter_pdf_status(reference, NOTIFICATION_PERMANENT_FAILURE) + updated_count = update_letter_pdf_status(reference, NOTIFICATION_VIRUS_SCAN_FAILED) if updated_count != 1: raise Exception( diff --git a/app/letters/utils.py b/app/letters/utils.py index 83af29b86..c79a09a8e 100644 --- a/app/letters/utils.py +++ b/app/letters/utils.py @@ -26,11 +26,11 @@ def get_folder_name(_now, is_test_or_scan_letter=False): return folder_name -def get_letter_pdf_filename(reference, crown, is_test_or_scan_letter=False): +def get_letter_pdf_filename(reference, crown, is_scan_letter=False): now = datetime.utcnow() upload_file_name = LETTERS_PDF_FILE_LOCATION_STRUCTURE.format( - folder=get_folder_name(now, is_test_or_scan_letter), + folder=get_folder_name(now, is_scan_letter), reference=reference, duplex="D", letter_class="2", @@ -58,22 +58,19 @@ def get_reference_from_filename(filename): return filename_parts[1] -def upload_letter_pdf(notification, pdf_data, is_test_letter=False): +def upload_letter_pdf(notification, pdf_data): current_app.logger.info("PDF Letter {} reference {} created at {}, {} bytes".format( notification.id, notification.reference, notification.created_at, len(pdf_data))) upload_file_name = get_letter_pdf_filename( notification.reference, notification.service.crown, - is_test_or_scan_letter=is_test_letter or notification.template.is_precompiled_letter) + is_scan_letter=notification.template.is_precompiled_letter) - if is_test_letter: - bucket_name = current_app.config['TEST_LETTERS_BUCKET_NAME'] + if notification.template.is_precompiled_letter: + bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME'] else: - if notification.template.is_precompiled_letter: - bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME'] - else: - bucket_name = current_app.config['LETTERS_PDF_BUCKET_NAME'] + bucket_name = current_app.config['LETTERS_PDF_BUCKET_NAME'] s3upload( filedata=pdf_data, @@ -88,9 +85,10 @@ def upload_letter_pdf(notification, pdf_data, is_test_letter=False): return upload_file_name -def move_scanned_pdf_to_letters_pdf_bucket(filename): +def move_scanned_pdf_to_test_or_live_pdf_bucket(filename, is_test_letter=False): source_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME'] - target_bucket_name = current_app.config['LETTERS_PDF_BUCKET_NAME'] + target_bucket_config = 'TEST_LETTERS_BUCKET_NAME' if is_test_letter else 'LETTERS_PDF_BUCKET_NAME' + target_bucket_name = current_app.config[target_bucket_config] s3 = boto3.resource('s3') copy_source = {'Bucket': source_bucket_name, 'Key': filename} diff --git a/app/v2/notifications/post_notifications.py b/app/v2/notifications/post_notifications.py index 5a3bf6e41..7e39ad116 100644 --- a/app/v2/notifications/post_notifications.py +++ b/app/v2/notifications/post_notifications.py @@ -227,7 +227,7 @@ def process_letter_notification(*, letter_data, api_key, template, reply_to_text if precompiled: try: - if should_send: + if should_send or (precompiled and api_key.key_type == KEY_TYPE_TEST): status = NOTIFICATION_PENDING_VIRUS_CHECK letter_content = base64.b64decode(letter_data['content']) pages = pdf_page_count(io.BytesIO(letter_content)) @@ -273,8 +273,16 @@ def process_letter_notification(*, letter_data, api_key, template, reply_to_text ) else: if precompiled and api_key.key_type == KEY_TYPE_TEST: - upload_letter_pdf(notification, letter_content, is_test_letter=True) - update_notification_status_by_reference(notification.reference, NOTIFICATION_DELIVERED) + filename = upload_letter_pdf(notification, letter_content) + + # call task to add the filename to anti virus queue + notify_celery.send_task( + name=TaskNames.SCAN_FILE, + kwargs={'filename': filename}, + queue=QueueNames.ANTIVIRUS, + ) + else: + update_notification_status_by_reference(notification.reference, NOTIFICATION_DELIVERED) return notification diff --git a/tests/app/celery/test_letters_pdf_tasks.py b/tests/app/celery/test_letters_pdf_tasks.py index b13eeac8f..e492540ed 100644 --- a/tests/app/celery/test_letters_pdf_tasks.py +++ b/tests/app/celery/test_letters_pdf_tasks.py @@ -22,9 +22,12 @@ from app.celery.letters_pdf_tasks import ( ) from app.letters.utils import get_letter_pdf_filename from app.models import ( + KEY_TYPE_NORMAL, + KEY_TYPE_TEST, Notification, NOTIFICATION_CREATED, - NOTIFICATION_PERMANENT_FAILURE, + NOTIFICATION_DELIVERED, + NOTIFICATION_VIRUS_SCAN_FAILED, NOTIFICATION_SENDING ) @@ -317,15 +320,22 @@ def test_letter_in_created_state_fails_if_notification_doesnt_exist(sample_notif assert letter_in_created_state(filename) is False -def test_process_letter_task_check_virus_scan_passed(sample_letter_notification, mocker): +@pytest.mark.parametrize('key_type,is_test_letter,noti_status', [ + (KEY_TYPE_NORMAL, False, NOTIFICATION_CREATED), + (KEY_TYPE_TEST, True, NOTIFICATION_DELIVERED) +]) +def test_process_letter_task_check_virus_scan_passed( + sample_letter_notification, mocker, key_type, is_test_letter, noti_status +): filename = 'NOTIFY.{}'.format(sample_letter_notification.reference) sample_letter_notification.status = 'pending-virus-check' - mock_move_pdf = mocker.patch('app.celery.letters_pdf_tasks.move_scanned_pdf_to_letters_pdf_bucket') + sample_letter_notification.key_type = key_type + mock_move_pdf = mocker.patch('app.celery.letters_pdf_tasks.move_scanned_pdf_to_test_or_live_pdf_bucket') process_virus_scan_passed(filename) - mock_move_pdf.assert_called_once_with(filename) - assert sample_letter_notification.status == NOTIFICATION_CREATED + mock_move_pdf.assert_called_once_with(filename, is_test_letter=is_test_letter) + assert sample_letter_notification.status == noti_status def test_process_letter_task_check_virus_scan_failed(sample_letter_notification, mocker): @@ -336,4 +346,4 @@ def test_process_letter_task_check_virus_scan_failed(sample_letter_notification, process_virus_scan_failed(filename) mock_delete_pdf.assert_called_once_with(filename) - assert sample_letter_notification.status == NOTIFICATION_PERMANENT_FAILURE + assert sample_letter_notification.status == NOTIFICATION_VIRUS_SCAN_FAILED diff --git a/tests/app/letters/test_letter_utils.py b/tests/app/letters/test_letter_utils.py index 981f97ddf..2a071bb0c 100644 --- a/tests/app/letters/test_letter_utils.py +++ b/tests/app/letters/test_letter_utils.py @@ -11,7 +11,7 @@ from app.letters.utils import ( get_letter_pdf_filename, get_letter_pdf, upload_letter_pdf, - move_scanned_pdf_to_letters_pdf_bucket + move_scanned_pdf_to_test_or_live_pdf_bucket ) from app.models import KEY_TYPE_NORMAL, KEY_TYPE_TEST, PRECOMPILED_TEMPLATE_NAME from app.variables import Retention @@ -71,7 +71,7 @@ def test_get_letter_pdf_filename_returns_correct_filename( @freeze_time("2017-12-04 17:29:00") def test_get_letter_pdf_filename_returns_correct_filename_for_test_letters( notify_api, mocker): - filename = get_letter_pdf_filename(reference='foo', crown='C', is_test_or_scan_letter=True) + filename = get_letter_pdf_filename(reference='foo', crown='C', is_scan_letter=True) assert filename == 'NOTIFY.FOO.D.2.C.C.20171204172900.PDF' @@ -125,7 +125,7 @@ def test_upload_letter_pdf_to_correct_bucket( filename = get_letter_pdf_filename( reference=sample_letter_notification.reference, crown=sample_letter_notification.service.crown, - is_test_or_scan_letter=is_precompiled_letter + is_scan_letter=is_precompiled_letter ) upload_letter_pdf(sample_letter_notification, b'\x00\x01') @@ -140,11 +140,17 @@ def test_upload_letter_pdf_to_correct_bucket( @mock_s3 +@pytest.mark.parametrize('is_test_letter,bucket_config_name', [ + (False, 'LETTERS_PDF_BUCKET_NAME'), + (True, 'TEST_LETTERS_BUCKET_NAME') +]) @freeze_time(FROZEN_DATE_TIME) -def test_move_scanned_letter_pdf_to_processing_bucket(notify_api): +def test_move_scanned_letter_pdf_to_processing_bucket( + notify_api, is_test_letter, bucket_config_name +): filename = 'test.pdf' source_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME'] - target_bucket_name = current_app.config['LETTERS_PDF_BUCKET_NAME'] + target_bucket_name = current_app.config[bucket_config_name] conn = boto3.resource('s3', region_name='eu-west-1') source_bucket = conn.create_bucket(Bucket=source_bucket_name) @@ -153,7 +159,7 @@ def test_move_scanned_letter_pdf_to_processing_bucket(notify_api): s3 = boto3.client('s3', region_name='eu-west-1') s3.put_object(Bucket=source_bucket_name, Key=filename, Body=b'pdf_content') - move_scanned_pdf_to_letters_pdf_bucket(filename) + move_scanned_pdf_to_test_or_live_pdf_bucket(filename, is_test_letter=is_test_letter) assert '2018-03-14/' + filename in [o.key for o in target_bucket.objects.all()] assert filename not in [o.key for o in source_bucket.objects.all()] diff --git a/tests/app/v2/notifications/test_post_letter_notifications.py b/tests/app/v2/notifications/test_post_letter_notifications.py index c73094c68..851cf2e3c 100644 --- a/tests/app/v2/notifications/test_post_letter_notifications.py +++ b/tests/app/v2/notifications/test_post_letter_notifications.py @@ -5,7 +5,7 @@ from flask import json from flask import url_for import pytest -from app.config import QueueNames +from app.config import TaskNames, QueueNames from app.models import ( Job, Notification, @@ -399,8 +399,9 @@ def test_post_letter_notification_is_delivered_and_has_pdf_uploaded_to_test_lett mocker ): sample_letter_service = create_service(service_permissions=['letter', 'precompiled_letter']) - s3mock = mocker.patch('app.v2.notifications.post_notifications.upload_letter_pdf') + s3mock = mocker.patch('app.v2.notifications.post_notifications.upload_letter_pdf', return_value='test.pdf') mocker.patch('app.v2.notifications.post_notifications.pdf_page_count', return_value=1) + mock_celery = mocker.patch("app.letters.rest.notify_celery.send_task") data = { "reference": "letter-reference", "content": "bGV0dGVyLWNvbnRlbnQ=" @@ -413,8 +414,13 @@ def test_post_letter_notification_is_delivered_and_has_pdf_uploaded_to_test_lett precompiled=True) notification = Notification.query.one() - assert notification.status == NOTIFICATION_DELIVERED - s3mock.assert_called_once_with(ANY, b'letter-content', is_test_letter=True) + assert notification.status == NOTIFICATION_PENDING_VIRUS_CHECK + s3mock.assert_called_once_with(ANY, b'letter-content') + mock_celery.assert_called_once_with( + name=TaskNames.SCAN_FILE, + kwargs={'filename': 'test.pdf'}, + queue=QueueNames.ANTIVIRUS + ) def test_post_letter_notification_persists_notification_reply_to_text(