diff --git a/assets/lambda/code/scan/lambda.py b/assets/lambda/code/scan/lambda.py index 761e8138..6e6cccb6 100644 --- a/assets/lambda/code/scan/lambda.py +++ b/assets/lambda/code/scan/lambda.py @@ -4,7 +4,6 @@ import glob import json import os -import pwd import shutil import subprocess import time @@ -13,6 +12,8 @@ import boto3 import botocore + +# pyrefly: ignore [missing-import] from aws_lambda_powertools import Logger, Metrics logger = Logger() @@ -175,8 +176,21 @@ def set_status(bucket, key, status, version_id=None): if version_id: tagging_args["VersionId"] = version_id - s3_client.put_object_tagging(**tagging_args) - metrics.add_metric(name=status, unit="Count", value=1) + try: + s3_client.put_object_tagging(**tagging_args) + metrics.add_metric(name=status, unit="Count", value=1) + except botocore.exceptions.ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "") + if error_code in ("NoSuchKey", "404"): + logger.info( + f"Could not set scan-status for {key}" + f"{' version ' + version_id if version_id else ''} " + "because the object no longer exists (likely deleted mid-scan)" + ) + else: + logger.warning( + "Error setting status: %s", e.response["Error"]["Message"] + ) def get_status(bucket_name, object_key, version_id=None): @@ -357,7 +371,7 @@ def freshclam_update(input_bucket, input_key, download_path, definitions_path): f"--config-file={conf}", "--stdout", "-u", - f"{pwd.getpwuid(os.getuid()).pw_name}", + f"{__import__('pwd').getpwuid(os.getuid()).pw_name}", f"--datadir={definitions_path}", ] update_summary = subprocess.run( diff --git a/assets/lambda/code/scan/test_lambda.py b/assets/lambda/code/scan/test_lambda.py index 8c0594a6..c0956cfd 100644 --- a/assets/lambda/code/scan/test_lambda.py +++ b/assets/lambda/code/scan/test_lambda.py @@ -182,6 +182,41 @@ def test_set_status_unversioned(s3_stubber): assert_that(s3_stubber).has_no_pending_responses() +def test_set_status_put_tagging_nosuchkey(s3_stubber): + """If the object is deleted mid-scan, put_object_tagging raises NoSuchKey. + set_status should handle this gracefully instead of raising.""" + # Arrange + bucket = "bucket" + key = "key" + status = scan_lambda.ScanStatus.CLEAN + + # Stub the get_object_tagging call (fetching old tags) - succeeds + s3_stubber.add_response( + "get_object_tagging", {"TagSet": []}, {"Bucket": bucket, "Key": key} + ) + + # Stub the put_object_tagging call to fail because the object + # was deleted between the scan finishing and the tagging attempt + expected_tags = {"TagSet": [{"Key": "scan-status", "Value": status}]} + s3_stubber.add_client_error( + "put_object_tagging", + service_error_code="NoSuchKey", + service_message="The specified key does not exist.", + http_status_code=404, + expected_params={ + "Bucket": bucket, + "Key": key, + "Tagging": expected_tags, + }, + ) + + # Act & Assert - should not raise + scan_lambda.set_status(bucket, key, status) + + # Assert + assert_that(s3_stubber).has_no_pending_responses() + + def test_get_status_versioned_found(s3_stubber): # Arrange expected_status = scan_lambda.ScanStatus.CLEAN