22

✏️ Editing: ondemand.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>
"""
import logging
import os
import socket
import time
from operator import itemgetter
from typing import (
    Any,
    Callable,
    Dict,
    Iterable,
    List,
    Optional,
    Set,
    Tuple,
    Type,
)

from defence360agent.contracts.config import (
    GENERIC_SENSOR_SOCKET_PATH,
    Core,
    Malware,
    MalwareScanIntensity,
)
from defence360agent.contracts.permissions import (
    MS_ON_DEMAND_SCAN,
    check_permission,
)
from defence360agent.rpc_tools import ValidationError
from defence360agent.rpc_tools.lookup import (
    CommonEndpoints,
    RootEndpoints,
    bind,
)
from defence360agent.utils import antivirus_mode, get_abspath_from_user_dir
from imav.malwarelib.config import (
    ExitDetachedScanType,
    MalwareScanResourceType,
    QueuedScanState,
)
from imav.malwarelib.model import MalwareScan
from imav.malwarelib.scan.ai_bolit.detached import (
    AiBolitDetachedScan,
)
from imav.malwarelib.scan.crontab import get_crontab
from imav.malwarelib.scan.detached import (
    PROCESS_START_TIME,
    DetachedOperation,
    DetachedState,
)
from imav.malwarelib.scan.queue_supervisor_sync import (
    QueueSupervisorSync as ScanQueue,
)

logger = logging.getLogger(__name__)

if antivirus_mode.enabled:
    ABORTABLE_DETACHED_OPERATIONS: Set[Type[DetachedOperation]] = {
        AiBolitDetachedScan,
    }
else:
    from imav.malwarelib.scan.mds.detached import (
        MDSDetachedCleanup,
        MDSDetachedRestore,
        MDSDetachedScan,
    )

ABORTABLE_DETACHED_OPERATIONS: Set[Type[DetachedOperation]] = {
        AiBolitDetachedScan,
        MDSDetachedCleanup,
        MDSDetachedRestore,
        MDSDetachedScan,
    }

def _get_prepared_scan_list(
    since, to, limit, offset, order_by, queue: ScanQueue, user=None
) -> Tuple[int, List[Dict[str, Any]]]:
    db_kwargs = {}
    user_paths = []
    if user is not None:
        user_paths = [str(get_abspath_from_user_dir(user)), get_crontab(user)]
        db_kwargs["paths"] = user_paths

max_count, scans_from_db = MalwareScan.ondemand_list(
        since, to, limit, offset, order_by, **db_kwargs
    )
    for scan in scans_from_db:
        scan["scan_status"] = QueuedScanState.stopped.value
        if scan["started"] is None or scan["completed"] is None:
            scan["duration"] = None
        else:
            scan["duration"] = scan["completed"] - scan["started"]

queued_scans = {}
    if user is None:
        queued_scans = queue.scan_summaries()
    else:  # only user scan
        user_queued_scans = [
            scan for scan, _ in queue.get_scans_from_paths(paths=user_paths)
        ]
        if user_queued_scans:
            queued_scans = queue.scan_summaries(scans=user_queued_scans)

# Filter out incomplete scans
    incomplete_scans = []
    # Add incomplete scans with actual info
    for scanid, scan in queued_scans.items():
        if scan["started"] is None:
            scan["duration"] = None
        else:
            scan["duration"] = int(time.time()) - scan["started"]
        scan["completed"] = None
        incomplete_scans.append({"scanid": scanid, **scan})

complete_scans = [
        scan for scan in scans_from_db if scan["scanid"] not in queued_scans
    ]

scans = incomplete_scans + complete_scans

for scan in scans:
        # FIXME: remove total? Ask UI team.
        if scan.get("total") is None:
            scan["total"] = scan["total_resources"]
        if scan.get("started") is None:
            scan["started"] = scan["created"]
        if scan.get("created") is None:
            scan["created"] = scan["started"]
    scans = [u for u in scans if since <= u["created"] <= to]

if order_by:
        for order in reversed(order_by):
            scans.sort(key=itemgetter(order.column_name), reverse=order.desc)
    return max_count, scans[:limit]

# TODO: Change CommonEndpoints -> UserOnlyEndpoints
# cause in fact it is used by users:
# it is not available via CLI and called from user UI.
class OnDemandUserEndpoints(CommonEndpoints):
    def __init__(self, sink):
        super().__init__(sink)
        self.queue = ScanQueue(sink=sink)

@bind("malware", "on-demand", "start-user")
    async def ondemand_start(self, user):
        await check_permission(MS_ON_DEMAND_SCAN, user)

path = str(get_abspath_from_user_dir(user))
        scan_args = {
            "intensity_cpu": MalwareScanIntensity.USER_CPU,
            "intensity_io": MalwareScanIntensity.USER_IO,
            "intensity_ram": MalwareScanIntensity.USER_RAM,
            "initiator": user,
        }
        resource_types = [MalwareScanResourceType.FILE]
        if antivirus_mode.disabled and Malware.DATABASE_SCAN_ENABLED:
            resource_types.insert(0, MalwareScanResourceType.DB)
        for resource_type in resource_types:
            await self.queue.put(
                paths=[path],
                resource_type=resource_type,
                **scan_args,
            )

@bind("malware", "on-demand", "stop-user")
    async def ondemand_stop(self, user) -> None:
        await check_permission(MS_ON_DEMAND_SCAN, user)

current_scan = self.queue.queue.current_scan
        path = str(get_abspath_from_user_dir(user))
        for scan, _ in self.queue.get_scans_from_paths([path]):
            self.queue.remove([scan.scanid])
            kill = scan.scanid == current_scan.scanid
            await scan.detached_scan.handle_aborted_process(
                sink=self._sink,
                exit_type=ExitDetachedScanType.STOPPED,
                kill=kill,
                scan_started=scan.started,
            )

@bind("malware", "on-demand", "status-user")
    async def ondemand_status(self, user) -> Dict[str, Dict[str, str]]:
        path = str(get_abspath_from_user_dir(user))
        scan = self.queue.queue.find(path=path)
        status = (
            scan.status()
            if scan is not None
            else {"status": QueuedScanState.stopped.value}
        )
        result = {"items": {"status": status["status"]}}
        if "progress" in status:
            result["items"]["progress"] = status["progress"]
        return result

@bind("malware", "on-demand", "list-user")
    async def ondemand_list(
        self, user, since, to, limit, offset, order_by=None
    ) -> Tuple[int, List[Dict[str, Any]]]:
        """Get list of user scans.

