SWAMPy

✎ CARVE FLESH

EDITING: model.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 itertools
import operator
import os
import re
from dataclasses import dataclass
from functools import reduce
from operator import attrgetter
from pathlib import Path
from time import time
from typing import Any, Dict, Iterable, List, Set, Tuple, cast

from peewee import (
    SQL,
    BooleanField,
    Case,
    CharField,
    Check,
    Expression,
    FloatField,
    ForeignKeyField,
    IntegerField,
    ModelSelect,
    PrimaryKeyField,
    TextField,
    fn,
)
from playhouse.shortcuts import model_to_dict

from defence360agent.contracts.config import UserType
from defence360agent.model import Model, instance
from defence360agent.model.simplification import (
    FilenameField,
    ScanPathField,
    apply_order_by,
)
from defence360agent.utils import (
    execute_iterable_expression,
    get_abspath_from_user_dir,
    get_results_iterable_expression,
    split_for_chunk,
)
from imav.api.cleanup_revert import RemoteRevertHitInfo
from imav.malwarelib.config import (
    FAILED_TO_CLEANUP,
    MalwareHitStatus,
    MalwareScanResourceType,
    MalwareScanType,
    VulnerabilityHitStatus,
)
from imav.malwarelib.scan.crontab import get_crontab
from imav.malwarelib.scan.mds.report import MalwareDatabaseHitInfo
from imav.malwarelib.utils.sqlite_limits import sqlite_max_variable_number

class MalwareScan(Model):
    """Represents a batch of files scanned for malware

Usually a single AI-BOLIT execution.
    See :class:`.MalwareScanType` for possible kinds of scans.
    """

class Meta:
        database = instance.db
        db_table = "malware_scans"

#: An id of a scan, unique per server.
    scanid = CharField(primary_key=True)
    #: Scan start timestamp.
    started = IntegerField(null=False)
    #: Scan completion timestamp.
    completed = IntegerField(null=True)
    #: Scan type - reflects how and why the files were scanned.
    #: Must be one of :class:`.MalwareScanType`.
    type = CharField(
        null=False,
        constraints=[
            Check(
                "type in {}".format(
                    (
                        MalwareScanType.ON_DEMAND,
                        MalwareScanType.REALTIME,
                        MalwareScanType.MALWARE_RESPONSE,
                        MalwareScanType.BACKGROUND,
                        MalwareScanType.RESCAN,
                        MalwareScanType.USER,
                        MalwareScanType.RESCAN_OUTDATED,
                    )
                )
            )
        ],
    )
    #: The number of resources scanned.
    total_resources = IntegerField(null=False, default=0)
    #: For some types of scan - the directory or a file that was scanned.
    path = ScanPathField(null=True, default="")
    #: If not `null`, the scan did not finish successfully.
    #: Can be one of :class:`.ExitDetachedScanType` if scan was aborted or
    #: stopped by user, or an arbitrary error message for other kinds
    #: of issues.
    error = TextField(null=True, default=None)
    #: The number of malicious files found
    total_malicious = IntegerField(null=False, default=0)
    resource_type = CharField(
        null=False,
        constraints=[
            Check(
                "resource_type in {}".format(
                    (
                        MalwareScanResourceType.DB.value,
                        MalwareScanResourceType.FILE.value,
                    )
                )
            )
        ],
    )
    #: user who started the scan (None for root user)
    initiator = CharField(null=True)
    #: Timestamp when the entry was created/updated
    timestamp = IntegerField(null=False, default=lambda: int(time()))

@classmethod
    def ondemand_list(
        cls,
        since,
        to,
        limit,
        offset,
        order_by=None,
        *,
        types=(
            MalwareScanType.ON_DEMAND,
            MalwareScanType.BACKGROUND,
            MalwareScanType.USER,
        ),
        paths=None,
    ):
        query = (
            cls.select(
                cls.total_resources,
                cls.path,
                cls.scanid,
                cls.started,
                cls.completed,
                cls.error,
                cls.total_malicious,
                cls.type.alias("scan_type"),
                cls.resource_type,
            )
            .where(cls.type.in_(types))
            .where(cls.started >= since)
            .where(cls.started <= to)
        )

if paths:
            query = query.where(cls.path.in_(paths))

query = (
            query.group_by(
                cls.total_resources, cls.path, cls.scanid, cls.started
            )
            .order_by(MalwareScan.started.desc())
            .limit(limit)
            .offset(offset)
        )

if order_by is not None:
            query = apply_order_by(order_by, cls, query)

