helper.py

"""General helper functions for the app."""

import os
import re
import socket

import netaddr
from django.contrib.contenttypes.models import ContentType
from nautobot.dcim.filters import DeviceFilterSet
from nautobot.dcim.models import Device
from netaddr.core import AddrFormatError
from nornir_nautobot.exceptions import NornirNautobotException

from nautobot_device_onboarding.exceptions import OnboardException

FIELDS_PK = {
    "location",
    "role",
}

FIELDS_NAME = {"tags"}
MARKDOWN_ESCAPE_RE = re.compile(r"(?<!\\)([\*_\[\]\(\)`])")


def format_log_message(message):
    """Format a message for inclusion in Nautobot job logs."""
    return (
        MARKDOWN_ESCAPE_RE.sub(r"\\\1", message)
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace("\r\n", "<br>")
        .replace("\n", "<br>")
    )


def get_job_filter(data=None):
    """Helper function to return a the filterable list of OS's based on platform.name and a specific custom value."""
    if not data:
        data = {}
    query = {}

    # Translate instances from FIELDS set to list of primary keys
    for field in FIELDS_PK:
        if data.get(field):
            query[field] = [str(data[field].id)]

    # Translate instances from FIELDS set to list of names
    for field in FIELDS_NAME:
        if data.get(field):
            query[field] = [str(data[field].id)]
    # Handle case where object is from single device run all.
    if data.get("devices") and isinstance(data["devices"], Device):
        query.update({"id": [str(data["devices"].pk)]})
    elif data.get("devices"):
        query.update({"id": data["devices"].values_list("pk", flat=True)})
    base_qs = Device.objects.all()
    # {'debug': False, 'namespace': <Namespace: Global>, 'devices': <ConfigContextModelQuerySet [<Device: demo-cisco-xe>]>, 'location': None, 'device_role': None, 'tag': None, 'port': 22, 'timeout': 30}
    devices_filtered = DeviceFilterSet(data=query, queryset=base_qs)

    if not devices_filtered.qs.exists():
        raise NornirNautobotException(
            "`E3016:` The provided job parameters didn't match any devices detected. Please select the correct job parameters to correctly match devices."
        )
    devices_no_platform = devices_filtered.qs.filter(platform__isnull=True)
    if devices_no_platform.exists():
        raise NornirNautobotException(
            f"`E3017:` The following device(s) {', '.join([device.name for device in devices_no_platform])} have no platform defined. Platform is required."
        )

    return devices_filtered.qs


def onboarding_task_fqdn_to_ip(address):
    """Method to assure OT has FQDN resolved to IP address and rewritten into OT.

    If it is a DNS name, attempt to resolve the DNS address and assign the IP address to the
    name.

    Returns:
        None

    Raises:
        OnboardException("fail-general"):
            When a prefix was entered for an IP address
        OnboardException("fail-dns"):
            When a Name lookup via DNS fails to resolve an IP address
    """
    try:
        # If successful, this is an IP address and can pass
        netaddr.IPAddress(address)
        return address
    # Raise an Exception for Prefix values
    except ValueError as err:
        raise OnboardException(f"fail-general - ERROR appears a prefix was entered: {address}") from err
    # An AddrFormatError exception means that there is not an IP address in the field, and should continue on
    except AddrFormatError:
        try:
            # Perform DNS Lookup
            return socket.gethostbyname(address)
        except socket.gaierror as err:
            # DNS Lookup has failed, Raise an exception for unable to complete DNS lookup
            raise OnboardException(f"fail-dns - ERROR failed to complete DNS lookup: {address}") from err


def check_for_required_file(directory, filename):
    """
    Checks if a file named 'filename' exists within the specified directory.

    Args:
        directory: The path to the directory to check.
        filename: The name of the file to check for.

    Returns:
        True if the 'index' file exists in the directory, False otherwise.
    """
    try:
        for f_name in os.listdir(directory):
            if f_name == filename:
                return True
        return False
    except FileNotFoundError:
        return False


def add_content_type(job, model_to_add, target_object):
    """Add a content type to the valid content types of a target object.

    Args:
        job: The job object used for logging.
        model_to_add: The model class to get the content type for.
        target_object: The object to which the content type will be added.

    Raises:
        OnboardException: If adding the content type fails.
    """
    try:
        job.logger.info(
            "Adding %s content type to valid content types for location type %s",
            model_to_add.__name__,
            target_object,
        )
        content_type = ContentType.objects.get_for_model(model_to_add)
        target_object.content_types.add(content_type)
    except Exception as e:
        err_msg = f"Failed to add {model_to_add.__name__} to valid content types for {target_object}: {e}"
        job.logger.error(err_msg)
        raise OnboardException("fail-general - " + err_msg) from e