- Parses info about scans from DB
        - Updates info about scans in scan_queue and parses it

:param since: scan start timestamp
        :param to: scan end timestamp
        :param limit: count of scans to print
        :param offset: offset of scans in DB to print
        :param order_by: name of column, by which sort data
        :return: (number of returned scans, list of data about scans)
        """
        assert user.strip() not in ("", "root"), f"Unexpected user '{user}'"
        return _get_prepared_scan_list(
            since,
            to,
            limit,
            offset,
            order_by,
            queue=self.queue.queue,
            user=user,
        )

class OnDemandEndpoints(RootEndpoints):
    def __init__(self, sink):
        super().__init__(sink)
        self.queue = ScanQueue(sink=sink)

@bind("malware", "on-demand", "start")
    async def ondemand_start(self, path, scan_file, scan_db, **scan_args):
        if not scan_file and not scan_db:
            raise ValidationError(
                "Either --scan-file or --scan-db should be specified"
            )
        if not self.queue.is_empty():
            raise ValidationError("On-demand scan is already running")

if scan_db:
            await self.queue.put(
                paths=[path],
                resource_type=MalwareScanResourceType.DB,
                initiator=None,
                **split_args(scan_args),
            )
        if scan_file:
            await self.queue.put(
                paths=[path],
                resource_type=MalwareScanResourceType.FILE,
                initiator=None,
                **split_args(scan_args),
            )

@bind("malware", "on-demand", "stop")
    async def ondemand_stop(self, all: bool) -> None:
        """CLI method to remove scans from ScanQueue

- Handles removed scans as aborted if they are detached

