Merge pull request #1802 from alphagov/rc_add_virus_scan_error_task

Added a new task to handle any error cases with the anti-virus

app/celery/letters_pdf_tasks.py

@@ -22,16 +22,16 @@ from app.dao.notifications_dao import (
     dao_update_notifications_by_reference,
 )
 from app.letters.utils import (
-    delete_pdf_from_letters_scan_bucket,
     get_reference_from_filename,
     move_scanned_pdf_to_test_or_live_pdf_bucket,
-    upload_letter_pdf
-)
+    upload_letter_pdf,
+    move_failed_pdf, ScanErrorType)
 from app.models import (
     KEY_TYPE_TEST,
     NOTIFICATION_CREATED,
     NOTIFICATION_DELIVERED,
     NOTIFICATION_VIRUS_SCAN_FAILED,
+    NOTIFICATION_TECHNICAL_FAILURE
 )
 
 
@@ -180,7 +180,7 @@ def process_virus_scan_passed(filename):
 @notify_celery.task(name='process-virus-scan-failed')
 def process_virus_scan_failed(filename):
     current_app.logger.exception('Virus scan failed: {}'.format(filename))
-    delete_pdf_from_letters_scan_bucket(filename)
+    move_failed_pdf(filename, ScanErrorType.FAILURE)
     reference = get_reference_from_filename(filename)
     updated_count = update_letter_pdf_status(reference, NOTIFICATION_VIRUS_SCAN_FAILED)
 
@@ -192,6 +192,21 @@ def process_virus_scan_failed(filename):
         )
 
 
