Edit

[EDIT] FILE: mrs_uploader.py

"""
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 <https://www.gnu.org/licenses/>.

This software is also available under ImunifyAV commercial license,
see <https://www.imunify360.com/legal/eula>
"""
from __future__ import annotations

import asyncio
import re
import time
from asyncio import CancelledError
from logging import getLogger

from defence360agent.api import inactivity
from defence360agent.contracts.config import Malware as Config
from defence360agent.contracts.license import LicenseError
from defence360agent.contracts.messages import GeneralMetrics, MessageType
from defence360agent.contracts.plugins import (
    MessageSink,
    MessageSource,
    expect,
)
from defence360agent.utils import (
    create_task_and_log_exceptions,
    safe_cancel_task,
)
from imav.contracts.messages import MalwareMRSUpload
from imav.malwarelib.utils import malware_response
from imav.malwarelib.utils.persistent_mrs_queue import PersistentMRSQueue

logger = getLogger(__name__)

METRICS_REPORT_INTERVAL = 300  # seconds

class MRSUploader(MessageSink, MessageSource):
    ERR_MSG = "Failed to submit a file"
    SUSP_PATTERN = re.compile(r"(?:suspicious\..+|[CS]MW-SUS-.+|SMW-HEUR-ELF)")

def __init__(self):
        self._queue = PersistentMRSQueue()
        self._new_items = asyncio.Event()
        self._shutdown = asyncio.Event()

async def create_source(self, loop, sink):
        self._sink = sink
        self._loop = loop
        self._queue.cleanup_exhausted()
        self._queue.requeue_stale_processing(older_than_seconds=0)
        if self._queue.pending_count():
            self._new_items.set()
        self._upload_task = create_task_and_log_exceptions(loop, self.upload)
        await self._maybe_emit_stale_metrics()
        self._metrics_task = create_task_and_log_exceptions(
            loop, self._report_metrics
        )

async def create_sink(self, loop):
        pass

async def shutdown(self):
        self._shutdown.set()
        self._new_items.set()
        await safe_cancel_task(self._upload_task)
        await self._emit_metrics()
        await safe_cancel_task(self._metrics_task)

async def _maybe_emit_stale_metrics(self):
        """Emit metrics on startup if last report is older than interval."""
        last_report = self._queue.get_last_metrics_report_time()
        if last_report is None or (
            time.time() - last_report > METRICS_REPORT_INTERVAL
        ):
            await self._emit_metrics()

async def _report_metrics(self):
        """Periodically emit queue metrics as GENERAL_METRICS."""
        while not self._shutdown.is_set():
            try:
                await asyncio.wait_for(
                    self._shutdown.wait(),
                    timeout=METRICS_REPORT_INTERVAL,
                )
            except asyncio.TimeoutError:
                pass
            if not self._shutdown.is_set():
                await self._emit_metrics()

async def _emit_metrics(self):
        """Collect and emit all queue metrics."""
        try:
            counters = self._queue.get_counters()
            gauges = self._queue.get_gauges()
            metrics = {**counters, **gauges}
            if not any(metrics.values()):
                return
            payload = [
                {"name": name, "value": value}
                for name, value in metrics.items()
            ]
            await self._sink.process_message(GeneralMetrics(payload))
            self._queue.set_last_metrics_report_time(time.time())
        except Exception:
            logger.exception("Failed to emit MRS queue metrics")

def _separate_hits_by_type(self, results) -> tuple:
        malicious = []
        suspicious = []
        extended_suspicious = []
        for file, data in results.items():
            is_extended_suspicious = False
            is_suspicious = False
            is_malicious = False
            for hit in data["hits"]:
                is_extended_suspicious |= hit.get("extended_suspicious", False)
                is_suspicious |= bool(
                    hit["suspicious"]
                    and self.SUSP_PATTERN.match(hit["matches"])
                )
                is_malicious |= not hit["suspicious"]
            hit_info = malware_response.HitInfo(file, data["hash"])
            if is_extended_suspicious:
                extended_suspicious.append(hit_info)
            elif is_suspicious:
                suspicious.append(hit_info)
            elif is_malicious:
                malicious.append(hit_info)
        return malicious, suspicious, extended_suspicious

@expect(MessageType.MalwareScan)
    async def process_scan(self, message):
        results = message["results"]
        if results is None:
            return
        if not Config.SEND_FILES:
            logger.info("Uploading files to MRS is disabled")
            return
        (
            malicious_hits,
            suspicious_hits,
            extended_suspicious_hits,
        ) = self._separate_hits_by_type(results)
        if malicious_hits:
            await self._sink.process_message(
                MalwareMRSUpload(
                    hits=malicious_hits, upload_reason="malicious"
                )
            )

if suspicious_hits:
            # Move uploading to another task
            await self._sink.process_message(
                MalwareMRSUpload(
                    hits=suspicious_hits, upload_reason="suspicious"
                )
            )
        if extended_suspicious_hits:  # pragma: no cover
            await self._sink.process_message(
                MalwareMRSUpload(
                    hits=extended_suspicious_hits,
                    upload_reason="extended-suspicious",
                )
            )
        errors = message["summary"].get("errors")
        if errors:
            error_hits = [
                malware_response.HitInfo(hit["file"], hit["hash"])
                for hit in errors
            ]
            await self._sink.process_message(
                MalwareMRSUpload(hits=error_hits, upload_reason="scan_error")
            )