return query.count(clear_limit=True), list(query.dicts())

class MalwareHit(Model):
    """Represents a malicious or suspicious file."""

class Meta:
        database = instance.db
        db_table = "malware_hits"

#: An id of a scan, unique per server.
    id = PrimaryKeyField()
    #: A reference to :class:`MalwareScan`.
    scanid = ForeignKeyField(
        MalwareScan, null=False, related_name="hits", on_delete="CASCADE"
    )
    #: The owner of the file.
    owner = CharField(null=False)
    #: The user a file belongs to (is in user's home but owned by another user)
    user = CharField(null=False)
    #: The original path to the file.
    orig_file = FilenameField(null=False)
    #: The type of infection (signature).
    type = CharField(null=False)

# TODO: Deprecated (DEF-34575)
    #: Whether the file is malicious or just suspicious.
    #: Suspicious files are not displayed in UI but sent for analysis to MRS.
    malicious = BooleanField(null=False, default=False)

#: The hash of the files as provided by AI-BOLIT.
    hash = CharField(null=True)
    #: The size of the file.
    size = CharField(null=True)
    #: The exact timestamp when AI-BOLIT has detected the file.
    #:
    #: FIXME: unused? It looks like it was intended to resolve some possible
    #: race conditions with parallel scans, but we don't actually use it
    #: from the DB - we only compare the value in scan report
    #: with :attr:`cleaned_at`.
    timestamp = FloatField(null=True)

#: The current status of the file.
    #: Must be one of :class:`.MalwareHitStatus`.
    status = CharField(default=MalwareHitStatus.FOUND)
    #: Timestamp when the file was last cleaned.
    cleaned_at = FloatField(null=True)
    resource_type = CharField(
        null=False,
        constraints=[
            Check(
                "resource_type in {}".format(
                    (
                        MalwareScanResourceType.DB.value,
                        MalwareScanResourceType.FILE.value,
                    )
                )
            )
        ],
    )
    app_name = CharField(null=True)
    db_host = CharField(null=True)
    db_port = CharField(null=True)
    db_name = CharField(null=True)
    snippet = CharField(null=True)

@property
    def signature_id(self) -> str:
        return cast(str, self.type)

@property
    def orig_file_path(self):
        orig_file = cast(str, self.orig_file)
        return Path(orig_file)

class OrderBy:
        @staticmethod
        def status():
            return (
                Case(
                    MalwareHit.status,
                    (
                        (MalwareHitStatus.CLEANUP_PENDING, 0),
                        (MalwareHitStatus.CLEANUP_STARTED, 1),
                        (MalwareHitStatus.FOUND, 2),
                        (MalwareHitStatus.CLEANUP_DONE, 4),
                        (MalwareHitStatus.CLEANUP_REMOVED, 5),
                    ),
                    100,
                ),
            )

@classmethod
    def _hits_list(
        cls,
        clauses,
        since=0,
        to=None,
        limit=None,
        offset=None,
        search=None,
        by_scan_id=None,
        user=None,
        order_by=None,
        by_status=None,
        ids=None,
        site_search=None,
        user_sites=None,
        **kwargs,
    ) -> Tuple[int, List[dict[str, Any]]]:
        hits = cls.select(cls, MalwareScan).join(MalwareScan)
        to = to or time()

# Collect all WHERE conditions in a list to build a flat expression tree
        # It is workaround for `cursor stack overflowed` error
        # Since it is counted as one function call, when it is single clause
        # Otherwise sql driver will rise stack for each condition
        where_conditions = [
            clauses,
            (MalwareScan.started >= since) & (MalwareScan.started <= to),
        ]

if search is not None:
            pattern = f"%{search}%"
            where_conditions.append(
                SQL("CAST(orig_file AS TEXT) LIKE ?", (pattern,))
                | (cls.user**pattern)
            )
        if user is not None:
            where_conditions.append(MalwareHit.user == user)
        if by_scan_id is not None:
            where_conditions.append(MalwareScan.scanid == by_scan_id)
        if by_status is not None:
            where_conditions.append(MalwareHit.status << by_status)

if user_sites and site_search:
            where_conditions.append(cls.resource_type == "file")
            where_conditions.append(cls.orig_file.startswith(site_search))

# Find all paths that are more specific than the current search path.
            longer_paths = [
                p
                for p in user_sites
                if len(p) > len(site_search) and p.startswith(site_search)
            ]