:param all
        """
        current_scan = self.queue.queue.current_scan
        if not all and current_scan:
            scans_to_stop = [current_scan]
        else:
            scans_to_stop = list(reversed(self.queue.queue.scans))

scan_ids = [scan.scanid for scan in scans_to_stop]
        self.queue.remove(scan_ids=scan_ids)
        for scan in scans_to_stop:
            kill = scan.scanid == current_scan.scanid
            await scan.detached_scan.handle_aborted_process(
                sink=self._sink,
                kill=kill,
                exit_type=ExitDetachedScanType.STOPPED,
                scan_started=scan.started,
            )

@bind("malware", "on-demand", "status")
    async def ondemand_status(self):
        status = self.queue.status()
        return {"items": status}

@bind("malware", "on-demand", "list")
    async def ondemand_list(
        self, since, to, limit, offset, order_by=None
    ) -> Tuple[int, List[Dict[str, Any]]]:
        """CLI method to print list of scans.

- Parses info about scans from DB
        - Updates info about scans in scan_queue and parses it

:param since: scan start timestamp
        :param to: scan end timestamp
        :param limit: count of scans to print
        :param offset: offset of scans in DB to print
        :param order_by: name of column, by which sort data
        :return: (number of returned scans, list of data about scans)
        """
        return _get_prepared_scan_list(
            since, to, limit, offset, order_by, queue=self.queue.queue
        )

@bind("malware", "on-demand", "queue", "put")
    async def ondemand_queue_put(
        self, paths, scan_file, scan_db, prioritize=False, **scan_args
    ):
        if not scan_file and not scan_db:
            raise ValidationError(
                "Either --scan-file or --scan-db should be specified"
            )

if scan_db:
            await self.queue.put(
                paths=paths,
                resource_type=MalwareScanResourceType.DB,
                prioritize=prioritize,
                **split_args(scan_args),
            )
        if scan_file:
            await self.queue.put(
                paths=paths,
                resource_type=MalwareScanResourceType.FILE,
                prioritize=prioritize,
                **split_args(scan_args),
            )

@bind("malware", "on-demand", "queue", "remove")
    async def ondemand_queue_remove(self, scan_ids: Optional[List[str]]):
        current_scan = self.queue.queue.current_scan
        if not scan_ids:
            queued_scans = [current_scan] if current_scan else []
            self.queue.remove()
        else:
            queued_scans = self.queue.queue.find_all(scan_ids)
            self.queue.remove(scan_ids)

for scan in queued_scans:
            await scan.detached_scan.handle_aborted_process(
                sink=self._sink,
                exit_type=ExitDetachedScanType.STOPPED,
                kill=scan is current_scan,
            )

@bind("malware", "on-demand", "check-detached")
    async def ondemand_check_detached(self):
        """
        Check if there are failed on-demand scans and start the agent so
        that we can handle aborted scans and proceed with the next scan
        in the queue
        """
        logger.info("Checking detached scan directory for failed scans")
        for detached_op in _active_detached_ops(dir_lister=_list_dir):
            _state = detached_op.get_detached_process_state(
                start_time=PROCESS_START_TIME
            )
            if _state == DetachedState.ABORTED:
                logger.info("Found failed scan. Waking up agent")
                _try_wake_up_agent()
                break
            if _state == DetachedState.FINISHED:
                logger.info("Finished scan found. Waking up agent")
                _try_wake_up_agent()
                break
        else:  # no break
            logger.info("No failed on-demand scans found")

def _list_dir(_dir):
    try:
        yield from os.listdir(_dir)
    except FileNotFoundError:
        return []

def _active_detached_ops(*, dir_lister: Callable[[str], Iterable[str]]):
    for detached_op_cls in ABORTABLE_DETACHED_OPERATIONS:
        for _id in dir_lister(detached_op_cls.DETACHED_DIR_CLS.DETACHED_DIR):
            yield detached_op_cls(_id)

def _try_wake_up_agent():
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
            s.settimeout(Core.DEFAULT_SOCKET_TIMEOUT)
            s.connect(GENERIC_SENSOR_SOCKET_PATH)
            s.send(b'{"method": "MALWARE_CHECK_DETACHED_SCANS"}\n')
    except (ConnectionRefusedError, FileNotFoundError, socket.timeout):
        pass

def _split_mask(mask):
    if mask is not None:
        return list(map(str.strip, mask.split(",")))

def split_args(scan_args):
    args = dict(scan_args)
    if "file_mask" in args:
        args["file_patterns"] = _split_mask(args.pop("file_mask"))
    if "ignore_mask" in args:
        args["exclude_patterns"] = _split_mask(args.pop("ignore_mask"))

return args

MR SKB