import logging from peewee import CharField, FloatField from defence360agent.utils import importer MalwareHitStatus = importer.get( module="imav.malwarelib.config", name="MalwareHitStatus", default=None ) logger = logging.getLogger(__name__) def migrate(migrator, database, fake=False, **kwargs): MalwareHit = migrator.orm["malware_hits"] migrator.add_fields( MalwareHit, status=CharField(default=MalwareHitStatus.FOUND), cleaned_at=FloatField(null=True), ) def rollback(migrator, database, fake=False, **kwargs): MalwareHit = migrator.orm["malware_hits"] migrator.remove_fields(MalwareHit, "status", "cleaned_at")