if longer_paths:
                # To avoid recursion depth errors from chaining thousands of `&` operators,
                # we build a single SQL object containing all the NOT LIKE clauses.
                # This creates a flat structure that Peewee and the DB can handle.
                sql_chunks = []
                sql_params = []
                for path in longer_paths:
                    # Peewee's .startswith() translates to `LIKE 'path%'`, so we build
                    # the negative condition manually.
                    sql_chunks.append("orig_file NOT LIKE ?")
                    sql_params.append(f"{path}%")

combined_sql_string = " AND ".join(sql_chunks)

where_conditions.append(
                    SQL(f"({combined_sql_string})", sql_params)
                )

# Use reduce to safely combine the flat list of conditions.
        # The resulting expression tree will be shallow and will not overflow.
        full_clauses = reduce(operator.and_, where_conditions)

# `max_count` is used for pagination, must not include `ids`
        max_count_clauses = full_clauses
        if ids is not None:
            full_clauses &= MalwareHit.id.in_(ids)

ordered = hits.where(full_clauses).limit(limit).offset(offset)

if order_by is not None:
            ordered = apply_order_by(order_by, MalwareHit, ordered)

max_count = cls._hits_num(max_count_clauses)
        result = [row.as_dict() for row in ordered]

return max_count, result

@classmethod
    def _hits_num(
        cls, clauses=None, since=None, to=None, user=None, order_by=None
    ):
        if since and to:
            clauses &= (MalwareScan.started >= since) & (
                MalwareScan.started <= to
            )
        if user is not None:
            clauses &= cls.user == user
        q = cls.select(fn.COUNT(cls.id)).join(MalwareScan).where(clauses)
        if order_by is not None:
            q = apply_order_by(order_by, MalwareHit, q)
        return q.scalar()

@classmethod
    def malicious_num(cls, since, to, user=None):
        return cls._hits_num(
            (cls.status.not_in(MalwareHitStatus.CLEANUP) & cls.malicious),
            since,
            to,
            user,
        )

@classmethod
    def malicious_list(cls, *args, ignore_cleaned=False, **kwargs):
        clauses = cls.malicious
        if ignore_cleaned:
            clauses &= cls.status.not_in(MalwareHitStatus.CLEANUP)
        return cls._hits_list(clauses, *args, **kwargs)

@classmethod
    def malicious_summary(cls, since=0, to=None, user=None, **kwargs):
        to = time() if to is None else to
        clauses = (
            cls.malicious
            & (MalwareScan.started >= since)
            & (MalwareScan.started <= to)
        )
        if user is not None:
            clauses &= cls.user == user
        rows = (
            cls.select(cls.status, fn.COUNT(cls.id).alias("count"))
            .join(MalwareScan)
            .where(clauses)
            .group_by(cls.status)
        )
        counts = {row.status: row.count for row in rows}
        last_scan = MalwareScan.select(fn.MAX(MalwareScan.started)).scalar()
        return {"counts": counts, "last_scan": last_scan}

@classmethod
    def set_status(cls, hits, status, cleaned_at=None):
        hits = [row.id for row in hits]

def expression(ids, cls, status, cleaned_at):
            fields_to_update = {
                "status": status,
            }
            if cleaned_at is not None:
                fields_to_update["cleaned_at"] = cleaned_at

return cls.update(**fields_to_update).where(cls.id.in_(ids))

return execute_iterable_expression(
            expression, hits, cls, status, cleaned_at
        )

@classmethod
    def delete_instances(cls, to_delete: list):
        to_delete = [row.id for row in to_delete]

def expression(ids):
            return cls.delete().where(cls.id.in_(ids))

return execute_iterable_expression(expression, to_delete)

@classmethod
    def update_instances(cls, to_update: list):
        for data in to_update:
            for _instance, new_fields_data in data.items():
                for field, value in new_fields_data.items():
                    setattr(_instance, field, value)
                _instance.save()

@classmethod
    def is_infected(cls) -> Expression:
        clauses = (
            cls.status.in_(
                [
                    MalwareHitStatus.FOUND,
                ]
            )
            & cls.malicious
        )
        return clauses

@classmethod
    def is_suspicious(cls):
        return ~cls.malicious

@classmethod
    def malicious_select(
        cls, ids=None, user=None, cleanup=False, restore=False, **kwargs
    ):
        def expression(chunk_of_ids, cls, user):
            clauses = cls.malicious
            if chunk_of_ids is not None:
                clauses &= cls.id.in_(chunk_of_ids)
            elif cleanup:
                clauses &= cls.status.not_in(MalwareHitStatus.CLEANUP)
            elif restore:
                clauses &= cls.status.in_(MalwareHitStatus.RESTORABLE)
            if user is not None:
                if isinstance(user, str):
                    user = [user]
                clauses &= cls.user.in_(user)
            return cls.select().where(clauses)

