Внедрение KESL на ВМ2 + замена CLAMAV на KESL

VM2_services/codebase/services/message-safety/app/worker.py

@@ -12,8 +12,9 @@ from app.config import validate_config
 from app.contracts import Attachment
 from app.db import FileVerdictCache, SafetyAudit, engine_and_sessions
 from app.file_pipeline import (
-    ClamAvInstream,
+    Antivirus,
     DependencyFailure,
+    KeslSocketScanner,
     ObjectChanged,
     collect_and_hash,
     detect_format,
@@ -34,7 +35,7 @@ tracer = trace.get_tracer("message-safety.worker")
 
 class Worker:
     def __init__(
-        self, repository: Repository, reader: S3VersionReader, antivirus: ClamAvInstream, artifacts
+        self, repository: Repository, reader: S3VersionReader, antivirus: Antivirus, artifacts
     ) -> None:
         self.repository, self.reader, self.antivirus, self.artifacts = (
             repository,
@@ -94,10 +95,15 @@ class Worker:
             )
             record_dependency("s3", "get_object", "success")
             rule = detect_format(body, attachment.mime_type)
+            signatures_version = task.signatures_version
             if not rule:
-                malware = await self.antivirus.scan(one_chunk(body))
-                record_dependency("clamav", "scan", "success")
-                rule = "file.malware_detected" if malware else None
+                scan = await self.antivirus.scan(
+                    one_chunk(body),
+                    scan_timeout=row.config["antivirus"]["scan_timeout_sec"],
+                )
+                signatures_version = scan.signatures_version
+                record_dependency("antivirus", "scan", "success")
+                rule = "file.malware_detected" if scan.threat else None
             with tracer.start_as_current_span("message_safety.worker.finalize"):
                 finished = await self.repository.finish(
                     task.id,
@@ -105,6 +111,7 @@ class Worker:
                     task.lease_generation,
                     allow=rule is None,
                     rule_id=rule or "safety.all_checks_passed",
+                    signatures_version=signatures_version,
                 )
             if finished:
                 record_worker("allow" if rule is None else "deny", task_age)
@@ -116,7 +123,7 @@ class Worker:
                         rules_version=task.rules_version,
                         detector_version=task.detector_version,
                         scanner_engine=task.scanner_engine,
-                        signatures_version=task.signatures_version,
+                        signatures_version=signatures_version,
                         verdict="allow" if rule is None else "deny",
                         rule_id=rule or "safety.all_checks_passed",
                         reason_code=None if rule is None else "message_blocked",
@@ -192,7 +199,7 @@ async def serve() -> None:
                 settings.s3_access_key.get_secret_value(),
                 settings.s3_secret_key.get_secret_value(),
             ),
-            ClamAvInstream(settings.clamav_host, settings.clamav_port),
+            KeslSocketScanner(settings.antivirus_socket),
             settings.artifacts_dir,
         )
         async with asyncio.TaskGroup() as group: