✨(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:
76
src/backend/core/tests/test_malware_detection.py
Normal file
76
src/backend/core/tests/test_malware_detection.py
Normal 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)
|
||||
Reference in New Issue
Block a user