return list(
            get_results_iterable_expression(
                expression, ids, cls, user, exec_expr_with_empty_iter=True
            )
        )

@classmethod
    def get_hits(cls, files, *, statuses=None):
        def expression(files):
            clauses = cls.orig_file.in_(files)
            if statuses:
                clauses &= cls.status.in_(statuses)
            return cls.select().where(clauses)

return get_results_iterable_expression(expression, files)

@classmethod
    def get_db_hits(cls, hits_info: Set[MalwareDatabaseHitInfo]):
        paths = [entry.path for entry in hits_info]
        apps = [entry.app_name for entry in hits_info]
        paths_apps = [(entry.path, entry.app_name) for entry in hits_info]
        hits = list(
            MalwareHit.select()
            .where(MalwareHit.orig_file.in_(paths))
            .where(MalwareHit.app_name.in_(apps))
        )
        hits = [
            hit for hit in hits if (hit.orig_file, hit.app_name) in paths_apps
        ]
        return hits

@classmethod
    def get_db_hits_for_remote_revert(
        cls, hits_info: list[RemoteRevertHitInfo]
    ):
        paths = [entry.app_root_path for entry in hits_info]
        apps = [entry.app_name for entry in hits_info]
        signatures = [entry.sig_id for entry in hits_info]
        paths_apps = [
            (entry.app_root_path, entry.app_name, entry.sig_id)
            for entry in hits_info
        ]
        hits: list[MalwareHit] = list(
            MalwareHit.db_hits()
            .where(MalwareHit.orig_file.in_(paths))
            .where(MalwareHit.app_name.in_(apps))
            .where(MalwareHit.type.in_(signatures))
            .where(MalwareHit.status.in_(MalwareHitStatus.RESTORABLE))
        )
        hits = [
            hit
            for hit in hits
            if (hit.orig_file, hit.app_name, hit.signature_id) in paths_apps
        ]
        return hits

@classmethod
    def delete_hits(cls, files):
        def expression(files):
            return cls.delete().where(cls.orig_file.in_(files))

return execute_iterable_expression(expression, files)

def refresh(self):
        return type(self).get(self._pk_expr())

@classmethod
    def refresh_hits(cls, hits: Iterable[MalwareHit], include_scan_info=False):
        def expression(hits):
            query = cls.select()
            if include_scan_info:  # use a single query to get scan info
                query = cls.select(cls, MalwareScan).join(MalwareScan)
            return query.where(cls.id.in_([hit.id for hit in hits]))

return list(get_results_iterable_expression(expression, hits))

@classmethod
    def cleaned_since(cls, timestamp: int):
        return cls.select().where(
            (cls.cleaned_at >= timestamp)
            & (cls.status.in_(MalwareHitStatus.CLEANED))
            & (cls.cleaned_at.is_null(False))
        )

@classmethod
    def db_hits(cls) -> ModelSelect:
        return cls.select().where(
            cls.resource_type == MalwareScanResourceType.DB.value
        )

@classmethod
    def db_hits_pending_cleanup(cls) -> ModelSelect:
        """Return db hits that are in queue for cleanup"""
        return cls.db_hits().where(
            cls.status == MalwareHitStatus.CLEANUP_PENDING,
        )

@classmethod
    def db_hits_under_cleanup(cls) -> ModelSelect:
        """Return db hits for which the cleanup is in progress"""
        return cls.db_hits().where(
            cls.status == MalwareHitStatus.CLEANUP_STARTED
        )

@classmethod
    def db_hits_under_restoration(cls) -> ModelSelect:
        """Return db hits for which the restore is in progress"""
        return cls.db_hits().where(
            cls.status == MalwareHitStatus.CLEANUP_RESTORE_STARTED
        )

@classmethod
    def db_hits_under_cleanup_in(cls, hit_info_set):
        """
        Return db hits for which the cleanup is in progress
        specified by the provided set of MalwareDatabaseHitInfo
        """
        # FIXME: Use peewee.ValuesList when peewee is updated
        # to obtain all hits using one query without additional processing
        path_set = {hit_info.path for hit_info in hit_info_set}
        app_name_set = {hit_info.app_name for hit_info in hit_info_set}
        path_app_name_set = {
            (hit_info.path, hit_info.app_name) for hit_info in hit_info_set
        }
        query = (
            cls.db_hits_under_cleanup()
            .where(cls.orig_file.in_(path_set))
            .where(cls.app_name.in_(app_name_set))
        )
        return [
            hit
            for hit in query
            if (hit.orig_file, hit.app_name) in path_app_name_set
        ]

