Adds New Delete File Endpoint and Refactors Endpoints, Adds VirusChecker

storage_service/service/storage/amazon_s3_service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from storage_service.service.storage_service import StorageService
+from storage_service.depends.depend_virus_checker_service import dependency_virus_checker_service
+from storage_service.service.storage.storage_service import StorageService
 from storage_service.utils.enums.file_type import FileType
 from storage_service.utils.file_handler import FILE_HANDLER
 
@@ -11,6 +12,9 @@ from typing import Any
 
 
 class AmazonS3Service(StorageService):
+
+    virus_checker_service = dependency_virus_checker_service()
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -39,8 +43,15 @@ class AmazonS3Service(StorageService):
     def get_temp_read_link(self, file_name) -> dict[str, str | None]:
         return {"presigned_url": self._get_presigned_read_url(file_name)}
 
+    def delete_file(self, file_name: str) -> None:
+        self._delete_file(file_name)
+
     def process_file(self, file_name: str, file_type: FileType = FileType.PNG) -> None:
         file_bytes = self._get_file_obj(file_name)
+
+        if not self.virus_checker_service.check_virus(file_bytes):
+            self._delete_file(file_name)
+
         handler = FILE_HANDLER[file_type]["handler"]
 
         self._upload_file(file_name, handler(file_bytes))
@@ -78,6 +89,9 @@ class AmazonS3Service(StorageService):
     def _upload_file(self, file_name: str, file_bytes: io.BytesIO) -> None:
         self.s3.upload_fileobj(file_bytes, Bucket=self.bucket_name, Key=file_name)
 
+    def _delete_file(self, file_name: str) -> None:
+        self.s3.delete_object(Bucket=self.bucket_name, Key=file_name)
+
     @staticmethod
     def __validate_config(**kwargs):
         if not kwargs.get("bucket_name"):

storage_service/service/storage/storage_service.py

@@ -20,6 +20,10 @@ class StorageService(ABC):
     def get_temp_read_link(self, file_name) -> dict[str, str | None]:
         pass
 
+    @abstractmethod
+    def delete_file(self, file_name: str) -> None:
+        pass
+
     @abstractmethod
     def process_file(self, file_name: str, file_type: FileType) -> None:
         pass

storage_service/service/virus_checker/virus_checker_service.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+from io import BytesIO
+
+
+class VirusCheckerService(ABC):
+    @abstractmethod
+    def check_virus(self, file_data: BytesIO) -> bool:
+        pass

storage_service/service/virus_checker/virus_total_service.py

@@ -0,0 +1,35 @@
+from io import BytesIO
+from virustotal_python import Virustotal
+
+from storage_service.service.virus_checker.virus_checker_service import VirusCheckerService
+
+
+class VirusTotalService(VirusCheckerService):
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+
+    def check_virus(self, file_data: BytesIO) -> bool:
+        files = {"file": ("image_file", file_data)}
+
+        with Virustotal(self.api_key) as vtotal:
+            resp = vtotal.request("files", files=files, method="POST")
+
+            file_attributes = self._get_analysis(resp.json()["data"]["id"])
+
+            return self._is_valid_file(file_attributes["data"]["attributes"]["stats"])
+
+    def _get_analysis(self, file_id: str) -> dict:
+        with Virustotal(self.api_key) as vtotal:
+            resp = vtotal.request(f"analyses/{file_id}")
+
+            return resp.json()
+
+    @staticmethod
+    def _is_valid_file(file_stats: dict) -> bool:
+        if 'malicious' in file_stats and file_stats['malicious'] > 0:
+            return False
+
+        if 'suspicious' in file_stats and file_stats['suspicious'] > 0:
+            return False
+
+        return True