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