@classmethod
    def db_hits_pending_cleanup_restore(cls):
        return cls.db_hits().where(
            cls.status.in_(
                [
                    MalwareHitStatus.CLEANUP_RESTORE_PENDING,
                    MalwareHitStatus.CLEANUP_REMOTE_RESTORE_PENDING,
                ]
            )
        )

@classmethod
    def db_hits_under_cleanup_restore(cls):
        return cls.db_hits().where(
            cls.status == MalwareHitStatus.CLEANUP_RESTORE_STARTED
        )

@staticmethod
    def group_by_attribute(
        *hit_list_list: List["MalwareHit"], attribute: str
    ) -> Dict[str, List["MalwareHit"]]:
        hit_list = sorted(
            (hit for hit in itertools.chain.from_iterable(hit_list_list)),
            key=attrgetter(attribute),
        )
        return {
            attr_value: list(hits)
            for attr_value, hits in itertools.groupby(
                hit_list,
                key=attrgetter(attribute),
            )
        }

def as_dict(self):
        return {
            "id": self.id,
            "username": self.user,
            "file": self.orig_file,
            "created": self.scanid.started,
            "scan_id": self.scanid_id,
            "scan_type": self.scanid.type,
            "resource_type": self.resource_type,
            "type": self.type,
            "hash": self.hash,
            "size": self.size,
            "malicious": self.malicious,
            "status": self.status,
            "cleaned_at": self.cleaned_at,
            "extra_data": {},
            "db_name": self.db_name,
            "app_name": self.app_name,
            "db_host": self.db_host,
            "db_port": self.db_port,
            "snippet": self.snippet,
            "table_fields": (
                list(
                    MalwareHistory.select(
                        MalwareHistory.table_name,
                        MalwareHistory.table_field,
                        MalwareHistory.table_row_inf,
                    )
                    .where(
                        MalwareHistory.app_name == self.app_name,
                        MalwareHistory.db_host == self.db_host,
                        MalwareHistory.db_port == self.db_port,
                        MalwareHistory.db_name == self.db_name,
                        MalwareHistory.path == self.orig_file,
                        MalwareHistory.resource_type == self.resource_type,
                        MalwareHistory.scan_id == self.scanid,
                        MalwareHistory.signature_id == self.type,
                        MalwareHistory.table_name.is_null(False),
                        MalwareHistory.table_field.is_null(False),
                        MalwareHistory.table_row_inf.is_null(False),
                    )
                    .dicts()
                )
                if self.resource_type == MalwareScanResourceType.DB.value
                else []
            ),
        }

def __repr__(self):
        if self.app_name:
            return "%s(orig_file=%r, app_name=%r)" % (
                self.__class__.__name__,
                self.orig_file,
                self.app_name,
            )
        return "%s(orig_file=%r)" % (self.__class__.__name__, self.orig_file)

@dataclass(frozen=True)
class MalwareHitAlternate:
    """
    Used as a replacement for MalwareHit for file hits only
    """

scanid: str
    orig_file: str
    # app_name is always None for file hits
    app_name: None
    owner: str
    user: str
    size: int
    hash: str
    type: str
    timestamp: int
    malicious: bool

@classmethod
    def create(cls, scanid, filename, data):
        return cls(
            scanid=scanid,
            orig_file=filename,
            app_name=None,
            owner=data["owner"],
            user=data["user"],
            size=data["size"],
            hash=data["hash"],
            type=data["hits"][0]["matches"],
            timestamp=data["hits"][0]["timestamp"],
            malicious=not data["hits"][0]["suspicious"],
        )

@property
    def orig_file_path(self):
        return Path(os.fsdecode(self.orig_file))

class MalwareIgnorePath(Model):
    """A path that must be excluded from all scans"""

class Meta:
        database = instance.db
        db_table = "malware_ignore_path"
        indexes = ((("path", "resource_type"), True),)  # True refers to unique

CACHE = None

id = PrimaryKeyField()
    #: The path itself. Wildcards or patterns are NOT supported.
    path = CharField()
    resource_type = CharField(
        null=False, constraints=[Check("resource_type in ('file','db')")]
    )
    #: Timestamp when it was added.
    added_date = IntegerField(null=False, default=lambda: int(time()))

