import peewee


def migrate(migrator, database, fake=False, **kwargs):
    malware_hits = migrator.orm["malware_hits"]
    migrator.remove_fields(malware_hits, "vendor")


def rollback(migrator, database, fake=False, **kwargs):
    malware_hits = migrator.orm["malware_hits"]
    migrator.add_fields(
        malware_hits, vendor=peewee.CharField(null=False, default="ai-bolit")
    )