관리-도구

편집 파일: 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