@classmethod
    def _update_cache(cls):
        items = list(cls.select().order_by(cls.path).dicts())
        cls.CACHE = items

@classmethod
    def create(cls, **kwargs):
        cls.CACHE = None
        return super(MalwareIgnorePath, cls).create(**kwargs)

@classmethod
    def delete(cls):
        cls.CACHE = None
        return super(MalwareIgnorePath, cls).delete()

@classmethod
    def paths_count_and_list(
        cls,
        limit=None,
        offset=None,
        search=None,
        resource_type: str | None = None,
        user=None,
        since=None,
        to=None,
        order_by=None,
    ):
        q = cls.select().order_by(cls.path)
        if since is not None:
            q = q.where(cls.added_date >= since)
        if to is not None:
            q = q.where(cls.added_date <= to)
        if search is not None:
            q = q.where(cls.path.contains(search))
        if resource_type is not None:
            q = q.where(cls.resource_type == resource_type)
        if offset is not None:
            q = q.offset(offset)
        if limit is not None:
            q = q.limit(limit)
        if order_by is not None:
            q = apply_order_by(order_by, cls, q)
        if user is not None:
            user_home = get_abspath_from_user_dir(user)
            q = q.where(
                (cls.path.startswith(str(user_home) + "/"))
                | (cls.path == str(user_home))
                | (cls.path == str(get_crontab(user)))
            )

max_count = q.count(clear_limit=True)
        return (
            max_count,
            [model_to_dict(row) for row in q],
        )

@classmethod
    def path_list(cls, *args, **kwargs) -> List[str]:
        _, path_list = cls.paths_count_and_list(*args, **kwargs)
        return [row["path"] for row in path_list]

@classmethod
    async def is_path_ignored(cls, check_path):
        """Checks whether path stored in MalwareIgnorePath cache or
        if it's belongs to path from cache or if it matches patters from cache

:param str check_path: path to check
        :return: bool: is ignored according MalwareIgnorePath
        """
        if cls.CACHE is None:
            cls._update_cache()
        path = Path(check_path)
        if cls.CACHE:
            for p in cls.CACHE:
                await asyncio.sleep(0)
                ignored_path = Path(p["path"])
                if (path == ignored_path) or (ignored_path in path.parents):
                    return True
        return False

class MalwareHistory(Model):
    """Records every event related to :class:`MalwareHit` records"""

class Meta:
        database = instance.db
        db_table = "malware_history"

#: The path of the file.
    path = FilenameField(null=False)
    app_name = CharField(null=True)
    resource_type = CharField(
        null=False,
        constraints=[
            Check(
                "resource_type in {}".format(
                    (
                        MalwareScanResourceType.DB.value,
                        MalwareScanResourceType.FILE.value,
                    )
                )
            )
        ],
        default=MalwareScanResourceType.FILE.value,
    )
    #: What happened with the file. Should be one of :class:`.MalwareEvent`.
    event = CharField(null=False)
    #: What kind of scan has detected the file, or `manual` for manual actions.
    #: See :class:`.MalwareScanType`.
    cause = CharField(null=False)
    #: The name of the user who has triggered the event.
    initiator = CharField(null=False)
    #: A snapshot of :attr:`MalwareHit.owner`
    file_owner = CharField(null=False)
    #: A snapshot of :attr:`MalwareHit.user`
    file_user = CharField(null=False)
    #: Timestamp when the event took place.
    ctime = IntegerField(null=False, default=lambda: int(time()))
    #: Database host name (for db type scan).
    db_host = CharField(null=True)
    #: Database port (for db type scan).
    db_port = CharField(null=True)
    #: Database name (for db type scan).
    db_name = CharField(null=True)
    #: Infected table name (for db type scan)
    table_name = CharField(null=True)
    #: Infected field name (for db type scan)
    table_field = CharField(null=True)
    #: Infected table row id (for db type scan)
    table_row_inf = IntegerField(null=True)
    #: Scan ID reference (for generating `table_fields`)
    scan_id = CharField(null=True)
    #: Signature ID
    signature_id = CharField(null=True)

@classmethod
    def get_history(
        cls, since, to, limit, offset, user=None, search=None, order_by=None
    ):
        clauses = (cls.ctime >= since) & (cls.ctime <= to)
        if search:
            clauses &= (cls.event.contains(search)) | (
                SQL("(INSTR(path, ?))", (search,))
            )

