import peewee as pw


def migrate(migrator, database, fake=False, **kwargs):
    MalwareScan = migrator.orm["malware_scans"]
    migrator.add_fields(
        MalwareScan,
        total_malicious=pw.IntegerField(null=False, default=0),
    )


def rollback(migrator, database, fake=False, **kwargs):
    MalwareScan = migrator.orm["malware_scans"]
    migrator.remove_fields(MalwareScan, "total_malicious")