Some imprvmnts (#15)

Co-authored-by: ruatrk8 <kirill.trishin@raiffeisen.ru>

Author: LsHanaha

safe_s3_storage/exceptions.py

@@ -13,6 +13,10 @@ class KasperskyScanEngineThreatDetectedError(BaseError):
     file_name: str
 
 
+@dataclasses.dataclass
+class KasperskyScanEngineConnectionStatusError(BaseError): ...
+
+
 @dataclasses.dataclass
 class NotAllowedMimeTypeError(BaseError):
     file_name: str

safe_s3_storage/file_validator.py

@@ -40,7 +40,7 @@ def _split_file_base_name_and_extensions(file_name: str) -> tuple[str, str | Non
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
 class FileValidator:
     kaspersky_scan_engine: KasperskyScanEngineClient | None = None
-    allowed_mime_types: list[str]
+    allowed_mime_types: list[str] | None = None
     scan_images_with_antivirus: bool = True
     max_file_size_bytes: int = 10 * 1024 * 1024  # 10 MB
     max_image_size_bytes: int = 50 * 1024 * 1024  # 50 MB
@@ -58,7 +58,7 @@ def _validate_mime_type(self, *, file_name: str, file_content: bytes) -> str:
                 mime_type = "application/octet-stream"
             else:
                 mime_type = "text/plain"
-        if mime_type in self.allowed_mime_types:
+        if self.allowed_mime_types is None or mime_type in self.allowed_mime_types:
             return mime_type
 
         raise exceptions.NotAllowedMimeTypeError(

safe_s3_storage/kaspersky_scan_engine.py

@@ -1,13 +1,18 @@
 import base64
 import dataclasses
 import enum
+import logging
 import typing
 
 import httpx
 import pydantic
 import stamina
 
-from safe_s3_storage.exceptions import KasperskyScanEngineThreatDetectedError
+from safe_s3_storage.exceptions import KasperskyScanEngineConnectionStatusError, KasperskyScanEngineThreatDetectedError
+
+
+kaspersky_logger: typing.Final = logging.getLogger(__name__)
+kaspersky_logger.setLevel(logging.ERROR)
 
 
 class KasperskyScanEngineRequest(pydantic.BaseModel):
@@ -47,9 +52,12 @@ async def scan_memory(self, *, file_name: str, file_content: bytes) -> None:
         payload: typing.Final = KasperskyScanEngineRequest(
             timeout=str(self.timeout_ms), object=base64.b64encode(file_content).decode(), name=self.client_name
         ).model_dump(mode="json")
-        response: typing.Final = await stamina.retry(on=httpx.HTTPError, attempts=self.max_retries)(
-            self._send_scan_memory_request
-        )(payload)
+        try:
+            response: typing.Final = await stamina.retry(on=httpx.HTTPError, attempts=self.max_retries)(
+                self._send_scan_memory_request
+            )(payload)
+        except httpx.HTTPStatusError as exc:
+            raise KasperskyScanEngineConnectionStatusError from exc
         validated_response: typing.Final = KasperskyScanEngineResponse.model_validate_json(response)
         if validated_response.scanResult == KasperskyScanEngineScanResult.DETECT:
             raise KasperskyScanEngineThreatDetectedError(response=response, file_name=file_name)

safe_s3_storage/s3_service.py

@@ -3,7 +3,7 @@
 import typing
 
 from types_aiobotocore_s3 import S3Client
-from types_aiobotocore_s3.type_defs import GetObjectOutputTypeDef
+from types_aiobotocore_s3.type_defs import GetObjectOutputTypeDef, HeadObjectOutputTypeDef
 
 from safe_s3_storage.exceptions import FailedToReplaceS3BaseUrlWithProxyBaseUrlError, InvalidS3PathError
 from safe_s3_storage.file_validator import ValidatedFile
@@ -92,3 +92,12 @@ async def create_file_url(
                 s3_file_presigned_url=presigned_url, proxy_base_url=proxy_base_url
             )
         return proxy_base_url.removesuffix("/") + presigned_url_without_prefix
+
+    async def delete_file(self, *, s3_path: str) -> bool:
+        bucket_name, object_key = _extract_bucket_name_and_object_key(s3_path)
+        await self.s3_client.delete_object(Bucket=bucket_name, Key=object_key)
+        return True
+
+    async def collect_file_head(self, *, s3_path: str) -> HeadObjectOutputTypeDef:
+        bucket_name, object_key = _extract_bucket_name_and_object_key(s3_path)
+        return await self.s3_client.head_object(Bucket=bucket_name, Key=object_key)

tests/test_file_validator.py

@@ -4,8 +4,10 @@
 import faker
 import httpx
 import pytest
+from httpx import codes as status_codes
 
 from safe_s3_storage import exceptions
+from safe_s3_storage.exceptions import KasperskyScanEngineConnectionStatusError
 from safe_s3_storage.file_validator import (
     _IMAGE_CONVERSION_FORMAT_TO_MIME_TYPE_AND_EXTENSION_MAP,
     FileValidator,
@@ -56,7 +58,21 @@ def get_mocked_kaspersky_scan_engine_client(*, faker: faker.Faker, ok_response:
         service_url=faker.url(schemes=["http"]),
         client_name=faker.pystr(),
         httpx_client=httpx.AsyncClient(
-            transport=httpx.MockTransport(lambda _: httpx.Response(200, json=scan_response.model_dump(mode="json")))
+            transport=httpx.MockTransport(
+                lambda _: httpx.Response(status_codes.OK, json=scan_response.model_dump(mode="json"))
+            )
+        ),
+    )
+
+
+def get_mocked_kaspersky_scan_engine_client_bad_response(
+    *, faker: faker.Faker, status_code: int = status_codes.OK
+) -> KasperskyScanEngineClient:
+    return KasperskyScanEngineClient(
+        service_url=faker.url(schemes=["http"]),
+        client_name=faker.pystr(),
+        httpx_client=httpx.AsyncClient(
+            transport=httpx.MockTransport(lambda _: httpx.Response(status_code, json="")),
         ),
     )
 
@@ -86,6 +102,12 @@ async def test_fails_to_convert_image(self, faker: faker.Faker, png_file: bytes)
                 file_name=faker.file_name(), file_content=png_file[:50]
             )
 
+    async def test_all_mime_types_allowed(self, faker: faker.Faker) -> None:
+        validated_file: typing.Final = await FileValidator(allowed_mime_types=None).validate_file(
+            file_name=faker.file_name(), file_content=generate_binary_content(faker)
+        )
+        assert validated_file is not None
+
     @pytest.mark.parametrize("image_conversion_format", list(ImageConversionFormat))
     async def test_ok_image(
         self, faker: faker.Faker, png_file: bytes, image_conversion_format: ImageConversionFormat
@@ -154,3 +176,10 @@ async def test_antivirus_passes_on_images(self, faker: faker.Faker, png_file: by
             kaspersky_scan_engine=get_mocked_kaspersky_scan_engine_client(faker=faker, ok_response=True),
             allowed_mime_types=["image/png"],
         ).validate_file(file_name=faker.file_name(), file_content=png_file)
+
+    async def test_antivirus_no_connection(self, faker: faker.Faker, png_file: bytes) -> None:
+        kasper: typing.Final = get_mocked_kaspersky_scan_engine_client_bad_response(
+            faker=faker, status_code=status_codes.GATEWAY_TIMEOUT
+        )
+        with pytest.raises(KasperskyScanEngineConnectionStatusError):
+            await kasper.scan_memory(file_name=faker.file_name(), file_content=png_file)

tests/test_s3_service.py

@@ -82,6 +82,24 @@ async def test_fails_to_parse_s3_path(self, faker: faker.Faker) -> None:
             await S3Service(s3_client=mock.Mock()).read_file(s3_path=faker.pystr())
 
 
+class TestS3ServiceDelete:
+    async def test_ok_delete(self, faker: faker.Faker) -> None:
+        bucket_name, s3_key = faker.pystr(), faker.pystr()
+        s3_client_mock: typing.Final = mock.AsyncMock(delete_object=mock.AsyncMock(return_value={}))
+
+        await S3Service(s3_client=s3_client_mock).delete_file(s3_path=f"{bucket_name}/{s3_key}")
+        s3_client_mock.delete_object.assert_called_once_with(Bucket=bucket_name, Key=s3_key)
+
+
+class TestS3ServiceHead:
+    async def test_ok_head(self, faker: faker.Faker) -> None:
+        bucket_name, s3_key = faker.pystr(), faker.pystr()
+        s3_client_mock: typing.Final = mock.AsyncMock(head_object=mock.AsyncMock(return_value={}))
+
+        await S3Service(s3_client=s3_client_mock).collect_file_head(s3_path=f"{bucket_name}/{s3_key}")
+        s3_client_mock.head_object.assert_called_once_with(Bucket=bucket_name, Key=s3_key)
+
+
 class TestS3ServiceCreateFileUrl:
     async def test_call(self, faker: faker.Faker) -> None:
         bucket_name, s3_key, display_file_name = faker.pystr(), faker.pystr(), faker.pystr()