/opt/imunify360/venv/lib/python3.11/site-packages/imav/malwarelib/plugins
Edit: /opt/imunify360/venv/lib/python3.11/site-packages/imav/malwarelib/plugins/detached_scan.py (13358B)
"""
This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License,
or (at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see
.
Copyright © 2019 Cloud Linux Software Inc.
This software is also available under ImunifyAV commercial license,
see
"""
import shutil
import time
from logging import getLogger
from typing import Any, Optional, Union
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import (
MessageSink,
MessageSource,
expect,
)
from defence360agent.internals.the_sink import TheSink
from defence360agent.utils import Scope
from imav.contracts.messages import MalwareDatabaseScan
from imav.malwarelib.config import (
MalwareScanResourceType,
MalwareScanType,
)
from imav.malwarelib.model import MalwareScan as MalwareScanModel
from imav.malwarelib.scan import (
ScanAlreadyCompleteError,
ScanInfoError,
)
from imav.malwarelib.scan.ai_bolit.detached import (
AiBolitDetachedScan,
)
from imav.malwarelib.scan.detached import DetachedScan
from imav.malwarelib.scan.mds.detached import MDSDetachedScan
from imav.malwarelib.scan.queue_supervisor_sync import QueueSupervisorSync
from imav.malwarelib.scan.scan_result import aggregate_result
from imav.malwarelib.utils.user_list import fill_results_owner
logger = getLogger(__name__)
class DetachedScanPlugin(MessageSink, MessageSource):
PROCESSING_ORDER = MessageSink.ProcessingOrder.PRE_PROCESS_MESSAGE
SCOPE = Scope.AV
loop = None
sink: TheSink
results_cache: dict[str, dict[str, Any]] = {}
# scan_ids whose complete() just succeeded. A duplicate MalwareScanComplete
# that then finds the report gone (it is removed right after the first,
# successful completion) must not be finalized as aborted -- the scan
# completed fine. Bounded; scan_ids are unique and the single-slot queue
# completes one scan at a time, so a small cap is ample.
_completed_scan_ids: dict[str, None] = {}
_MAX_COMPLETED_SCAN_IDS = 256
async def create_source(self, loop, sink):
self.loop = loop
self.sink = sink
async def create_sink(self, loop):
pass
@expect(MessageType.MalwareScan, async_lock=True)
async def complete_scan(self, message):
message_type = MalwareScanMessageInfo(message)
if not message_type.is_detached:
total_malicious = await self._count_total_malicious(message)
message["summary"]["total_malicious"] = total_malicious
return message
elif message_type.is_summary:
return await self._handle_summary(message)
# message_type.is_result
return await self._handle_results(message)
async def _handle_summary(self, message):
scan_id = message["summary"]["scanid"]
# If summary arrives after results, results are read from cache
if scan_id in self.results_cache:
message["summary"]["completed"] = time.time()
message["results"] = self.results_cache.pop(scan_id)
total_malicious = await self._count_total_malicious(message)
message["summary"]["total_malicious"] = total_malicious
queued_scan = QueueSupervisorSync.queue.find(
scanid=message["summary"]["scanid"]
)
if queued_scan:
message["summary"]["scan_args"] = queued_scan.args
QueueSupervisorSync.queue.remove(queued_scan)
return message
async def _handle_results(self, message):
message = await self.aggregate_result(message)
message_type = MalwareScanMessageInfo(message)
summary = message["summary"]
logger.info("Scan stopped")
queued_scan = QueueSupervisorSync.queue.find(scanid=summary["scanid"])
if message_type.summary_from_db is None:
if queued_scan:
summary["file_patterns"] = queued_scan.args["file_patterns"]
summary["exclude_patterns"] = queued_scan.args[
"exclude_patterns"
]
summary["scan_args"] = queued_scan.args
QueueSupervisorSync.queue.remove(queued_scan)
if summary.get("path") or summary.get("error"):
# Scan failed
summary["total_malicious"] = 0
await self._recheck_scan_queue()
return message
# Summary is not in DB yet, save results to cache
scan_id = message["summary"]["scanid"]
self.results_cache[scan_id] = message["results"]
# Report an error to Sentry if cache grows
cache_size = len(self.results_cache)
if cache_size > 1:
logger.error("MalwareScan cache size is %d", cache_size)
return
scan = message_type.summary_from_db
summary["scanid"] = scan.scanid
summary["path"] = scan.path
summary["started"] = scan.started
summary["completed"] = time.time()
if summary.get("total_files") is None:
summary["total_files"] = scan.total_resources
summary["type"] = scan.type
summary["error"] = summary.get("error", None)
message["summary"] = summary
total_malicious = await self._count_total_malicious(message)
message["summary"]["total_malicious"] = total_malicious
if queued_scan:
summary["file_patterns"] = queued_scan.args["file_patterns"]
summary["exclude_patterns"] = queued_scan.args["exclude_patterns"]
summary["scan_args"] = queued_scan.args
QueueSupervisorSync.queue.remove(queued_scan)
await self._recheck_scan_queue()
return message
@staticmethod
async def _count_total_malicious(message) -> int:
return len(
[
k
for k, v in message["results"].items()
if v["hits"][0]["suspicious"] is False
]
)
@staticmethod
def _get_detached_scan(
resource_type: Optional[Union[str, MalwareScanResourceType]],
scan_id: str,
) -> DetachedScan:
return AiBolitDetachedScan(scan_id)
@expect(MessageType.MalwareScanComplete)
async def complete_detached_scan(self, message):
scan_id = message.get("scan_id")
resource_type = message.get("resource_type")
detached_scan = self._get_detached_scan(resource_type, scan_id)
try:
scan_message = await detached_scan.complete()
# Record the success before the finally clause removes the report.
# A duplicate completion that then finds the report gone is checked
# against this below, so it is never finalized as aborted.
self._mark_completed(scan_id)
except ScanAlreadyCompleteError as err:
# A duplicate MalwareScanComplete is normal: the scan completed, its
# report was consumed, then complete() runs again when AV is woken
# up by AiBolit. If this scan completed (possibly concurrently --
# messages are not serialized per scan), the success owns it.
# Otherwise, if it is still queued its report is gone yet it never
# reached a terminal state and would block the queue forever --
# finalize it.
if scan_id in self._completed_scan_ids:
logger.warning(
"Scan %s already completed; ignoring duplicate completion"
":\n%s",
scan_id,
err,
)
elif await self._finalize_stuck_scan(detached_scan, scan_id):
logger.warning(
"Scan %s has no report but is still queued;"
" finalized it as aborted: %s",
scan_id,
err,
)
else:
logger.warning(
"Cannot complete scan %s, assuming it is already complete"
":\n%s",
scan_id,
err,
)
return
except ScanInfoError as err:
if scan_id in self._completed_scan_ids:
logger.warning(
"Scan %s already completed; ignoring duplicate completion"
":\n%s",
scan_id,
err,
)
elif await self._finalize_stuck_scan(detached_scan, scan_id):
logger.warning(
"Scan %s has no scan_info but is still queued;"
" finalized it as aborted: %s",
scan_id,
err,
)
else:
logger.error(
"Cannot complete %s scan %s, assuming it was not started"
":\n%s",
detached_scan.RESOURCE_TYPE.value,
scan_id,
err,
)
return
finally:
shutil.rmtree(str(detached_scan.detached_dir), ignore_errors=True)
await self.sink.process_message(scan_message)
def _mark_completed(self, scan_id: str) -> None:
"""Remember a scan whose complete() just succeeded (bounded set)."""
self._completed_scan_ids[scan_id] = None
# dicts preserve insertion order -> drop the oldest entries past the cap.
while len(self._completed_scan_ids) > self._MAX_COMPLETED_SCAN_IDS:
del self._completed_scan_ids[next(iter(self._completed_scan_ids))]
async def _finalize_stuck_scan(self, detached_scan, scan_id) -> bool:
"""Route a still-queued scan whose report/scan_info is gone through the
aborted path so it is recorded failed and stops blocking the queue;
return whether it was still queued and got finalized."""
# A scan that completed (even concurrently) is filtered out by the
# completed-id guard in the caller, so reaching here means it is stuck.
queued = QueueSupervisorSync.queue.find(scanid=scan_id)
if queued is None:
return False
# kill=False: there is no live worker to kill (its output is already
# gone), and kill=True would block ~30s waiting for a pid file. Pass the
# queued scan's start time so the aborted summary is persisted -- the
# store drops summaries with a falsy "started".
await detached_scan.handle_aborted_process(
sink=self.sink,
kill=False,
scan_started=getattr(queued, "started", None),
)
return True
@classmethod
async def aggregate_result(cls, message):
message["results"] = aggregate_result(message["results"])
await fill_results_owner(message["results"])
return message
async def _recheck_scan_queue(self):
await self.sink.process_message(MessageType.MalwareScanQueueRecheck())
class MalwareScanMessageInfo:
"""A helper class that allows to receive information about scan
from MalwareScan message.
"""
def __init__(self, message):
self.message = message
self._summary_from_db = None
self.scan_id = self.message["summary"]["scanid"]
@property
def is_detached(self):
summary = self.message["summary"]
return summary.get("type") in (
MalwareScanType.ON_DEMAND,
MalwareScanType.BACKGROUND,
MalwareScanType.USER,
None,
)
@property
def is_summary(self):
return self.message["results"] is None
@property
def summary_from_db(self):
if not self._summary_from_db:
summary_from_db = (
MalwareScanModel.select()
.where(MalwareScanModel.scanid == self.scan_id)
.limit(1)
)
if summary_from_db:
self._summary_from_db = summary_from_db[0]
return self._summary_from_db
class DetachedScanPluginIm360(DetachedScanPlugin):
SCOPE = Scope.IM360
@staticmethod
def _get_detached_scan(
resource_type: Optional[Union[str, MalwareScanResourceType]],
scan_id: str,
) -> DetachedScan:
if resource_type is not None and (
MalwareScanResourceType(resource_type)
is MalwareScanResourceType.DB
):
return MDSDetachedScan(scan_id)
return AiBolitDetachedScan(scan_id)
@expect(MessageType.MalwareDatabaseScan)
async def complete_scan_db(self, message: MalwareDatabaseScan):
queued_scan = QueueSupervisorSync.queue.find(scanid=message["scan_id"])
if queued_scan:
QueueSupervisorSync.queue.remove(queued_scan)
await self._recheck_scan_queue()