✨(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:
51
src/backend/core/malware_detection.py
Normal file
51
src/backend/core/malware_detection.py
Normal file
@@ -0,0 +1,51 @@
|
||||
"""Malware detection callbacks"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
from lasuite.malware_detection.enums import ReportStatus
|
||||
|
||||
from core.enums import DocumentAttachmentStatus
|
||||
from core.models import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def malware_detection_callback(file_path, status, error_info, **kwargs):
|
||||
"""Malware detection callback"""
|
||||
|
||||
if status == ReportStatus.SAFE:
|
||||
logger.info("File %s is safe", file_path)
|
||||
# Get existing metadata
|
||||
s3_client = default_storage.connection.meta.client
|
||||
bucket_name = default_storage.bucket_name
|
||||
head_resp = s3_client.head_object(Bucket=bucket_name, Key=file_path)
|
||||
metadata = head_resp.get("Metadata", {})
|
||||
metadata.update({"status": DocumentAttachmentStatus.READY})
|
||||
# Update status in metadata
|
||||
s3_client.copy_object(
|
||||
Bucket=bucket_name,
|
||||
CopySource={"Bucket": bucket_name, "Key": file_path},
|
||||
Key=file_path,
|
||||
ContentType=head_resp.get("ContentType"),
|
||||
Metadata=metadata,
|
||||
MetadataDirective="REPLACE",
|
||||
)
|
||||
return
|
||||
|
||||
document_id = kwargs.get("document_id")
|
||||
logger.error(
|
||||
"File %s for document %s is infected with malware. Error info: %s",
|
||||
file_path,
|
||||
document_id,
|
||||
error_info,
|
||||
)
|
||||
|
||||
# Remove the file from the document and change the status to unsafe
|
||||
document = Document.objects.get(pk=document_id)
|
||||
document.attachments.remove(file_path)
|
||||
document.save(update_fields=["attachments"])
|
||||
|
||||
# Delete the file from the storage
|
||||
default_storage.delete(file_path)
|
||||
Reference in New Issue
Block a user