EVOLUTION-NINJA

Edit File: detached.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 os
import time
from abc import ABC, abstractmethod
from asyncio import gather
from pathlib import Path
from typing import Type

import psutil

from defence360agent.contracts.messages import Message
from defence360agent.defence360 import logger
from imav.contracts.messages import MalwareScanComplete
from imav.malwarelib.config import MalwareScanResourceType

PROCESS_START_TIME = 300  # sec

class DetachedState:
    FINISHED = "finished"
    ABORTED = "aborted"
    RUNNING = "running"
    NO_DIR = "no_dir"
    NO_PROCESS = "no_process"

class DetachedDir(ABC):
    DETACHED_DIR: str

done_file: Path
    err_file: Path
    log_file: Path
    pid_file: Path
    progress_file: Path

FILES = {
        "done_file": "done",
        "err_file": "err_file",
        "log_file": "log_file",
        "pid_file": "pid",
        "progress_file": "progress_file",
    }

def __init__(self, detached_id):
        """
        NOTE: Initialization should not create any files
        """
        self.path = Path(self.DETACHED_DIR, detached_id)
        for attr, file_name in self.FILES.items():
            setattr(self, attr, self.path / file_name)

def __enter__(self):
        self.path.mkdir(parents=True, exist_ok=True)
        self.log_file.touch()
        self.err_file.touch()
        return self

def __exit__(self, exc_type, exc_val, exc_tb):
        pass

def __str__(self):
        return str(self.path)

def __getattr__(self, name):
        # Delegate pathlib methods to the internal directory
        return getattr(self.path, name)

class DetachedOperation(ABC):
    NAME = None
    DETACHED_DIR_CLS: Type[DetachedDir]

@classmethod
    async def check_detached_operation_dir(cls, sink):
        restore_dir = cls.DETACHED_DIR_CLS.DETACHED_DIR
        check_coros = []
        if os.path.exists(restore_dir):
            entry: os.DirEntry
            with os.scandir(restore_dir) as it:
                for entry in it:
                    if entry.is_dir():
                        check_coros.append(
                            cls(entry.name).check_detached_operation(sink)
                        )

await gather(*check_coros)

def __init__(self, detached_id):
        self.detached_id = detached_id
        self.detached_dir = self.DETACHED_DIR_CLS(self.detached_id)

def __eq__(self, other):
        return self.detached_id == other.detached_id

def __hash__(self):
        # Two different DetachedScan's with the same ids
        # are practically the same
        return hash(self.detached_id)

@abstractmethod
    async def handle_aborted_process(self, *_, **__) -> None:
        pass

def get_detached_process_state(self, start_time=PROCESS_START_TIME) -> str:
        if not self.detached_dir.is_dir():
            logger.debug("No such scan directory: %s", self.detached_dir)
            return DetachedState.NO_DIR

if self.detached_dir.done_file.is_file():
            return DetachedState.FINISHED

try:
            pid = self.get_pid()
        except (FileNotFoundError, ValueError) as err:
            detached_dir_mtime = self.detached_dir.stat().st_mtime
            if time.time() - detached_dir_mtime > start_time:
                logger.warning(
                    "Pid file wasn't created or pid file is empty: %s,"
                    " error: %s",
                    self.detached_dir.pid_file,
                    err,
                )
                return DetachedState.ABORTED
            return DetachedState.NO_PROCESS

try:
            proc = psutil.Process(pid)
            if self.process_is_suitable(proc):
                return DetachedState.RUNNING
            else:  # not suitable, some new process with the same pid running
                logger.warning(
                    "The running process (pid=%s) does not match "
                    "the detached operation",
                    pid,
                )
        except (FileNotFoundError, psutil.Error) as err:
            logger.warning(
                "Can't get the process for the detached operation"
                " (pid=%s): %s",
                pid,
                err,
            )

return DetachedState.ABORTED

def get_pid(self):
        return int(self.detached_dir.pid_file.read_text())

@property
    @abstractmethod
    def on_complete_message(self) -> Message:
        pass

async def check_detached_operation(self, sink):
        """Updates operation to a consistent state.

Algorithm:
        1) If scan_dir doesn't exist:
            - scan_dir doesn't already create
            - scan ended and MalwareScanComplete is in process
        2) pid file doesn't exist:
            - since scan_dir mtime has passed more than 5 minutes
                -> something went wrong & handle_aborted_process
            - ... less than 5 minutes -> waiting for ai-bolit
        3) done file exists -> scan completed
        4) done file doesn't exist:
            - process executes -> scan running
            - process dead -> scan aborted
        """
        operation_state = self.get_detached_process_state()
        if operation_state == DetachedState.FINISHED:
            await sink.process_message(self.on_complete_message)
        elif operation_state == DetachedState.ABORTED:
            await self.handle_aborted_process(sink=sink)

@classmethod
    @abstractmethod
    def process_is_suitable(cls, proc) -> bool:
        pass

@abstractmethod
    async def complete(self) -> Message:
        pass

class DetachedScan(DetachedOperation, ABC):
    NAME = "scan"

@property
    @abstractmethod
    def RESOURCE_TYPE(self) -> MalwareScanResourceType:
        pass

@property
    @abstractmethod
    def progress(self) -> int:
        pass

@property
    @abstractmethod
    def phase(self) -> str:
        pass

@property
    @abstractmethod
    def total_resources(self) -> int:
        pass

@property
    def on_complete_message(self) -> Message:
        return MalwareScanComplete(
            scan_id=self.detached_id, resource_type=self.RESOURCE_TYPE
        )