if user:
            clauses &= cls.file_user == user
        query = cls.select().where(clauses).limit(limit).offset(offset).dicts()
        if order_by is not None:
            query = apply_order_by(order_by, MalwareHistory, query)

list_result = list(query)

return query.count(clear_limit=True), list_result

@classmethod
    def save_event(cls, **kwargs):
        cls.insert(
            initiator=kwargs.pop("initiator", None) or UserType.ROOT,
            cause=kwargs.pop("cause", None) or MalwareScanType.MANUAL,
            resource_type=kwargs.pop("resource_type", None)
            or MalwareScanResourceType.FILE.value,
            **kwargs,
        ).execute()

@classmethod
    def save_events(cls, hits: List[dict]):
        with instance.db.atomic():
            chunk_size = max(
                1,
                sqlite_max_variable_number(cls._meta.database)
                // len(cls._meta.columns),
            )
            for hits_chunk in split_for_chunk(hits, chunk_size=chunk_size):
                cls.insert_many(hits_chunk).execute()

@classmethod
    def get_failed_cleanup_events_count(cls, paths: list, *, since: int):
        return (
            cls.select(cls.path, fn.COUNT())
            .where(
                cls.path.in_(paths)
                & (cls.event == FAILED_TO_CLEANUP)
                & (cls.ctime >= since)
            )
            .group_by(cls.path)
            .tuples()
        )

class VulnerabilityHit(Model):
    """Represents a vulnerable file."""

class Meta:
        database = instance.db
        db_table = "vulnerability_hits"

#: An id of a detection, unique per server.
    id = PrimaryKeyField()
    #: A reference to :class:`MalwareScan`.
    scanid = ForeignKeyField(
        MalwareScan,
        null=False,
        related_name="vulnerabilities",
        on_delete="CASCADE",
    )
    #: The owner of the file.
    owner = CharField(null=False)
    #: The user a file belongs to (is in user's home but owned by another user)
    user = CharField(null=False)
    #: The original path to the file.
    orig_file = FilenameField(null=False)
    #: The type of infection (signature).
    type = CharField(null=False)
    #: The hash of the files as provided by AI-BOLIT.
    hash = CharField(null=True)
    #: The size of the file.
    size = CharField(null=True)
    #: The exact timestamp when AI-BOLIT has detected the file.
    timestamp = FloatField(null=True)
    #: The current status of the file.
    #: Must be one of :class:`.VulnerabilityHitStatus`.
    status = CharField(default=VulnerabilityHitStatus.VULNERABLE)
    #: Timestamp when the file was last patched.
    patched_at = FloatField(null=True)

@property
    def orig_file_path(self):
        orig_file = cast(str, self.orig_file)
        return Path(orig_file)

@classmethod
    def match(cls, signature: str) -> bool:
        return signature.startswith("VULN-ESUS-")

@classmethod
    def get_vulnerability_ids(cls, signature: str) -> list:
        if match := re.search(r"VULN-ESUS-([\d,]+)", signature):
            return match.groups()[0].split(",")
        return []

return get_results_iterable_expression(expression, files)

@classmethod
    def delete_hits(cls, files):
        def expression(files):
            return cls.delete().where(cls.orig_file.in_(files))

return execute_iterable_expression(expression, files)

@classmethod
    def _hits_list(
        cls,
        since=0,
        to=None,
        limit=None,
        offset=None,
        search=None,
        by_scan_id=None,
        user=None,
        order_by=None,
        by_status=None,
        ids=None,
        **kwargs,
    ):
        hits = cls.select(cls, MalwareScan).join(MalwareScan)

to = to or time()
        pattern = "%{}%".format(search)
        full_clauses = (MalwareScan.started >= since) & (
            MalwareScan.started <= to
        )
        if search is not None:
            full_clauses &= SQL(
                "CAST(orig_file AS TEXT) LIKE ?", (pattern,)
            ) | (cls.user**pattern)
        if user is not None:
            full_clauses &= cls.user == user
        if by_scan_id is not None:
            full_clauses &= cls.scanid == by_scan_id
        if by_status is not None:
            full_clauses &= cls.status << by_status
        # `max_count` is used for pagination, must not include `ids`
        max_count_clauses = full_clauses
        if ids is not None:
            full_clauses &= cls.id.in_(ids)

ordered = hits.where(full_clauses).limit(limit).offset(offset)

if order_by is not None:
            ordered = apply_order_by(order_by, cls, ordered)

