import peewee as pw


def migrate(migrator, database, fake=False, **kwargs):
    IgnoredByPort = migrator.orm["ignored_by_port_proto"]
    Country = migrator.orm["country"]
    migrator.add_fields(
        IgnoredByPort,
        network_address=pw.IntegerField(null=True),
        netmask=pw.IntegerField(null=True),
        version=pw.IntegerField(null=True),
        country=pw.ForeignKeyField(Country, null=True),
    )


def rollback(migrator, database, fake=False, **kwargs):
    IgnoredByPort = migrator.orm["ignored_by_port_proto"]
    migrator.remove_fields(
        IgnoredByPort, "network_address", "netmask", "version", "country"
    )