D7net
Home
Console
Upload
information
Create File
Create Folder
About
Tools
:
/
opt
/
imunify360
/
venv
/
lib64
/
python3.11
/
site-packages
/
imav
/
plugins
/
Filename :
restore_from_backup.py
back
Copy
""" 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> """ from functools import partial from typing import Dict from defence360agent.api import inactivity from defence360agent.contracts.messages import MessageType from defence360agent.contracts.plugins import ( MessageSink, MessageSource, expect, ) from imav.malwarelib.config import MalwareHitStatus from imav.malwarelib.model import MalwareHit from imav.malwarelib.subsys.malware import MalwareAction from defence360agent.utils import Scope from imav.malwarelib.subsys.malware import MalwareActionIm360 _group_by_status = partial(MalwareHit.group_by_attribute, attribute="status") class RestoreFromBackupPlugin(MessageSink, MessageSource): """A plugin that handles RestoreFromBackupTask messages.""" SCOPE = Scope.IM360 def __init__(self): self._loop, self._sink = None, None async def create_sink(self, loop): self._loop = loop async def create_source(self, loop, sink): self._sink = sink @expect(MessageType.RestoreFromBackupTask) async def restore_from_backup(self, message: Dict): """Restores hits in RestoreFromBackupTask message from backup. - sets restore_from_backup_started status for received hits in the database. - successfully restored files are removed from the database. - statuses in the database are reverted from restore_from_backup_started to its original value for files that failed to restore. """ hits = message.get("hits") if not hits: return initiator = message.get("initiator") MalwareHit.set_status( hits, MalwareHitStatus.RESTORE_FROM_BACKUP_STARTED ) with inactivity.track.task("restore_from_backup"): to_restore = {} to_restore.setdefault(initiator, []).extend(hits) restored, failed = [], [] for initiator, init_hits in to_restore.items(): results = await MalwareAction.restore_from_backup( hits=init_hits, initiator=initiator, sink=self._sink, ) restored.extend( [hit for hit, event in results.items() if event.successful] ) failed.extend( [ hit for hit, event in results.items() if not event.successful ] ) MalwareHit.delete_instances(restored) original_status = _group_by_status(failed) for status, failed_hits in original_status.items(): MalwareHit.set_status(failed_hits, status)