Skip to content

Commit a070e1d

Browse files
committed
✨(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.
1 parent 37d9ae8 commit a070e1d

File tree

6 files changed

+154
-1
lines changed

6 files changed

+154
-1
lines changed

docs/env.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,5 @@ These are the environmental variables you can set for the impress-backend contai
9898
| DJANGO_CSRF_TRUSTED_ORIGINS | CSRF trusted origins | [] |
9999
| REDIS_URL | cache url | redis://redis:6379/1 |
100100
| CACHES_DEFAULT_TIMEOUT | cache default timeout | 30 |
101+
| MALWARE_DETECTION_BACKEND | The malware detection backend use from the django-lasuite package | lasuite.malware_detection.backends.dummy.DummyBackend |
102+
| MALWARE_DETECTION_PARAMETERS | A dict containing all the parameters to initiate the malware detection backend | {"callback_path": "core.malware_detection.malware_detection_callback",} |

src/backend/core/enums.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44

55
import re
6+
from enum import StrEnum
67

78
from django.conf import global_settings, settings
89
from django.db import models
@@ -38,3 +39,10 @@ class MoveNodePositionChoices(models.TextChoices):
3839
LAST_SIBLING = "last-sibling", _("Last sibling")
3940
LEFT = "left", _("Left")
4041
RIGHT = "right", _("Right")
42+
43+
44+
class DocumentAttachmentStatus(StrEnum):
45+
"""Defines the possible statuses for an attachment."""
46+
47+
PROCESSING = "processing"
48+
READY = "ready"

src/backend/core/malware_detection.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
"""Malware detection callbacks"""
2+
3+
import logging
4+
5+
from django.core.files.storage import default_storage
6+
7+
from lasuite.malware_detection.enums import ReportStatus
8+
9+
from core.enums import DocumentAttachmentStatus
10+
from core.models import Document
11+
12+
logger = logging.getLogger(__name__)
13+
14+
15+
def malware_detection_callback(file_path, status, error_info, **kwargs):
16+
"""Malware detection callback"""
17+
18+
if status == ReportStatus.SAFE:
19+
logger.info("File %s is safe", file_path)
20+
# Get existing metadata
21+
s3_client = default_storage.connection.meta.client
22+
bucket_name = default_storage.bucket_name
23+
head_resp = s3_client.head_object(Bucket=bucket_name, Key=file_path)
24+
metadata = head_resp.get("Metadata", {})
25+
metadata.update({"status": DocumentAttachmentStatus.READY})
26+
# Update status in metadata
27+
s3_client.copy_object(
28+
Bucket=bucket_name,
29+
CopySource={"Bucket": bucket_name, "Key": file_path},
30+
Key=file_path,
31+
ContentType=head_resp.get("ContentType"),
32+
Metadata=metadata,
33+
MetadataDirective="REPLACE",
34+
)
35+
return
36+
37+
document_id = kwargs.get("document_id")
38+
logger.error(
39+
"File %s for document %s is infected with malware. Error info: %s",
40+
file_path,
41+
document_id,
42+
error_info,
43+
)
44+
45+
# Remove the file from the document and change the status to unsafe
46+
document = Document.objects.get(pk=document_id)
47+
document.attachments.remove(file_path)
48+
document.save(update_fields=["attachments"])
49+
50+
# Delete the file from the storage
51+
default_storage.delete(file_path)
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
"""Test malware detection callback."""
2+
3+
import random
4+
5+
from django.core.files.base import ContentFile
6+
from django.core.files.storage import default_storage
7+
8+
import pytest
9+
from lasuite.malware_detection.enums import ReportStatus
10+
11+
from core.enums import DocumentAttachmentStatus
12+
from core.factories import DocumentFactory
13+
from core.malware_detection import malware_detection_callback
14+
15+
pytestmark = pytest.mark.django_db
16+
17+
18+
@pytest.fixture
19+
def safe_file():
20+
"""Create a safe file."""
21+
file_path = "test.txt"
22+
default_storage.save(file_path, ContentFile("test"))
23+
yield file_path
24+
default_storage.delete(file_path)
25+
26+
27+
@pytest.fixture
28+
def unsafe_file():
29+
"""Create an unsafe file."""
30+
file_path = "unsafe.txt"
31+
default_storage.save(file_path, ContentFile("test"))
32+
yield file_path
33+
34+
35+
def test_malware_detection_callback_safe_status(safe_file):
36+
"""Test malware detection callback with safe status."""
37+
38+
document = DocumentFactory(attachments=[safe_file])
39+
40+
malware_detection_callback(
41+
safe_file,
42+
ReportStatus.SAFE,
43+
error_info={},
44+
document_id=document.id,
45+
)
46+
47+
document.refresh_from_db()
48+
49+
assert safe_file in document.attachments
50+
assert default_storage.exists(safe_file)
51+
52+
s3_client = default_storage.connection.meta.client
53+
bucket_name = default_storage.bucket_name
54+
head_resp = s3_client.head_object(Bucket=bucket_name, Key=safe_file)
55+
metadata = head_resp.get("Metadata", {})
56+
assert metadata["status"] == DocumentAttachmentStatus.READY
57+
58+
59+
def test_malware_detection_callback_unsafe_status(unsafe_file):
60+
"""Test malware detection callback with unsafe status."""
61+
62+
document = DocumentFactory(attachments=[unsafe_file])
63+
64+
malware_detection_callback(
65+
unsafe_file,
66+
random.choice(
67+
[status.value for status in ReportStatus if status != ReportStatus.SAFE]
68+
),
69+
error_info={"error": "test", "error_code": 4001},
70+
document_id=document.id,
71+
)
72+
73+
document.refresh_from_db()
74+
75+
assert unsafe_file not in document.attachments
76+
assert not default_storage.exists(unsafe_file)

src/backend/impress/settings.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,7 @@ class Base(Configuration):
317317
"django.contrib.staticfiles",
318318
# OIDC third party
319319
"mozilla_django_oidc",
320+
"lasuite.malware_detection",
320321
]
321322

322323
# Cache
@@ -680,6 +681,21 @@ class Base(Configuration):
680681
},
681682
}
682683

684+
MALWARE_DETECTION = {
685+
"BACKEND": values.Value(
686+
"lasuite.malware_detection.backends.dummy.DummyBackend",
687+
environ_name="MALWARE_DETECTION_BACKEND",
688+
environ_prefix=None,
689+
),
690+
"PARAMETERS": values.DictValue(
691+
default={
692+
"callback_path": "core.malware_detection.malware_detection_callback",
693+
},
694+
environ_name="MALWARE_DETECTION_PARAMETERS",
695+
environ_prefix=None,
696+
),
697+
}
698+
683699
API_USERS_LIST_LIMIT = values.PositiveIntegerValue(
684700
default=5,
685701
environ_name="API_USERS_LIST_LIMIT",

src/backend/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ dependencies = [
3333
"django-cors-headers==4.7.0",
3434
"django-countries==7.6.1",
3535
"django-filter==25.1",
36-
"django-lasuite==0.0.7",
36+
"django-lasuite[all]==0.0.8",
3737
"django-parler==2.3",
3838
"django-redis==5.4.0",
3939
"django-storages[s3]==1.14.6",

0 commit comments

Comments
 (0)