Viewing file:  187_fix_scan_unserialization.py (1.51 KB) -rw-r--r--
Select action/file-type:
 (+) |  (+) |  (+) | Code (+) | Session (+) |  (+) | SDB (+) |  (+) |  (+) |  (+) |  (+) |  (+) |
"""
Used to fix issue with inability to unserialize stored scans.
See DEF-23121 for details.
"""
import importlib
import logging
import pickle
from pathlib import Path
logger = logging.getLogger(__name__)
SCANS_PATH = Path("/var/imunify360/aibolit/scans.pickle")
IM360_MALWARELIB = "im360.malwarelib"
AV_MALWARELIB = "imav.malwarelib"
class AVUnpickler(pickle.Unpickler):
def find_class(self, module, name):
try:
return super().find_class(module, name)
except ModuleNotFoundError:
if module.startswith(IM360_MALWARELIB):
av_module = importlib.import_module(
module.replace(IM360_MALWARELIB, AV_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 IM360_MALWARELIB.encode() in SCANS_PATH.read_bytes():
try:
with SCANS_PATH.open("rb") as f:
obj = AVUnpickler(f).load()
except Exception as exc:
logger.exception(
"Failed to load pickle scans %s: %s", SCANS_PATH, exc
)
else:
dump(obj, SCANS_PATH)
def rollback(migrator, *_, fake=False, **__):
pass
|