22

✏️ Editing: persistent_mrs_queue.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 time
from logging import getLogger
from typing import Iterable

from peewee import Case, OperationalError, fn

from imav.malwarelib.model import (
    EXHAUSTED,
    PENDING,
    PROCESSING,
    WORKING_STATUSES,
    MRSQueueCounter,
    MRSUploadItem,
)
from imav.malwarelib.utils.sqlite_limits import sqlite_max_variable_number

logger = getLogger(__name__)

class PersistentMRSQueue:
    def __init__(
        self,
        storage_limit: int = 2_000_000,
        exhausted_ttl_seconds: float = 7 * 24 * 60 * 60,
        model: type[MRSUploadItem] = MRSUploadItem,
    ):
        self._storage_limit = storage_limit
        self._exhausted_ttl_seconds = exhausted_ttl_seconds
        self._model = model
        self._db = model._meta.database

def _increment_counter(self, name: str, delta: int) -> None:
        """Increment a persistent counter. Must be called inside atomic()."""
        if delta == 0:
            return
        (
            MRSQueueCounter.insert(name=name, value=delta)
            .on_conflict(
                conflict_target=[MRSQueueCounter.name],
                update={MRSQueueCounter.value: MRSQueueCounter.value + delta},
            )
            .execute()
        )

def put_many(self, items, upload_reason: str) -> int:
        rows = [
            {
                "hash": item.hash,
                "file_path": item.file,
                "upload_reason": upload_reason,
                "enqueued_at": time.time(),
                "status": PENDING,
            }
            for item in items
            if item.hash
        ]
        if not rows:
            return 0

chunk_size = max(
            1,
            sqlite_max_variable_number(self._db)
            // len(self._model._meta.columns),
        )

inserted = 0
        for i in range(0, len(rows), chunk_size):
            chunk = rows[i : i + chunk_size]
            try:
                with self._db.atomic():
                    chunk_inserted = (
                        self._model.insert_many(chunk)
                        .on_conflict_ignore()
                        .as_rowcount()
                        .execute()
                    )
                    self._increment_counter("enqueued_total", chunk_inserted)
                    self._increment_counter(
                        "deduplicated_total", len(chunk) - chunk_inserted
                    )
            except OperationalError:
                logger.exception(
                    "Failed to enqueue %d MRS upload items", len(chunk)
                )
                continue
            inserted += chunk_inserted

self.remove_exceeding_limit()
        return inserted

def take_batch(
        self, upload_reason: str, batch_size: int = 1000
    ) -> list[MRSUploadItem]:
        with self._db.atomic():
            query = (
                self._model.select()
                .where(
                    (self._model.status == PENDING)
                    & (self._model.upload_reason == upload_reason)
                )
                .order_by(self._model.enqueued_at, self._model.id)
                .limit(batch_size)
            )
            batch = list(query)
            if not batch:
                return []

batch_ids = [item.id for item in batch]
            now = time.time()
            for chunk in self._chunk_ids(batch_ids):
                (
                    self._model.update(
                        status=PROCESSING,
                        last_attempt_at=now,
                    )
                    .where(self._model.id.in_(chunk))
                    .execute()
                )

refreshed: list[MRSUploadItem] = []
            for chunk in self._chunk_ids(batch_ids):
                refreshed.extend(
                    self._model.select().where(self._model.id.in_(chunk))
                )
            refreshed.sort(key=lambda item: (item.enqueued_at, item.id))
            return refreshed

def mark_done(self, item_ids: list[int]) -> None:
        if not item_ids:
            return
        with self._db.atomic():
            for batch_ids in self._chunk_ids(item_ids):
                self._model.delete().where(
                    self._model.id.in_(batch_ids)
                ).execute()
            self._increment_counter("uploaded_total", len(item_ids))

def mark_failed(
        self, item_ids: list[int], error: str, max_attempts: int = 10
    ) -> None:
        if not item_ids:
            return

for batch_ids in self._chunk_ids(item_ids):
            with self._db.atomic():
                exhausted_count = (
                    self._model.select(fn.COUNT(self._model.id))
                    .where(
                        self._model.id.in_(batch_ids)
                        & ((self._model.attempts + 1) >= max_attempts)
                    )
                    .scalar()
                )
                (
                    self._model.update(
                        attempts=self._model.attempts + 1,
                        last_error=error,
                        last_attempt_at=time.time(),
                        status=Case(
                            None,
                            (
                                (
                                    (self._model.attempts + 1) >= max_attempts,
                                    EXHAUSTED,
                                ),
                            ),
                            PENDING,
                        ),
                    )
                    .where(self._model.id.in_(batch_ids))
                    .execute()
                )
                self._increment_counter("failed_total", len(batch_ids))
                self._increment_counter("exhausted_total", exhausted_count)