async def upload(self):
        while not self._shutdown.is_set():
            await self._new_items.wait()
            self._new_items.clear()

try:
                await self._consume_pending()
            except CancelledError:
                raise
            except Exception:
                logger.exception("Unexpected error in MRS upload consumer")
                if not self._shutdown.is_set():
                    # A failure mid-batch leaves rows stuck in PROCESSING;
                    # return them to PENDING so the next cycle picks them up
                    # instead of idling on _new_items until a restart.
                    self._queue.requeue_stale_processing(older_than_seconds=0)
                    self._new_items.set()
                    await asyncio.sleep(1)

async def _consume_pending(self):
        while not self._shutdown.is_set():
            pending_reasons = self._queue.pending_reasons()
            if not pending_reasons:
                return

for upload_reason in pending_reasons:
                if self._shutdown.is_set():
                    return

while not self._shutdown.is_set():
                    batch = self._queue.take_batch(
                        upload_reason,
                        batch_size=malware_response.DEFAULT_CHUNK_SIZE,
                    )
                    if not batch:
                        break
                    if await self._process_batch(batch, upload_reason):
                        return

async def _process_batch(self, batch, upload_reason: str):
        unknown_hashes = set()
        saw_known_hashes_result = False

async for chunk in malware_response.check_known_hashes(
            self._loop,
            (item.hash for item in batch),
            upload_reason,
            chunk_size=malware_response.DEFAULT_CHUNK_SIZE,
        ):
            saw_known_hashes_result = True
            unknown_hashes.update(chunk)

if not saw_known_hashes_result:
            # Generator yielded nothing — `check_known_hashes` short-circuited
            # before issuing any request (IAIDTokenError or similar systemic
            # failure). Penalising per-item attempts here would EXHAUST the
            # whole queue after ~max_attempts seconds of transient auth/CAS
            # downtime; treat it like the LicenseError branch and requeue.
            self._queue.requeue(
                [item.id for item in batch],
                "Failed to check whether hashes are already known to MRS",
            )
            self._new_items.set()
            await asyncio.sleep(1)
            return True

known_item_ids = [
            item.id for item in batch if item.hash not in unknown_hashes
        ]
        if known_item_ids:
            self._queue.mark_done(known_item_ids)

unknown_items = [item for item in batch if item.hash in unknown_hashes]
        if not unknown_items:
            logger.info("All files are known to MRS. Skipping uploading...")
            return False

return await self._upload_unknown_items(unknown_items, upload_reason)

async def _upload_unknown_items(self, items, upload_reason: str):
        uploaded_item_ids = []
        retry_requested = False
        with inactivity.track.task("mrs_upload"):
            for index, item in enumerate(items):
                if self._shutdown.is_set():
                    if uploaded_item_ids:
                        self._queue.mark_done(uploaded_item_ids)
                    self._queue.requeue(
                        [pending_item.id for pending_item in items[index:]]
                    )
                    return True

try:
                    await malware_response.upload_file(
                        item.file_path, upload_reason=upload_reason
                    )
                    uploaded_item_ids.append(item.id)
                except LicenseError as e:
                    logger.warning(
                        "Cannot upload files to MRS with reason %s: %s",
                        upload_reason,
                        e,
                    )
                    if uploaded_item_ids:
                        self._queue.mark_done(uploaded_item_ids)
                    self._queue.requeue(
                        [pending_item.id for pending_item in items[index:]],
                        str(e),
                    )
                    self._new_items.set()
                    await asyncio.sleep(5)
                    return True
                except malware_response.FileTooLargeError as e:
                    logger.warning(
                        "%s: %s (file: %s)",
                        self.ERR_MSG,
                        e,
                        item.file_path,
                    )
                    self._queue.mark_failed([item.id], str(e), max_attempts=1)
                except malware_response.UploadFailure as e:
                    logger.warning(
                        "%s: %s (file: %s)",
                        self.ERR_MSG,
                        e,
                        item.file_path,
                    )
                    self._queue.mark_failed([item.id], str(e))
                    retry_requested = True
                except FileNotFoundError as e:
                    err = "{}. {}".format(self.ERR_MSG, e.strerror)
                    logger.warning("%s: %s", err, e.filename)
                    self._queue.mark_failed([item.id], err, max_attempts=1)
                except Exception as e:
                    logger.exception(
                        "Unexpected MRS upload error for %s",
                        item.file_path,
                    )
                    self._queue.mark_failed([item.id], str(e))
                    retry_requested = True

if uploaded_item_ids:
            self._queue.mark_done(uploaded_item_ids)
        if retry_requested:
            self._new_items.set()
            await asyncio.sleep(1)
        return retry_requested

@expect(MessageType.MalwareMRSUpload)
    async def process_hits(self, message: MalwareMRSUpload):
        upload_reason = message.get("upload_reason", "suspicious")
        inserted = self._queue.put_many(message["hits"], upload_reason)
        if inserted == 0:
            return
        self._new_items.set()

CANCEL