1#@!#!123s

https 1#@!#!123s
: / opt / imunify360 / venv / lib / python3.11 / site-packages / im360 / migrations /
Filename : 005_ip_network_v6_fix.py
back
import logging
import peewee as pw

from defence360agent.utils.validate import IP, NumericIPVersion
from im360.utils.net import unpack_ip_network

logger = logging.getLogger(__name__)


def migrate(migrator, database, fake=False, **kwargs):
    """Should fix all IPv6 addresses in the database.

    wrong record {"ip": "2602:fb54:3b6::"}
    correct record {"ip": "2602:fb54:3b6::/64"}

    """
    if fake:
        # Nothing to do in fake mode, it used to prepare models to following
        # migrations
        return

    try:
        IPList = migrator.orm["iplist"]

        with database.atomic():
            for ip_obj in (
                IPList.select()
                .where(IPList.version == NumericIPVersion[IP.V6])
                .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 005_ip_network_v6_fix %r",
            e,
        )


def rollback(migrator, database, fake=False, **kwargs):
    """Write your rollback migrations here."""
    pass