ok
Mini Shell
"""
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/>.
Copyright © 2019 Cloud Linux Software Inc.
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 imav.malwarelib.config import MalwareScanResourceType
from defence360agent.contracts.messages import Message, MessageType
from defence360agent.defence360 import logger
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
@property
@abstractmethod
def DETACHED_DIR_CLS(self) -> Type[DetachedDir]:
pass
@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) -> 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:
if time.time() - self.detached_dir.stat().st_mtime > 300:
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
except (FileNotFoundError, psutil.Error):
pass
return DetachedState.ABORTED
def get_pid(self):
return int(self.detached_dir.pid_file.read_text())
@property
@abstractmethod
def on_complete_message(self):
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):
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 MessageType.MalwareScanComplete(
scan_id=self.detached_id, resource_type=self.RESOURCE_TYPE
)
Zerion Mini Shell 1.0