+@notify_celery.task(name='process-virus-scan-error')
+def process_virus_scan_error(filename):
+    current_app.logger.exception('Virus scan error: {}'.format(filename))
+    move_failed_pdf(filename, ScanErrorType.ERROR)
+    reference = get_reference_from_filename(filename)
+    updated_count = update_letter_pdf_status(reference, NOTIFICATION_TECHNICAL_FAILURE)
+
+    if updated_count != 1:
+        raise Exception(
+            "There should only be one letter notification for each reference. Found {} notifications".format(
+                updated_count
+            )
+        )
+
+
 def update_letter_pdf_status(reference, status):
     return dao_update_notifications_by_reference(
         references=[reference],

app/letters/utils.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timedelta
+from enum import Enum
 
 import boto3
 from flask import current_app
@@ -9,6 +10,11 @@ from app.models import KEY_TYPE_TEST
 from app.variables import Retention
 
 
+class ScanErrorType(Enum):
+    ERROR = 1
+    FAILURE = 2
+
+
 LETTERS_PDF_FILE_LOCATION_STRUCTURE = \
     '{folder}NOTIFY.{reference}.{duplex}.{letter_class}.{colour}.{crown}.{date}.pdf'
 
@@ -85,36 +91,22 @@ def upload_letter_pdf(notification, pdf_data):
     return upload_file_name
 
 
-def move_scanned_pdf_to_test_or_live_pdf_bucket(filename, is_test_letter=False):
+def move_scanned_pdf_to_test_or_live_pdf_bucket(source_filename, is_test_letter=False):
     source_bucket_name = current_app.config['LETTERS_SCAN_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}
-    target_filename = get_folder_name(datetime.utcnow(), is_test_letter) + filename
+    target_filename = get_folder_name(datetime.utcnow(), is_test_letter) + source_filename
 
-    target_bucket = s3.Bucket(target_bucket_name)
-    obj = target_bucket.Object(target_filename)
-
-    # Tags are copied across but the expiration time is reset in the destination bucket
-    # e.g. if a file has 5 days left to expire on a ONE_WEEK retention in the source bucket,
-    # in the destination bucket the expiration time will be reset to 7 days left to expire
-    obj.copy(copy_source, ExtraArgs={'ServerSideEncryption': 'AES256'})
-
-    s3.Object(source_bucket_name, filename).delete()
-
-    current_app.logger.info("Moved letter PDF: {}/{} to {}/{}".format(
-        source_bucket_name, filename, target_bucket_name, target_filename))
+    _move_s3_object(source_bucket_name, source_filename, target_bucket_name, target_filename)
 
 
-def delete_pdf_from_letters_scan_bucket(filename):
-    bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME']
+def move_failed_pdf(source_filename, scan_error_type):
+    scan_bucket = current_app.config['LETTERS_SCAN_BUCKET_NAME']
 
-    s3 = boto3.resource('s3')
-    s3.Object(bucket_name, filename).delete()
+    target_filename = ('ERROR/' if scan_error_type == ScanErrorType.ERROR else 'FAILURE/') + source_filename
 
-    current_app.logger.info("Deleted letter PDF: {}/{}".format(bucket_name, filename))
+    _move_s3_object(scan_bucket, source_filename, scan_bucket, target_filename)
 
 
 def get_letter_pdf(notification):
@@ -135,3 +127,21 @@ def get_letter_pdf(notification):
         file_content = obj.get()["Body"].read()
 
     return file_content
+
+
+def _move_s3_object(source_bucket, source_filename, target_bucket, target_filename):
+    s3 = boto3.resource('s3')
+    copy_source = {'Bucket': source_bucket, 'Key': source_filename}
+
+    target_bucket = s3.Bucket(target_bucket)
+    obj = target_bucket.Object(target_filename)
+
+    # Tags are copied across but the expiration time is reset in the destination bucket
+    # e.g. if a file has 5 days left to expire on a ONE_WEEK retention in the source bucket,
+    # in the destination bucket the expiration time will be reset to 7 days left to expire
+    obj.copy(copy_source, ExtraArgs={'ServerSideEncryption': 'AES256'})
+
+    s3.Object(source_bucket, source_filename).delete()
+
+    current_app.logger.info("Moved letter PDF: {}/{} to {}/{}".format(
+        source_bucket, source_filename, target_bucket, target_filename))

tests/app/celery/test_letters_pdf_tasks.py

@@ -19,8 +19,8 @@ from app.celery.letters_pdf_tasks import (
     letter_in_created_state,
     process_virus_scan_passed,
     process_virus_scan_failed,
-)
-from app.letters.utils import get_letter_pdf_filename
+    process_virus_scan_error)
+from app.letters.utils import get_letter_pdf_filename, ScanErrorType
 from app.models import (
     KEY_TYPE_NORMAL,
     KEY_TYPE_TEST,
@@ -28,8 +28,8 @@ from app.models import (
     NOTIFICATION_CREATED,
     NOTIFICATION_DELIVERED,
     NOTIFICATION_VIRUS_SCAN_FAILED,
-    NOTIFICATION_SENDING
-)
+    NOTIFICATION_SENDING,
+    NOTIFICATION_TECHNICAL_FAILURE)
 
 from tests.conftest import set_config_values
 
@@ -341,9 +341,20 @@ def test_process_letter_task_check_virus_scan_passed(
 def test_process_letter_task_check_virus_scan_failed(sample_letter_notification, mocker):
     filename = 'NOTIFY.{}'.format(sample_letter_notification.reference)
     sample_letter_notification.status = 'pending-virus-check'
-    mock_delete_pdf = mocker.patch('app.celery.letters_pdf_tasks.delete_pdf_from_letters_scan_bucket')
+    mock_move_failed_pdf = mocker.patch('app.celery.letters_pdf_tasks.move_failed_pdf')
 
     process_virus_scan_failed(filename)
 
-    mock_delete_pdf.assert_called_once_with(filename)
+    mock_move_failed_pdf.assert_called_once_with(filename, ScanErrorType.FAILURE)
     assert sample_letter_notification.status == NOTIFICATION_VIRUS_SCAN_FAILED
+
+
+def test_process_letter_task_check_virus_scan_error(sample_letter_notification, mocker):
+    filename = 'NOTIFY.{}'.format(sample_letter_notification.reference)
+    sample_letter_notification.status = 'pending-virus-check'
+    mock_move_failed_pdf = mocker.patch('app.celery.letters_pdf_tasks.move_failed_pdf')
+
+    process_virus_scan_error(filename)
+
+    mock_move_failed_pdf.assert_called_once_with(filename, ScanErrorType.ERROR)
+    assert sample_letter_notification.status == NOTIFICATION_TECHNICAL_FAILURE

tests/app/letters/test_letter_utils.py

@@ -11,8 +11,8 @@ from app.letters.utils import (
     get_letter_pdf_filename,
     get_letter_pdf,
     upload_letter_pdf,
-    move_scanned_pdf_to_test_or_live_pdf_bucket
-)
+    move_scanned_pdf_to_test_or_live_pdf_bucket,
+    ScanErrorType, move_failed_pdf)
 from app.models import KEY_TYPE_NORMAL, KEY_TYPE_TEST, PRECOMPILED_TEMPLATE_NAME
 from app.variables import Retention
 
@@ -163,3 +163,43 @@ def test_move_scanned_letter_pdf_to_processing_bucket(
 
     assert folder_date_name + filename in [o.key for o in target_bucket.objects.all()]
     assert filename not in [o.key for o in source_bucket.objects.all()]
+
+
+@mock_s3
+@freeze_time(FROZEN_DATE_TIME)
+def test_move_failed_pdf_error(notify_api):
+    filename = 'test.pdf'
+    source_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME']
+    target_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME']
+
+    conn = boto3.resource('s3', region_name='eu-west-1')
+    source_bucket = conn.create_bucket(Bucket=source_bucket_name)
+    target_bucket = conn.create_bucket(Bucket=target_bucket_name)
+
+    s3 = boto3.client('s3', region_name='eu-west-1')
+    s3.put_object(Bucket=source_bucket_name, Key=filename, Body=b'pdf_content')
+
+    move_failed_pdf(filename, ScanErrorType.ERROR)
+
+    assert 'ERROR/' + filename in [o.key for o in target_bucket.objects.all()]
+    assert filename not in [o.key for o in source_bucket.objects.all()]
+
+
+@mock_s3
+@freeze_time(FROZEN_DATE_TIME)
+def test_move_failed_pdf_scan_failed(notify_api):
+    filename = 'test.pdf'
+    source_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME']
+    target_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME']
+
+    conn = boto3.resource('s3', region_name='eu-west-1')
+    source_bucket = conn.create_bucket(Bucket=source_bucket_name)
+    target_bucket = conn.create_bucket(Bucket=target_bucket_name)
+
+    s3 = boto3.client('s3', region_name='eu-west-1')
+    s3.put_object(Bucket=source_bucket_name, Key=filename, Body=b'pdf_content')
+
+    move_failed_pdf(filename, ScanErrorType.FAILURE)
+
+    assert 'FAILURE/' + filename in [o.key for o in target_bucket.objects.all()]
+    assert filename not in [o.key for o in source_bucket.objects.all()]