/opt/imunify360/venv/lib/python3.11/site-packages/imav/migrations
Edit: /opt/imunify360/venv/lib/python3.11/site-packages/imav/migrations/001_fix_scan_unserialization.py (2923B)
"""
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
.
Copyright © 2019 Cloud Linux Software Inc.
This software is also available under ImunifyAV commercial license,
see
Used to fix issue with inability to unserialize stored scans.
See DEF-23158 for details.
"""
import collections
import importlib
import logging
import pickle
from pathlib import Path
logger = logging.getLogger(__name__)
SCANS_PATH = Path("/var/imunify360/aibolit/scans.pickle")
OLD_MALWARELIB = "defence360agent.malwarelib"
NEW_MALWARELIB = "imav.malwarelib"
class AVUnpickler(pickle.Unpickler):
def find_class(self, module, name):
try:
return super().find_class(module, name)
except (ImportError, ModuleNotFoundError):
if module.startswith(OLD_MALWARELIB):
av_module = importlib.import_module(
module.replace(OLD_MALWARELIB, NEW_MALWARELIB)
)
return getattr(av_module, name)
raise
def dump(obj, path):
temp_path = path.with_name(path.name + ".temp")
with temp_path.open("wb") as f:
pickle.dump(obj, f)
# to avoid the possibility of leaving a broken file,
# if any errors occurred above
temp_path.replace(path)
def migrate(migrator, *_, fake=False, **__):
if fake or not SCANS_PATH.exists():
return
if OLD_MALWARELIB.encode() in SCANS_PATH.read_bytes():
try:
with SCANS_PATH.open("rb") as f:
scans = AVUnpickler(f).load()
# convert class instance to a simple dict
updated = collections.deque(
{
"path": item.path,
"scanid": item.scanid,
"scan_type": item.scan_type,
"created": item.created,
"started": item.started,
"args": item.args,
"resource_type": item.resource_type.value,
"state": item.state.value,
}
for item in scans
)
except Exception as exc:
logger.exception(
"Failed to load pickle scans %s: %s", SCANS_PATH, exc
)
else:
dump(updated, SCANS_PATH)
def rollback(migrator, *_, fake=False, **__):
pass