import logging import peewee as pw from defence360agent.utils.validate import IP from im360.utils.net import unpack_ip_network logger = logging.getLogger(__name__) def migrate(migrator, database, fake=False, **kwargs): if fake: return try: IPList = migrator.orm["iplist"] with database.atomic(): for ip_obj in IPList.select().dicts(): try: net = unpack_ip_network( ip_obj["network_address"], ip_obj["netmask"], ip_obj["version"], ) ip_str = IP.ip_net_to_string(net) # If ip format already / skip it if ip_str == ip_obj["ip"]: continue IPList.update(ip=ip_str).where( # from table: # primary_key = CompositeKey( # "network_address", "netmask", "version", "listname" # ) (IPList.network_address == ip_obj["network_address"]) & (IPList.netmask == ip_obj["netmask"]) & (IPList.version == ip_obj["version"]) & (IPList.listname == ip_obj["listname"]) ).execute() except ValueError as e: logger.warning( "Error processing IP %s with netmask %s: %s", ip_obj["ip"], ip_obj.get("netmask", "N/A"), str(e), ) except pw.IntegrityError as e: logger.warning("Error updating IP: %s", e) except Exception as e: logger.error( "Something wrong happened in migration" " 004_ips_as_packed_network %r", e, ) def rollback(migrator, database, fake=False, **kwargs): """Write your rollback migrations here.""" pass