File: //opt/imunify360/venv/lib/python3.11/site-packages/defence360agent/migrations/120_scheduled_scan.py
import logging
from datetime import date, timedelta
import peewee as pw
from defence360agent.contracts.config import ConfigFile
from defence360agent.utils import importer
MalwareScanType = importer.get(
    module="imav.malwarelib.config", name="MalwareScanType", default=None
)
logger = logging.getLogger(__name__)
types = (
    MalwareScanType.ON_DEMAND,
    MalwareScanType.REALTIME,
    MalwareScanType.MALWARE_RESPONSE,
    MalwareScanType.BACKGROUND,
)
def _update_config(path=None):
    tomorrow = date.today() + timedelta(days=1)
    config = {
        "MALWARE_SCAN_SCHEDULE": {
            "day_of_month": tomorrow.day,
        }
    }
    try:
        config_file = ConfigFile(path=path)
        config_file.dict_to_config(config)
    except Exception:
        logger.exception("Failed to set malware scan schedule config")
def migrate(migrator, database, fake=False, **kwargs):
    MalwareScan = migrator.orm["malware_scans"]
    migrator.change_fields(
        MalwareScan,
        type=pw.CharField(
            null=False, constraints=[pw.Check("type in {}".format(types))]
        ),
    )
    if fake:
        return
    _update_config()
def rollback(migrator, database, fake=False, **kwargs):
    pass