max_count = (
            cls.select(fn.COUNT(cls.id))
            .join(MalwareScan)
            .where(max_count_clauses)
            .scalar()
        )
        result = [row.as_dict() for row in ordered]

return max_count, result

@classmethod
    def list(cls, *args, **kwargs):
        return cls._hits_list(*args, **kwargs)

@staticmethod
    def group_by_attribute(
        *hit_list_list: List["VulnerabilityHit"], attribute: str
    ) -> Dict[str, List["VulnerabilityHit"]]:
        hit_list = sorted(
            (hit for hit in itertools.chain.from_iterable(hit_list_list)),
            key=attrgetter(attribute),
        )
        return {
            attr_value: list(hits)
            for attr_value, hits in itertools.groupby(
                hit_list,
                key=attrgetter(attribute),
            )
        }

@classmethod
    def get_vulnerabilities_ids(cls, hits: list) -> list[int]:
        vuln_ids = set()
        for hit in hits:
            vuln_ids |= set(
                VulnerabilityHit.get_vulnerability_ids(hit["type"])
            )
        return list(vuln_ids)

@classmethod
    def set_status(cls, hits, status, patched_at=None):
        hits = [row.id for row in hits]

def expression(ids, cls, status, patched_at):
            fields_to_update = {
                "status": status,
            }
            if patched_at is not None:
                fields_to_update["patched_at"] = patched_at

return cls.update(**fields_to_update).where(cls.id.in_(ids))

return execute_iterable_expression(
            expression, hits, cls, status, patched_at
        )

def as_dict(self):
        return {
            "id": self.id,
            "username": self.user,
            "file_path": self.orig_file,
            "created": self.scanid.started,
            "scan_id": self.scanid_id,
            "scan_type": self.scanid.type,
            "type": self.type,
            "hash": self.hash,
            "size": self.size,
            "status": self.status,
            "patched_at": self.patched_at,
        }

class ImunifyPatchSubscription(Model):
    """Stores Imunify Patch user subscriptions."""

class Meta:
        database = instance.db
        table_name = "imunify_patch_subscriptions"

user_id = TextField(primary_key=True)

PENDING = "pending"
PROCESSING = "processing"
EXHAUSTED = "exhausted"
WORKING_STATUSES = (PENDING, PROCESSING)

class MRSUploadItem(Model):
    class Meta:
        database = instance.db
        db_table = "mrs_upload_queue"
        indexes = (
            (("hash", "upload_reason"), True),
            (("status", "upload_reason", "enqueued_at"), False),
            (("status", "last_attempt_at"), False),
        )

id = PrimaryKeyField()
    hash = CharField(max_length=64, null=False)
    file_path = TextField(null=False)
    upload_reason = CharField(max_length=32, null=False)
    enqueued_at = FloatField(null=False)
    status = CharField(max_length=16, null=False, default=PENDING)
    attempts = IntegerField(default=0)
    last_error = TextField(null=True)
    last_attempt_at = FloatField(null=True)

class MRSQueueCounter(Model):
    class Meta:
        database = instance.db
        db_table = "mrs_queue_counters"

name = CharField(max_length=64, primary_key=True)
    value = IntegerField(default=0)
    last_metrics_report_time = IntegerField(null=True)

CANCEL

Name	Type	Size	Modified	Actions
↩ ..	DIR	—	—
📁 __pycache__/	DIR	—	2026-08-14 05:35
📁 advanced/	DIR	—	2026-08-14 05:35
📁 api/	DIR	—	2026-08-14 05:35
📁 cleanup/	DIR	—	2026-08-14 05:35
📁 difflib/	DIR	—	2026-08-14 05:35
📁 plugins/	DIR	—	2026-08-14 05:35
📁 rpc/	DIR	—	2026-08-14 05:35
📁 scan/	DIR	—	2026-08-14 05:35
📁 subsys/	DIR	—	2026-08-14 05:35
📁 utils/	DIR	—	2026-08-14 05:35
📁 vulnerabilities/	DIR	—	2026-08-14 05:35
📄 __init__.py	PY	0 B	2026-08-05 11:17	EDIT
📄 config.py	PY	7.6 KB	2026-08-05 11:17	EDIT
📄 engine.py	PY	4.1 KB	2026-08-05 11:17	EDIT
📄 model.py	PY	37.9 KB	2026-08-05 11:17	EDIT
📄 mrs_upload_queue.py	PY	922 B	2026-08-05 11:17	EDIT