import peewee as pw


def migrate(migrator, database, fake=False, **kwargs):
    IPList = migrator.orm["iplist"]
    migrator.sql("UPDATE iplist SET manual=0 WHERE listname='GRAY'")
    migrator.sql("UPDATE iplist SET manual=1 WHERE listname='WHITE'")
    migrator.sql("UPDATE iplist SET manual=1 WHERE listname='BLACK'")
    migrator.sql(
        "UPDATE iplist SET listname='BLACK'"
        "WHERE listname='GRAY' AND no_captcha=1"
    )
    migrator.sql(
        "UPDATE iplist SET "
        "comment='Automatically blocked due to distributed attack', "
        "imported_from='Imunify360'"
        " WHERE listname='BLACK' AND manual=0"
    )
    migrator.remove_fields(IPList, "no_captcha")


def rollback(migrator, database, fake=False, **kwargs):
    IPList = migrator.orm["iplist"]
    migrator.add_fields(IPList, no_captcha=pw.BooleanField(default=False))