def requeue(self, item_ids: list[int], error: str | None = None) -> None:
        if not item_ids:
            return

update_data = {
            "status": PENDING,
            "last_attempt_at": time.time(),
        }
        if error is not None:
            update_data["last_error"] = error

for batch_ids in self._chunk_ids(item_ids):
            (
                self._model.update(**update_data)
                .where(self._model.id.in_(batch_ids))
                .execute()
            )

def requeue_stale_processing(
        self, older_than_seconds: float = 3600
    ) -> int:
        threshold = time.time() - older_than_seconds
        return (
            self._model.update(status=PENDING)
            .where(
                (self._model.status == PROCESSING)
                & (self._model.last_attempt_at < threshold)
            )
            .execute()
        )

def cleanup_exhausted(self) -> int:
        threshold = time.time() - self._exhausted_ttl_seconds
        return (
            self._model.delete()
            .where(
                (self._model.status == EXHAUSTED)
                & (
                    (self._model.last_attempt_at < threshold)
                    | self._model.last_attempt_at.is_null(True)
                )
            )
            .execute()
        )

def remove_exceeding_limit(self) -> int:
        overflow = self._working_count() - self._storage_limit
        if overflow <= 0:
            return 0

oldest_pending = list(
            self._model.select(self._model.id)
            .where(self._model.status == PENDING)
            .order_by(self._model.enqueued_at, self._model.id)
            .limit(overflow)
            .tuples()
        )
        item_ids = [item_id for (item_id,) in oldest_pending]
        if not item_ids:
            return 0

removed = 0
        for chunk in self._chunk_ids(item_ids):
            removed += (
                self._model.delete().where(self._model.id.in_(chunk)).execute()
            )
        return removed

def pending_reasons(self) -> list[str]:
        return [
            row.upload_reason
            for row in (
                self._model.select(self._model.upload_reason)
                .where(self._model.status == PENDING)
                .distinct()
            )
        ]

def pending_count(self) -> int:
        return self._count_by_status(PENDING)

def processing_count(self) -> int:
        return self._count_by_status(PROCESSING)

def failed_count(self) -> int:
        return self._count_by_status(EXHAUSTED)

def oldest_pending_age(self) -> float | None:
        oldest = (
            self._model.select(fn.MIN(self._model.enqueued_at))
            .where(self._model.status == PENDING)
            .scalar()
        )
        if oldest is None:
            return None
        return time.time() - oldest

def get_counters(self) -> dict[str, int]:
        """Return all counter name->value pairs (prefixed with mrs_queue.)."""
        return {
            f"mrs_queue.{row.name}": row.value
            for row in MRSQueueCounter.select()
            if not row.name.startswith("_")
        }

def get_gauges(self) -> dict[str, float]:
        """Return pending, processing counts and lag_seconds."""
        return {
            "mrs_queue.pending": float(self.pending_count()),
            "mrs_queue.processing": float(self.processing_count()),
            "mrs_queue.lag_seconds": self.oldest_pending_age() or 0.0,
        }

def get_last_metrics_report_time(self) -> float | None:
        """Return epoch timestamp of last metrics emission, or None."""
        SENTINEL = "_last_metrics_report"
        try:
            row = MRSQueueCounter.get_by_id(SENTINEL)
        except MRSQueueCounter.DoesNotExist:
            return None
        if row.last_metrics_report_time is None:
            return None
        return float(row.last_metrics_report_time)

def set_last_metrics_report_time(self, timestamp: float) -> None:
        """Record when metrics were last emitted."""
        SENTINEL = "_last_metrics_report"
        (
            MRSQueueCounter.insert(
                name=SENTINEL,
                value=0,
                last_metrics_report_time=int(timestamp),
            )
            .on_conflict(
                conflict_target=[MRSQueueCounter.name],
                update={
                    MRSQueueCounter.last_metrics_report_time: int(timestamp)
                },
            )
            .execute()
        )

def _count_by_status(self, status: str) -> int:
        return (
            self._model.select(fn.COUNT(self._model.id))
            .where(self._model.status == status)
            .scalar()
        )

def _working_count(self) -> int:
        return (
            self._model.select(fn.COUNT(self._model.id))
            .where(self._model.status.in_(WORKING_STATUSES))
            .scalar()
        )

@staticmethod
    def _chunk_ids(
        item_ids: Iterable[int], chunk_size: int = 500
    ) -> list[list[int]]:
        chunk = []
        chunks = []
        for item_id in item_ids:
            chunk.append(item_id)
            if len(chunk) >= chunk_size:
                chunks.append(chunk)
                chunk = []
        if chunk:
            chunks.append(chunk)
        return chunks

MR SKB