1#@!#!123s

: / opt / imunify360 / venv / lib / python3.11 / site-packages / im360 / migrations /
Filename : 004_ips_as_packed_network.py
back
import logging
import peewee as pw

from im360.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 <ip>/<netmask> 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