(backend) configure lasuite.malware_detection module

We want to use the malware_detection module from lasuite library. We add
a new setting MALWARE_DETECTION to configure the backend we want to use.
The callback is also added. It removes the file if it is not safe or
change it's status in the metadata to set it as ready.
This commit is contained in:
Manuel Raynaud
2025-05-05 15:58:36 +02:00
parent 37d9ae8cca
commit a070e1dd87
6 changed files with 154 additions and 1 deletions

View File

@@ -0,0 +1,76 @@
"""Test malware detection callback."""
import random
from django.core.files.base import ContentFile
from django.core.files.storage import default_storage
import pytest
from lasuite.malware_detection.enums import ReportStatus
from core.enums import DocumentAttachmentStatus
from core.factories import DocumentFactory
from core.malware_detection import malware_detection_callback
pytestmark = pytest.mark.django_db
@pytest.fixture
def safe_file():
"""Create a safe file."""
file_path = "test.txt"
default_storage.save(file_path, ContentFile("test"))
yield file_path
default_storage.delete(file_path)
@pytest.fixture
def unsafe_file():
"""Create an unsafe file."""
file_path = "unsafe.txt"
default_storage.save(file_path, ContentFile("test"))
yield file_path
def test_malware_detection_callback_safe_status(safe_file):
"""Test malware detection callback with safe status."""
document = DocumentFactory(attachments=[safe_file])
malware_detection_callback(
safe_file,
ReportStatus.SAFE,
error_info={},
document_id=document.id,
)
document.refresh_from_db()
assert safe_file in document.attachments
assert default_storage.exists(safe_file)
s3_client = default_storage.connection.meta.client
bucket_name = default_storage.bucket_name
head_resp = s3_client.head_object(Bucket=bucket_name, Key=safe_file)
metadata = head_resp.get("Metadata", {})
assert metadata["status"] == DocumentAttachmentStatus.READY
def test_malware_detection_callback_unsafe_status(unsafe_file):
"""Test malware detection callback with unsafe status."""
document = DocumentFactory(attachments=[unsafe_file])
malware_detection_callback(
unsafe_file,
random.choice(
[status.value for status in ReportStatus if status != ReportStatus.SAFE]
),
error_info={"error": "test", "error_code": 4001},
document_id=document.id,
)
document.refresh_from_db()
assert unsafe_file not in document.attachments
assert not default_storage.exists(unsafe_file)