OwlCyberSecurity - MANAGER

Edit File: 187_fix_scan_unserialization.py

"""
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