diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e91d92..1d55db8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to - ✨(backend) keep traces of failed malware analysis tasks - ✨(backend) save backend used in a malware analysis task +- ✨(backend) allow a malware detection backend to reschedule a task ## [0.0.19] - 2025-11-21 diff --git a/src/lasuite/malware_detection/backends/base.py b/src/lasuite/malware_detection/backends/base.py index 029e122..374b565 100644 --- a/src/lasuite/malware_detection/backends/base.py +++ b/src/lasuite/malware_detection/backends/base.py @@ -4,6 +4,8 @@ from abc import ABC, abstractmethod from django.utils.module_loading import import_string +from ..models import MalwareDetection + class BaseBackend(ABC): """Base class for all malware detection backends.""" @@ -34,3 +36,7 @@ class BaseBackend(ABC): @abstractmethod def launch_next_analysis(self) -> None: """Launch the next analysis.""" + + @abstractmethod + def reschedule_processing_task(self, malware_detection_record: MalwareDetection) -> None: + """Reschedule the processing task for a malware detection record.""" diff --git a/src/lasuite/malware_detection/backends/dummy.py b/src/lasuite/malware_detection/backends/dummy.py index e411ef9..480e611 100644 --- a/src/lasuite/malware_detection/backends/dummy.py +++ b/src/lasuite/malware_detection/backends/dummy.py @@ -1,6 +1,7 @@ """Module contains the dummy backend for the malware detection system.""" from ..enums import ReportStatus +from ..models import MalwareDetection from .base import BaseBackend @@ -13,3 +14,6 @@ class DummyBackend(BaseBackend): def launch_next_analysis(self) -> None: """Launch the next analysis.""" + + def reschedule_processing_task(self, malware_detection_record: MalwareDetection) -> None: + """Reschedule the processing task for a malware detection record.""" diff --git a/src/lasuite/malware_detection/backends/jcop.py b/src/lasuite/malware_detection/backends/jcop.py index 303643d..2e7de50 100644 --- a/src/lasuite/malware_detection/backends/jcop.py +++ b/src/lasuite/malware_detection/backends/jcop.py @@ -64,6 +64,25 @@ class JCOPBackend(BaseBackend): self.launch_next_analysis() + def reschedule_processing_task(self, malware_detection_record: MalwareDetection) -> None: + """Reschedule the processing task for a malware detection record.""" + if malware_detection_record.status != MalwareDetectionStatus.PROCESSING: + return + + if malware_detection_record.backend != self.backend_name: + return + + # assert first the file still exists in the system + if not default_storage.exists(malware_detection_record.path): + logger.info("File %s not found when rescheduling processing task", malware_detection_record.path) + malware_detection_record.delete() + return + + analyse_file_async.delay( + malware_detection_record.path, + **malware_detection_record.parameters, + ) + def launch_next_analysis(self) -> None: """Launch the next pending analysis.""" if ( diff --git a/tests/malware_detection/backends/test_jcop_backend.py b/tests/malware_detection/backends/test_jcop_backend.py index 477143c..dc424ef 100644 --- a/tests/malware_detection/backends/test_jcop_backend.py +++ b/tests/malware_detection/backends/test_jcop_backend.py @@ -883,3 +883,67 @@ def test_jcop_backend_delete_non_existing_detection(jcop_backend): jcop_backend.delete_detection("file.txt") assert MalwareDetection.objects.count() == 0 + + +def test_jcop_backend_reschedule_processing_task(jcop_generate_file_path, jcop_backend): + """Reschedule the processing task for a malware detection record.""" + file_path, _ = jcop_generate_file_path + malware_detection = factories.MalwareDetectionFactory( + path=file_path, + status=MalwareDetectionStatus.PROCESSING, + backend="lasuite.malware_detection.backends.jcop.JCOPBackend", + ) + + with mock.patch.object(analyse_file_async, "delay") as analyse_file_async_mock: + jcop_backend.reschedule_processing_task(malware_detection) + analyse_file_async_mock.assert_called_once_with( + file_path, + ) + + +def test_jcop_backend_reschedule_processing_missing_file(jcop_backend): + """Reschedule the processing task for a malware detection record with a missing file.""" + file_path = "file.txt" + malware_detection = factories.MalwareDetectionFactory( + path=file_path, + status=MalwareDetectionStatus.PROCESSING, + backend="lasuite.malware_detection.backends.jcop.JCOPBackend", + ) + + with mock.patch.object(analyse_file_async, "delay") as analyse_file_async_mock: + jcop_backend.reschedule_processing_task(malware_detection) + analyse_file_async_mock.assert_not_called() + + assert not MalwareDetection.objects.filter(path=file_path).exists() + + +def test_jcop_backend_reschedule_processing_not_processing(jcop_backend): + """Reschedule the processing task for a malware detection record with a not processing status.""" + file_path = "file.txt" + malware_detection = factories.MalwareDetectionFactory( + path=file_path, + status=MalwareDetectionStatus.PENDING, + backend="lasuite.malware_detection.backends.jcop.JCOPBackend", + ) + + with mock.patch.object(analyse_file_async, "delay") as analyse_file_async_mock: + jcop_backend.reschedule_processing_task(malware_detection) + analyse_file_async_mock.assert_not_called() + + assert MalwareDetection.objects.filter(path=file_path).exists() + + +def test_jcop_backend_reschedule_processing_not_jcop_backend(jcop_backend): + """Reschedule the processing task for a malware detection record with a not JCOP backend.""" + file_path = "file.txt" + malware_detection = factories.MalwareDetectionFactory( + path=file_path, + status=MalwareDetectionStatus.PROCESSING, + backend="lasuite.malware_detection.backends.dummy.DummyBackend", + ) + + with mock.patch.object(analyse_file_async, "delay") as analyse_file_async_mock: + jcop_backend.reschedule_processing_task(malware_detection) + analyse_file_async_mock.assert_not_called() + + assert MalwareDetection.objects.filter(path=file_path).exists() diff --git a/tests/malware_detection/test_check_analysis_pending_command.py b/tests/malware_detection/test_check_analysis_pending_command.py index 96032ad..fe8dcb0 100644 --- a/tests/malware_detection/test_check_analysis_pending_command.py +++ b/tests/malware_detection/test_check_analysis_pending_command.py @@ -5,6 +5,7 @@ from unittest import mock from django.core.management import call_command from lasuite.malware_detection.backends.base import BaseBackend +from lasuite.malware_detection.models import MalwareDetection mock_launch_next_analysis = mock.MagicMock() @@ -19,6 +20,9 @@ class TestBackend(BaseBackend): def analyse_file(self, file_path: str, **kwargs) -> None: """Analyse a file.""" + def reschedule_processing_task(self, malware_detection_record: MalwareDetection) -> None: + """Reschedule the processing task for a malware detection record.""" + def test_check_analysis_pending_command(settings): """Test the check_analysis_pending command."""